A nicer way to read chunks of bytes from a socket

I'm implementing a minimalistic LED-stripe-service. It just listens on a raw socket for the number of LEDs followed by their values. Only a single connection is possible.

I'm a bit unsatisfied by the part where I have to read the LED values from the socket via a temp slice. I'd like to have it more functional without the risk of any heap allocation happening within that tight loop. Some kind of chunks()-method would have been nice. Is such a frequent call of read_exact inefficient*? Would it be better to read all data into a buffer, split it, use it and throw it away?

*) I just answered that by myself by accident :slight_smile: BufReader

Also, I don't know how to split the nested program structure into multiple functions. I'd either have to

  • move led_bytes into the function, which will result in frequent heap allocations(?)
  • pass &mut led_bytes to the function, which is quite ugly

side note: I'd like to avoid unsafe code.

use std::net::TcpListener;
use std::io::prelude::*;

// all those constants will be dynamic in future implementations
const CHANNELS_PER_LED: usize = 3;
const BYTES_PER_CHANNEL: usize = 2;
const BITS_PER_SAMPLE: usize = 16;
const BYTES_PER_LED: usize = BYTES_PER_CHANNEL * CHANNELS_PER_LED;

fn main() {
    // let spi_options = SpidevOptions::new().bits_per_word(8).max_speed_hz(MAX_SPI_FREQUENCY).lsb_first(false).mode(SPI_MODE_0).build();
    // let mut spidev = Spidev::open("/dev/spidev0.0").expect("Cannot open SPI device");
    // spidev.configure(&spi_options).expect("Failed to set SPI options");

    let listener = TcpListener::bind("0.0.0.0:56501").expect("Cannot start listening for incoming connections");

    'next_connection: for stream in listener.incoming() {
        match stream {
            Ok(mut stream) => {
                println!("New connection ...");

                // gets overwritten in a tight loop; bigger scope shall prevent frequent heap allocations
                // size will be dynamic in future implementation
                let mut led_bytes = vec![0; BYTES_PER_LED];

                // read length of upcoming data set or quit
                let mut length_bytes = [0; 2];
                while let Ok(_) = stream.read_exact(&mut length_bytes) {
                    // join bytes as big endian (network byte order)
                    let led_count = (length_bytes[0] as usize) << 8 | length_bytes[1] as usize;
                    println!("led count: {:?}", led_count);

                    // read values of LEDs
                    for _ in 0..led_count {
                        // read all bytes required to describe the values of the LED's channels
                        if let Err(_) = stream.read_exact(&mut led_bytes) { continue 'next_connection }

                        // combine bytes as big endian (network byte order)
                        // Later I will specialize this for the allowed lengths [1 ... 8]
                        let led_data = led_bytes.iter().fold(0u64, |acc, &x| acc << 8 | x as u64);

                        // decode led data into separate channels
                        // this will get more flexible in future implementations
                        let (red, green, blue) = (
                            (led_data >> 2 * BITS_PER_SAMPLE) as u16,
                            (led_data >> 1 * BITS_PER_SAMPLE) as u16,
                            (led_data >> 0 * BITS_PER_SAMPLE) as u16,
                        );

                        // debug output
                        println!("rgb: {:4x} {:4x} {:4x}", red, green, blue);
                    }
                }
                println!("Connection closed");
            }
            Err(e) => {
                println!("Failed to accept connection from client: {:}", e);
            }
        }
    }

}

An example call:

$ echo -e "\x00\x01\xaa\xaa\xbb\xbb\xcc\xcc" | nc -I 1 -O 1 127.0.0.1 56501

Output

New connection ...
led count: 1
rgb: aaaa bbbb cccc
Connection closed

You can write struct which will own Buf<TcpStream> and necessary buffers, then all logic can be implemented in terms of methods, so your code will look like this.

UPD: updated code a bit

1 Like

Good compromise for reducing the scope of led_buf. I still hope I can get rid of it as it's only used locally for temporary purposes.

Hmm ... I just had an interesting idea :face_with_raised_eyebrow::

// read all bytes required to describe the values of the LED's channels
let mut led_bytes = &mut [0; MAX_BYTES_PER_LED][0..BYTES_PER_LED];
if let Err(_) = stream.read_exact(&mut led_bytes) { continue 'next_connection }

Can someone confirm, if this is heap-allocation-free?

Well, you need buffer to put your data in, so you have only two choices considering desire for buffer length changeable at runtime:

  1. large enough stack allocated buffer which will be bigger or equal to possible desirable buffer sizes
  2. heap allocated vector, which you'll need to somehow keep if you don't want unnecessary allocations (approach showed earlier one of solutions to make it nicier)

There is theoretically the third option, namely dynamically sized stack allocated buffers using alloca, but they can potentially have a negative impact on performance.

UPD:

Yes, it is.

That matches your first option.