🎓 A1Learner
🔍
0 of 45 lessons completed

Rust Tokio Connection

A Connection is a type that owns a socket and its buffer. A Connection offers two methods, read_frame and write_frame.

The struct

pub struct Connection {
    stream: BufWriter<TcpStream>,
    buffer: BytesMut,
}

pub fn new(stream: TcpStream) -> Connection {
    Connection {
        stream: BufWriter::new(stream),
        buffer: BytesMut::with_capacity(4 * 1024),
    }
}

Parse one line-frame from the buffer

Our frame shape is one line of text. No newline yet means: not a full frame.

fn parse_frame(&mut self) -> Option<String> {
    let newline = self.buffer.iter().position(|b| *b == b'\n')?;
    let mut line = self.buffer.split_to(newline + 1);
    line.truncate(newline);
    Some(String::from_utf8(line.to_vec()).unwrap())
}

Read until a whole frame exists

loop {
    if let Some(frame) = self.parse_frame() {
        return Ok(Some(frame));
    }
    if 0 == self.stream.get_mut().read_buf(&mut self.buffer).await? {
        if self.buffer.is_empty() {
            return Ok(None); // clean close
        } else {
            return Err("connection reset by peer".into());
        }
    }
}

Write a frame and flush

pub async fn write_frame(&mut self, frame: &str) -> io::Result<()> {
    self.stream.write_all(frame.as_bytes()).await?;
    self.stream.write_all(b"\n").await?;
    self.stream.flush().await?; // buffered writes leave only on flush
    Ok(())
}

A client is a Connection plus commands

The tutorial's client wraps a Connection and turns each command into a write-then-read pair:

struct Client {
    conn: Connection,
}

async fn get(&mut self, key: &str) -> Option<String> {
    self.conn.write_frame(&format!("get {}", key)).await.ok()?;
    self.conn.read_frame().await.ok()?
}

Example

// Cargo.toml: tokio = { version = "1", features = ["full"] }, bytes = "1"
use bytes::BytesMut;
use std::io;
use tokio::io::{AsyncReadExt, AsyncWriteExt, BufWriter};
use tokio::net::{TcpListener, TcpStream};

pub struct Connection {
    stream: BufWriter<TcpStream>,
    buffer: BytesMut,
}

impl Connection {
    pub fn new(stream: TcpStream) -> Connection {
        Connection {
            stream: BufWriter::new(stream),
            buffer: BytesMut::with_capacity(1024),
        }
    }

    fn parse_frame(&mut self) -> Option<String> {
        let newline = self.buffer.iter().position(|b| *b == b'\n')?;
        let mut line = self.buffer.split_to(newline + 1);
        line.truncate(newline);
        if line.last() == Some(&b'\r') {
            line.truncate(line.len() - 1);
        }
        Some(String::from_utf8(line.to_vec()).unwrap())
    }

    pub async fn read_frame(&mut self) -> io::Result<Option<String>> {
        loop {
            if let Some(frame) = self.parse_frame() {
                return Ok(Some(frame));
            }
            if 0 == self.stream.get_mut().read_buf(&mut self.buffer).await? {
                if self.buffer.is_empty() {
                    return Ok(None);
                } else {
                    return Err(io::Error::new(
                        io::ErrorKind::InvalidData,
                        "connection reset by peer",
                    ));
                }
            }
        }
    }

    pub async fn write_frame(&mut self, frame: &str) -> io::Result<()> {
        self.stream.write_all(frame.as_bytes()).await?;
        self.stream.write_all(b"\n").await?;
        self.stream.flush().await?;
        Ok(())
    }
}

#[tokio::main]
async fn main() -> io::Result<()> {
    let listener = TcpListener::bind("127.0.0.1:0").await?;
    let addr = listener.local_addr()?;

    let server = tokio::spawn(async move {
        let (socket, _) = listener.accept().await.unwrap();
        let mut conn = Connection::new(socket);
        let frame = conn.read_frame().await.unwrap().unwrap();
        println!("server received: {}", frame);
        conn.write_frame("pong").await.unwrap();
    });

    let stream = TcpStream::connect(addr).await?;
    let mut client = Connection::new(stream);
    client.write_frame("ping").await?;
    let reply = client.read_frame().await?.unwrap();
    println!("client received: {}", reply);

    server.await.unwrap();
    Ok(())
}

Key points

  • Connection = a TcpStream wrapped in a BufWriter, plus a BytesMut read buffer: one struct, both directions
  • read_frame returns Ok(Some(frame)) for a whole frame, Ok(None) when the peer closed cleanly, and Err when the peer vanished mid-frame
  • A Connection is used through &mut self.
  • So exactly one task drives it.
  • That is why servers spawn one task per connection.
  • The client side uses the same type.
  • A client struct holds a Connection.
  • It offers one async method per command.
  • Each method writes a command frame and reads one response frame.
  • Keep per-connection errors inside the connection's task.
  • Log them and end that task.
  • Never panic.
  • One broken client must not stop the server.
  • The tutorial gives the buffer 4 KB of starting capacity.
  • That is big enough for most frames.

In this example, the server binds port 0. The system picks a free port. Nothing already running can clash with it. The accepted socket is wrapped in a Connection. read_frame waits inside its loop. It reads chunks into the buffer until one whole line-frame has arrived. It prints ping, then write_frame sends pong back as one buffered write plus a flush. The client side wraps its stream in the very same Connection type: it writes ping and reads one frame back. One struct, both ends of the wire — that is the whole pattern.

Note

BytesMut grows it when a bigger frame arrives.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Rust Playground

Exercise

read_frame returns Ok(None). What does that mean?