perf: optimize memory allocation with buffer pools and Cstruct

- Add TCP buffer pools to Protocol.t (64KB recv, 128KB decompress)
- Update LZW to work directly with Cstruct (zero-copy decompression)
- Add decompress_to_buffer and decompress_cstruct functions to LZW
- Add Cstruct-based codec functions (decode_compress_from_cstruct,
decode_push_pull_msg_cstruct)
- Update handle_tcp_connection to use buffer pools instead of per-
connection allocation
- Eliminate Cstruct.to_string conversions in hot path

Closes: swim-hrd

lib/codec.ml

···
               else ""
             in
             Ok (header, nodes, user_state))
+
+let decode_compress_from_cstruct (buf : Cstruct.t) :
+    (int * Cstruct.t, Types.decode_error) result =
+  let data = Cstruct.to_string buf in
+  let _, msgpack = Msgpck.String.read data in
+  match msgpack with
+  | Msgpck.Map fields -> (
+      let algo =
+        match List.assoc_opt (Msgpck.String "Algo") fields with
+        | Some (Msgpck.Int i) -> i
+        | Some (Msgpck.Int32 i) -> Int32.to_int i
+        | _ -> -1
+      in
+      let compressed_buf =
+        match List.assoc_opt (Msgpck.String "Buf") fields with
+        | Some (Msgpck.Bytes s) -> Some (Cstruct.of_string s)
+        | Some (Msgpck.String s) -> Some (Cstruct.of_string s)
+        | _ -> None
+      in
+      match compressed_buf with
+      | Some cs -> Ok (algo, cs)
+      | None -> Error (Types.Msgpack_error "missing Buf field"))
+  | _ -> Error (Types.Msgpack_error "expected map for compress")
+
+let decode_push_pull_msg_cstruct (buf : Cstruct.t) :
+    ( push_pull_header * push_node_state list * Cstruct.t,
+      Types.decode_error )
+    result =
+  if Cstruct.length buf < 1 then Error Types.Truncated_message
+  else
+    let data = Cstruct.to_string buf in
+    let header_size, header_msgpack = Msgpck.String.read data in
+    match decode_push_pull_header header_msgpack with
+    | Error e -> Error (Types.Msgpack_error e)
+    | Ok header -> (
+        let rec read_nodes offset remaining acc =
+          if remaining <= 0 then Ok (List.rev acc, offset)
+          else if offset >= String.length data then
+            Error Types.Truncated_message
+          else
+            let rest = String.sub data offset (String.length data - offset) in
+            let node_size, node_msgpack = Msgpck.String.read rest in
+            match decode_push_node_state node_msgpack with
+            | Error e -> Error (Types.Msgpack_error e)
+            | Ok node ->
+                read_nodes (offset + node_size) (remaining - 1) (node :: acc)
+        in
+        match read_nodes header_size header.pp_nodes [] with
+        | Error e -> Error e
+        | Ok (nodes, offset) ->
+            let user_state =
+              if header.pp_user_state_len > 0 && offset < Cstruct.length buf
+              then
+                Cstruct.sub buf offset
+                  (min header.pp_user_state_len (Cstruct.length buf - offset))
+              else Cstruct.empty
+            in
+            Ok (header, nodes, user_state))

lib/lzw.ml

···
 let max_dict_size = 1 lsl max_code_bits
 
 type bit_reader = {
-  data : string;
+  data : Cstruct.t;
   mutable pos : int;
-  mutable bit_pos : int;
   mutable bits_buf : int;
   mutable bits_count : int;
 }
 
-let make_bit_reader data =
-  { data; pos = 0; bit_pos = 0; bits_buf = 0; bits_count = 0 }
+let make_bit_reader data = { data; pos = 0; bits_buf = 0; bits_count = 0 }
 
 let read_bits_lsb reader n =
   while reader.bits_count < n do
-    if reader.pos >= String.length reader.data then raise (Failure "eof")
+    if reader.pos >= Cstruct.length reader.data then raise Exit
     else begin
-      let byte = Char.code reader.data.[reader.pos] in
+      let byte = Cstruct.get_uint8 reader.data reader.pos in
       reader.bits_buf <- reader.bits_buf lor (byte lsl reader.bits_count);
       reader.bits_count <- reader.bits_count + 8;
       reader.pos <- reader.pos + 1
···
   reader.bits_count <- reader.bits_count - n;
   result
 
-let decompress ?(order = LSB) ?(lit_width = 8) data =
-  if order <> LSB then Error (Invalid_code 0)
-  else if lit_width <> 8 then Error (Invalid_code 0)
-  else
-    try
-      let reader = make_bit_reader data in
-      let output = Buffer.create (String.length data * 2) in
+let decompress_to_buffer ~src ~dst =
+  try
+    let reader = make_bit_reader src in
+    let out_pos = ref 0 in
+    let dst_len = Cstruct.length dst in
 
-      let dict = Array.make max_dict_size "" in
-      for i = 0 to 255 do
-        dict.(i) <- String.make 1 (Char.chr i)
-      done;
-      dict.(clear_code) <- "";
-      dict.(eof_code) <- "";
+    let dict = Array.make max_dict_size (Cstruct.empty, 0) in
+    for i = 0 to 255 do
+      dict.(i) <- (Cstruct.of_string (String.make 1 (Char.chr i)), 1)
+    done;
+    dict.(clear_code) <- (Cstruct.empty, 0);
+    dict.(eof_code) <- (Cstruct.empty, 0);
 
-      let dict_size = ref initial_dict_size in
-      let code_bits = ref 9 in
-      let prev_string = ref "" in
+    let dict_size = ref initial_dict_size in
+    let code_bits = ref 9 in
+    let prev_code = ref (-1) in
 
-      let add_to_dict s =
-        if !dict_size < max_dict_size then begin
-          dict.(!dict_size) <- s;
-          incr dict_size;
-          if !dict_size >= 1 lsl !code_bits && !code_bits < max_code_bits then
-            incr code_bits
-        end
-      in
+    let write_entry (entry, len) =
+      if !out_pos + len > dst_len then raise (Failure "overflow");
+      Cstruct.blit entry 0 dst !out_pos len;
+      out_pos := !out_pos + len
+    in
+
+    let add_to_dict first_byte =
+      if !dict_size < max_dict_size && !prev_code >= 0 then begin
+        let prev_entry, prev_len = dict.(!prev_code) in
+        let new_entry = Cstruct.create (prev_len + 1) in
+        Cstruct.blit prev_entry 0 new_entry 0 prev_len;
+        Cstruct.set_uint8 new_entry prev_len first_byte;
+        dict.(!dict_size) <- (new_entry, prev_len + 1);
+        incr dict_size;
+        if !dict_size >= 1 lsl !code_bits && !code_bits < max_code_bits then
+          incr code_bits
+      end
+    in
+
+    let reset_dict () =
+      dict_size := initial_dict_size;
+      code_bits := 9;
+      prev_code := -1
+    in
+
+    let rec decode_loop () =
+      let code = read_bits_lsb reader !code_bits in
+      if code = eof_code then ()
+      else if code = clear_code then begin
+        reset_dict ();
+        decode_loop ()
+      end
+      else begin
+        let entry, len, first_byte =
+          if code < !dict_size then
+            let e, l = dict.(code) in
+            (e, l, Cstruct.get_uint8 e 0)
+          else if code = !dict_size && !prev_code >= 0 then (
+            let prev_entry, prev_len = dict.(!prev_code) in
+            let first = Cstruct.get_uint8 prev_entry 0 in
+            let new_entry = Cstruct.create (prev_len + 1) in
+            Cstruct.blit prev_entry 0 new_entry 0 prev_len;
+            Cstruct.set_uint8 new_entry prev_len first;
+            (new_entry, prev_len + 1, first))
+          else raise (Failure "invalid")
+        in
+        write_entry (entry, len);
+        add_to_dict first_byte;
+        prev_code := code;
+        decode_loop ()
+      end
+    in
 
-      let reset_dict () =
-        dict_size := initial_dict_size;
-        code_bits := 9;
-        prev_string := ""
-      in
+    decode_loop ();
+    Ok !out_pos
+  with
+  | Exit -> Error Unexpected_eof
+  | Failure msg when msg = "overflow" -> Error Buffer_overflow
+  | Failure msg when msg = "invalid" -> Error (Invalid_code 0)
+  | _ -> Error (Invalid_code 0)
 
-      let rec decode_loop () =
-        let code = read_bits_lsb reader !code_bits in
-        if code = eof_code then ()
-        else if code = clear_code then begin
-          reset_dict ();
-          decode_loop ()
-        end
-        else begin
-          let current_string =
-            if code < !dict_size then dict.(code)
-            else if code = !dict_size then
-              !prev_string ^ String.make 1 !prev_string.[0]
-            else
-              raise
-                (Failure
-                   (Printf.sprintf "invalid code %d >= %d" code !dict_size))
-          in
-          Buffer.add_string output current_string;
-          if !prev_string <> "" then
-            add_to_dict (!prev_string ^ String.make 1 current_string.[0]);
-          prev_string := current_string;
-          decode_loop ()
-        end
-      in
+let decompress_cstruct src =
+  let estimated_size = max (Cstruct.length src * 4) 4096 in
+  let dst = Cstruct.create estimated_size in
+  match decompress_to_buffer ~src ~dst with
+  | Ok len -> Ok (Cstruct.sub dst 0 len)
+  | Error Buffer_overflow -> (
+      let larger = Cstruct.create (estimated_size * 4) in
+      match decompress_to_buffer ~src ~dst:larger with
+      | Ok len -> Ok (Cstruct.sub larger 0 len)
+      | Error e -> Error e)
+  | Error e -> Error e
 
-      decode_loop ();
-      Ok (Buffer.contents output)
-    with
-    | Failure msg when msg = "eof" -> Error Unexpected_eof
-    | Failure msg
-      when String.length msg > 12 && String.sub msg 0 12 = "invalid code" ->
-        Error (Invalid_code 0)
-    | _ -> Error (Invalid_code 0)
+let decompress ?(order = LSB) ?(lit_width = 8) data =
+  if order <> LSB then Error (Invalid_code 0)
+  else if lit_width <> 8 then Error (Invalid_code 0)
+  else
+    let src = Cstruct.of_string data in
+    match decompress_cstruct src with
+    | Ok cs -> Ok (Cstruct.to_string cs)
+    | Error e -> Error e
 
 let decompress_lsb8 data = decompress ~order:LSB ~lit_width:8 data

lib/lzw.mli

···
 type error = Invalid_code of int | Unexpected_eof | Buffer_overflow
 
 val error_to_string : error -> string
+val decompress_to_buffer : src:Cstruct.t -> dst:Cstruct.t -> (int, error) result
+val decompress_cstruct : Cstruct.t -> (Cstruct.t, error) result
 
 val decompress :
   ?order:order -> ?lit_width:int -> string -> (string, error) result

lib/protocol.ml

···
   probe_index : int Kcas.Loc.t;
   send_pool : Buffer_pool.t;
   recv_pool : Buffer_pool.t;
+  tcp_recv_pool : Buffer_pool.t;
+  tcp_decompress_pool : Buffer_pool.t;
   udp_sock : [ `Generic ] Eio.Net.datagram_socket_ty Eio.Resource.t;
   tcp_listener : [ `Generic ] Eio.Net.listening_socket_ty Eio.Resource.t;
   event_stream : node_event Eio.Stream.t;
···
       else None
   | _ -> None
 
+let decompress_payload_cstruct ~src ~dst =
+  match Codec.decode_compress_from_cstruct src with
+  | Error _ -> None
+  | Ok (algo, compressed) ->
+      if algo = 0 then
+        match Lzw.decompress_to_buffer ~src:compressed ~dst with
+        | Ok len -> Some len
+        | Error _ -> None
+      else None
+
 let handle_tcp_connection t flow =
-  let buf = Cstruct.create 65536 in
-  match read_exact flow buf 1 with
-  | Error _ -> ()
-  | Ok () -> (
-      let msg_type_byte = Cstruct.get_uint8 buf 0 in
-      let get_push_pull_payload () =
-        let n = read_available flow (Cstruct.shift buf 1) in
-        if n > 0 then Some (Cstruct.sub buf 1 n) else None
-      in
-      let payload_opt =
-        if msg_type_byte = Types.Wire.message_type_to_int Types.Wire.Encrypt_msg
-        then
-          match get_push_pull_payload () with
-          | Some encrypted -> (
-              match Crypto.decrypt ~key:t.cipher_key encrypted with
-              | Ok decrypted -> Some decrypted
-              | Error _ -> None)
-          | None -> None
-        else if
-          msg_type_byte = Types.Wire.message_type_to_int Types.Wire.Compress_msg
-        then
-          match get_push_pull_payload () with
-          | Some compressed -> (
-              let data = Cstruct.to_string compressed in
-              match decompress_payload data with
-              | Some decompressed ->
-                  if String.length decompressed > 0 then
-                    let inner_type = Char.code decompressed.[0] in
-                    if
-                      inner_type
-                      = Types.Wire.message_type_to_int Types.Wire.Push_pull_msg
-                    then
-                      Some
-                        (Cstruct.of_string
-                           (String.sub decompressed 1
-                              (String.length decompressed - 1)))
-                    else None
-                  else None
-              | None -> None)
-          | None -> None
-        else if
-          msg_type_byte
-          = Types.Wire.message_type_to_int Types.Wire.Has_label_msg
-        then
+  Buffer_pool.with_buffer t.tcp_recv_pool (fun buf ->
+      Buffer_pool.with_buffer t.tcp_decompress_pool (fun decomp_buf ->
           match read_exact flow buf 1 with
-          | Error _ -> None
-          | Ok () ->
-              let label_len = Cstruct.get_uint8 buf 0 in
-              if label_len > 0 then
-                match read_exact flow buf label_len with
-                | Error _ -> None
-                | Ok () -> (
-                    match read_exact flow buf 1 with
-                    | Error _ -> None
-                    | Ok () ->
-                        let inner_type = Cstruct.get_uint8 buf 0 in
-                        if
-                          inner_type
-                          = Types.Wire.message_type_to_int
-                              Types.Wire.Push_pull_msg
-                        then get_push_pull_payload ()
-                        else None)
-              else None
-        else if
-          msg_type_byte
-          = Types.Wire.message_type_to_int Types.Wire.Push_pull_msg
-        then get_push_pull_payload ()
-        else None
-      in
-      match payload_opt with
-      | None -> ()
-      | Some payload -> (
-          let data = Cstruct.to_string payload in
-          match Codec.decode_push_pull_msg data with
           | Error _ -> ()
-          | Ok (header, nodes, _user_state) -> (
-              merge_remote_state t nodes ~is_join:header.pp_join;
-              let resp_header, resp_nodes =
-                build_local_state t ~is_join:false
-              in
-              let response =
-                Codec.encode_push_pull_msg ~header:resp_header ~nodes:resp_nodes
-                  ~user_state:""
+          | Ok () -> (
+              let msg_type_byte = Cstruct.get_uint8 buf 0 in
+              let get_push_pull_payload () =
+                let n = read_available flow (Cstruct.shift buf 1) in
+                if n > 0 then Some (Cstruct.sub buf 1 n) else None
               in
-              let resp_buf =
-                if t.config.encryption_enabled then
-                  let plain = Cstruct.of_string response in
-                  let encrypted =
-                    Crypto.encrypt ~key:t.cipher_key ~random:t.secure_random
-                      plain
-                  in
-                  encrypted
-                else Cstruct.of_string response
+              let payload_opt =
+                if
+                  msg_type_byte
+                  = Types.Wire.message_type_to_int Types.Wire.Encrypt_msg
+                then
+                  match get_push_pull_payload () with
+                  | Some encrypted -> (
+                      match Crypto.decrypt ~key:t.cipher_key encrypted with
+                      | Ok decrypted -> Some decrypted
+                      | Error _ -> None)
+                  | None -> None
+                else if
+                  msg_type_byte
+                  = Types.Wire.message_type_to_int Types.Wire.Compress_msg
+                then
+                  match get_push_pull_payload () with
+                  | Some compressed -> (
+                      match
+                        decompress_payload_cstruct ~src:compressed
+                          ~dst:decomp_buf
+                      with
+                      | Some len ->
+                          if len > 0 then
+                            let inner_type = Cstruct.get_uint8 decomp_buf 0 in
+                            if
+                              inner_type
+                              = Types.Wire.message_type_to_int
+                                  Types.Wire.Push_pull_msg
+                            then Some (Cstruct.sub decomp_buf 1 (len - 1))
+                            else None
+                          else None
+                      | None -> None)
+                  | None -> None
+                else if
+                  msg_type_byte
+                  = Types.Wire.message_type_to_int Types.Wire.Has_label_msg
+                then
+                  match read_exact flow buf 1 with
+                  | Error _ -> None
+                  | Ok () ->
+                      let label_len = Cstruct.get_uint8 buf 0 in
+                      if label_len > 0 then
+                        match read_exact flow buf label_len with
+                        | Error _ -> None
+                        | Ok () -> (
+                            match read_exact flow buf 1 with
+                            | Error _ -> None
+                            | Ok () ->
+                                let inner_type = Cstruct.get_uint8 buf 0 in
+                                if
+                                  inner_type
+                                  = Types.Wire.message_type_to_int
+                                      Types.Wire.Push_pull_msg
+                                then get_push_pull_payload ()
+                                else None)
+                      else None
+                else if
+                  msg_type_byte
+                  = Types.Wire.message_type_to_int Types.Wire.Push_pull_msg
+                then get_push_pull_payload ()
+                else None
               in
-              try Eio.Flow.write flow [ resp_buf ] with _ -> ())))
+              match payload_opt with
+              | None -> ()
+              | Some payload -> (
+                  match Codec.decode_push_pull_msg_cstruct payload with
+                  | Error _ -> ()
+                  | Ok (header, nodes, _user_state) -> (
+                      merge_remote_state t nodes ~is_join:header.pp_join;
+                      let resp_header, resp_nodes =
+                        build_local_state t ~is_join:false
+                      in
+                      let response =
+                        Codec.encode_push_pull_msg ~header:resp_header
+                          ~nodes:resp_nodes ~user_state:""
+                      in
+                      let resp_buf =
+                        if t.config.encryption_enabled then
+                          let plain = Cstruct.of_string response in
+                          let encrypted =
+                            Crypto.encrypt ~key:t.cipher_key
+                              ~random:t.secure_random plain
+                          in
+                          encrypted
+                        else Cstruct.of_string response
+                      in
+                      try Eio.Flow.write flow [ resp_buf ] with _ -> ())))))
 
 let run_tcp_listener t =
   while not (is_shutdown t) do
···
           recv_pool =
             Buffer_pool.create ~size:config.udp_buffer_size
               ~count:config.recv_buffer_count;
+          tcp_recv_pool = Buffer_pool.create ~size:65536 ~count:4;
+          tcp_decompress_pool = Buffer_pool.create ~size:131072 ~count:4;
           udp_sock;
           tcp_listener;
           event_stream = Eio.Stream.create 100;