736 Advanced

736-typestate-connection — Typestate Connection

Functional Programming

Tutorial

The Problem

Network connection objects are the canonical example of runtime state misuse: calling send() before connect() or after close() is a logic error that only manifests at runtime. Production code adds defensive if !self.connected { return Err(...) } checks everywhere. The typestate pattern eliminates these checks: send and recv only exist on TcpConn<Connected>, so calling them on a disconnected connection is a compile error. Used in production in the tokio-serial and embedded-hal crates.

🎯 Learning Outcomes

• Model TCP connection lifecycle as typestate: Disconnected, Connecting, Connected, Closed

• Enforce that send and recv are only callable in the Connected state

• Return Result from transition methods to handle IO errors without breaking the typestate invariant

• Track connection statistics (bytes_sent, bytes_recv) safely within the typed connection

• Understand how typestate composes with Result: connect() returns Result<TcpConn<Connected>, _>

Code Example

#![allow(clippy::all)]
/// 736: TCP Connection modelled as typestate
/// Send/recv only available on Connected; connect only on Disconnected.
use std::marker::PhantomData;

// ── State markers ─────────────────────────────────────────────────────────────
pub struct Disconnected;
pub struct Connecting;
pub struct Connected;
pub struct Closed;

// ── Connection ────────────────────────────────────────────────────────────────

pub struct TcpConn<State> {
    host: String,
    port: u16,
    // In a real impl, this would hold a socket fd
    bytes_sent: usize,
    bytes_recv: usize,
    _state: PhantomData<State>,
}

impl TcpConn<Disconnected> {
    pub fn new(host: impl Into<String>, port: u16) -> Self {
        TcpConn {
            host: host.into(),
            port,
            bytes_sent: 0,
            bytes_recv: 0,
            _state: PhantomData,
        }
    }

    /// Transition: Disconnected → Connected
    pub fn connect(self) -> Result<TcpConn<Connected>, String> {
        println!("Connecting to {}:{} ...", self.host, self.port);
        // In reality: TcpStream::connect(...)
        Ok(TcpConn {
            host: self.host,
            port: self.port,
            bytes_sent: 0,
            bytes_recv: 0,
            _state: PhantomData,
        })
    }
}

impl TcpConn<Connected> {
    /// Send data — only available when Connected.
    pub fn send(mut self, data: &[u8]) -> Result<TcpConn<Connected>, String> {
        println!("[{}:{}] → {} bytes", self.host, self.port, data.len());
        self.bytes_sent += data.len();
        Ok(self)
    }

    /// Receive data — only available when Connected.
    pub fn recv(mut self) -> Result<(Vec<u8>, TcpConn<Connected>), String> {
        let fake_data = b"HTTP/1.1 200 OK\r\n".to_vec();
        println!("[{}:{}] ← {} bytes", self.host, self.port, fake_data.len());
        self.bytes_recv += fake_data.len();
        Ok((fake_data, self))
    }

    /// Transition: Connected → Closed
    pub fn close(self) -> TcpConn<Closed> {
        println!(
            "Closing {}:{} (sent={}, recv={})",
            self.host, self.port, self.bytes_sent, self.bytes_recv
        );
        TcpConn {
            host: self.host,
            port: self.port,
            bytes_sent: self.bytes_sent,
            bytes_recv: self.bytes_recv,
            _state: PhantomData,
        }
    }

    pub fn peer(&self) -> String {
        format!("{}:{}", self.host, self.port)
    }
}

impl TcpConn<Closed> {
    pub fn bytes_sent(&self) -> usize {
        self.bytes_sent
    }
    pub fn bytes_recv(&self) -> usize {
        self.bytes_recv
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn connect_then_close() {
        let conn = TcpConn::<Disconnected>::new("localhost", 8080);
        let conn = conn.connect().unwrap();
        let closed = conn.close();
        assert_eq!(closed.bytes_sent(), 0);
        assert_eq!(closed.bytes_recv(), 0);
    }

    #[test]
    fn send_recv_accumulates_bytes() {
        let conn = TcpConn::<Disconnected>::new("localhost", 8080)
            .connect()
            .unwrap();
        let conn = conn.send(b"hello world").unwrap();
        let (_data, conn) = conn.recv().unwrap();
        let closed = conn.close();
        assert_eq!(closed.bytes_sent(), 11);
        assert!(closed.bytes_recv() > 0);
    }

    #[test]
    fn peer_returns_host_and_port() {
        let conn = TcpConn::<Disconnected>::new("example.com", 443)
            .connect()
            .unwrap();
        assert_eq!(conn.peer(), "example.com:443");
        conn.close();
    }
}

(* 736: TCP Connection as Typestate — OCaml simulation *)

type disconnected = Disconnected
type connecting   = Connecting
type connected    = Connected
type closed       = Closed

(* We simulate state using a phantom wrapper *)
type 'state conn = {
  host: string;
  port: int;
  state_label: string;
}

let new_conn host port : disconnected conn =
  { host; port; state_label = "Disconnected" }

let connect (c: disconnected conn) : connected conn =
  Printf.printf "Connecting to %s:%d...\n" c.host c.port;
  { c with state_label = "Connected" }

let send (c: connected conn) msg : connected conn =
  Printf.printf "[%s:%d] Sent: %s\n" c.host c.port msg;
  c

let recv (c: connected conn) : (string * connected conn) =
  Printf.printf "[%s:%d] Recv: <data>\n" c.host c.port;
  ("response", c)

let close (c: connected conn) : closed conn =
  Printf.printf "Closing connection to %s:%d\n" c.host c.port;
  { c with state_label = "Closed" }

let () =
  let conn = new_conn "example.com" 80 in
  let conn = connect conn in
  let conn = send conn "GET / HTTP/1.1" in
  let (_data, conn) = recv conn in
  let _closed = close conn in
  ()
  (* Invalid: send conn "another message" -- OCaml would flag this for closed conn *)

Key Differences

Move semantics: Rust's consuming self prevents double-use; OCaml's modules must use abstract types and hide constructors to achieve the same effect.

Error handling: Rust's Result<TcpConn<Connected>, E> combines typestate with error propagation naturally; OCaml uses result type or exceptions similarly.

Async integration: OCaml's Lwt_unix returns promise-wrapped connection types; Rust's tokio::net::TcpStream uses async fn with the same typestate ideas.

Runtime cost: Both approaches have zero runtime cost for the state tracking — no runtime enum, no branch.

OCaml Approach

OCaml models connection typestate using abstract types in separate modules. A Disconnected.t and Connected.t are distinct types exposed through module signatures. The connect : Disconnected.t -> (Connected.t, exn) result function enforces the transition. OCaml's Lwt and Async add monadic sequencing, making it natural to chain connect >>= send >>= recv >>= close.

Full Source

#![allow(clippy::all)]
/// 736: TCP Connection modelled as typestate
/// Send/recv only available on Connected; connect only on Disconnected.
use std::marker::PhantomData;

// ── State markers ─────────────────────────────────────────────────────────────
pub struct Disconnected;
pub struct Connecting;
pub struct Connected;
pub struct Closed;

// ── Connection ────────────────────────────────────────────────────────────────

pub struct TcpConn<State> {
    host: String,
    port: u16,
    // In a real impl, this would hold a socket fd
    bytes_sent: usize,
    bytes_recv: usize,
    _state: PhantomData<State>,
}

impl TcpConn<Disconnected> {
    pub fn new(host: impl Into<String>, port: u16) -> Self {
        TcpConn {
            host: host.into(),
            port,
            bytes_sent: 0,
            bytes_recv: 0,
            _state: PhantomData,
        }
    }

    /// Transition: Disconnected → Connected
    pub fn connect(self) -> Result<TcpConn<Connected>, String> {
        println!("Connecting to {}:{} ...", self.host, self.port);
        // In reality: TcpStream::connect(...)
        Ok(TcpConn {
            host: self.host,
            port: self.port,
            bytes_sent: 0,
            bytes_recv: 0,
            _state: PhantomData,
        })
    }
}

impl TcpConn<Connected> {
    /// Send data — only available when Connected.
    pub fn send(mut self, data: &[u8]) -> Result<TcpConn<Connected>, String> {
        println!("[{}:{}] → {} bytes", self.host, self.port, data.len());
        self.bytes_sent += data.len();
        Ok(self)
    }

    /// Receive data — only available when Connected.
    pub fn recv(mut self) -> Result<(Vec<u8>, TcpConn<Connected>), String> {
        let fake_data = b"HTTP/1.1 200 OK\r\n".to_vec();
        println!("[{}:{}] ← {} bytes", self.host, self.port, fake_data.len());
        self.bytes_recv += fake_data.len();
        Ok((fake_data, self))
    }

    /// Transition: Connected → Closed
    pub fn close(self) -> TcpConn<Closed> {
        println!(
            "Closing {}:{} (sent={}, recv={})",
            self.host, self.port, self.bytes_sent, self.bytes_recv
        );
        TcpConn {
            host: self.host,
            port: self.port,
            bytes_sent: self.bytes_sent,
            bytes_recv: self.bytes_recv,
            _state: PhantomData,
        }
    }

    pub fn peer(&self) -> String {
        format!("{}:{}", self.host, self.port)
    }
}

impl TcpConn<Closed> {
    pub fn bytes_sent(&self) -> usize {
        self.bytes_sent
    }
    pub fn bytes_recv(&self) -> usize {
        self.bytes_recv
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn connect_then_close() {
        let conn = TcpConn::<Disconnected>::new("localhost", 8080);
        let conn = conn.connect().unwrap();
        let closed = conn.close();
        assert_eq!(closed.bytes_sent(), 0);
        assert_eq!(closed.bytes_recv(), 0);
    }

    #[test]
    fn send_recv_accumulates_bytes() {
        let conn = TcpConn::<Disconnected>::new("localhost", 8080)
            .connect()
            .unwrap();
        let conn = conn.send(b"hello world").unwrap();
        let (_data, conn) = conn.recv().unwrap();
        let closed = conn.close();
        assert_eq!(closed.bytes_sent(), 11);
        assert!(closed.bytes_recv() > 0);
    }

    #[test]
    fn peer_returns_host_and_port() {
        let conn = TcpConn::<Disconnected>::new("example.com", 443)
            .connect()
            .unwrap();
        assert_eq!(conn.peer(), "example.com:443");
        conn.close();
    }
}

(* 736: TCP Connection as Typestate — OCaml simulation *)

type disconnected = Disconnected
type connecting   = Connecting
type connected    = Connected
type closed       = Closed

(* We simulate state using a phantom wrapper *)
type 'state conn = {
  host: string;
  port: int;
  state_label: string;
}

let new_conn host port : disconnected conn =
  { host; port; state_label = "Disconnected" }

let connect (c: disconnected conn) : connected conn =
  Printf.printf "Connecting to %s:%d...\n" c.host c.port;
  { c with state_label = "Connected" }

let send (c: connected conn) msg : connected conn =
  Printf.printf "[%s:%d] Sent: %s\n" c.host c.port msg;
  c

let recv (c: connected conn) : (string * connected conn) =
  Printf.printf "[%s:%d] Recv: <data>\n" c.host c.port;
  ("response", c)

let close (c: connected conn) : closed conn =
  Printf.printf "Closing connection to %s:%d\n" c.host c.port;
  { c with state_label = "Closed" }

let () =
  let conn = new_conn "example.com" 80 in
  let conn = connect conn in
  let conn = send conn "GET / HTTP/1.1" in
  let (_data, conn) = recv conn in
  let _closed = close conn in
  ()
  (* Invalid: send conn "another message" -- OCaml would flag this for closed conn *)

✓ Tests Rust test suite

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn connect_then_close() {
        let conn = TcpConn::<Disconnected>::new("localhost", 8080);
        let conn = conn.connect().unwrap();
        let closed = conn.close();
        assert_eq!(closed.bytes_sent(), 0);
        assert_eq!(closed.bytes_recv(), 0);
    }

    #[test]
    fn send_recv_accumulates_bytes() {
        let conn = TcpConn::<Disconnected>::new("localhost", 8080)
            .connect()
            .unwrap();
        let conn = conn.send(b"hello world").unwrap();
        let (_data, conn) = conn.recv().unwrap();
        let closed = conn.close();
        assert_eq!(closed.bytes_sent(), 11);
        assert!(closed.bytes_recv() > 0);
    }

    #[test]
    fn peer_returns_host_and_port() {
        let conn = TcpConn::<Disconnected>::new("example.com", 443)
            .connect()
            .unwrap();
        assert_eq!(conn.peer(), "example.com:443");
        conn.close();
    }
}

Exercises

Add a Reconnecting state between Closed and Connected with a reconnect() method that retries up to N times.

Implement a TlsConn<State> that wraps TcpConn<Connected> and adds a tls_handshake() transition to TlsConn<Secured> before allowing encrypted send/recv.

Write a generic Pipeline<C: Connected> that accepts any Connected connection type and sends a sequence of protocol messages, returning accumulated statistics.

Open Source Repos

functional-rust

View the source for this example on GitHub — OCaml and Rust side by side in the repo.

Rust