Fix Wsd

lib/wsd.ml

@@ -1,20 +1,76 @@
-type t = Faraday.t
+module IOVec = Httpaf.IOVec
 
-let schedule ?mask t ~kind payload ~off ~len =
-  Websocket.Frame.schedule_serialize ?mask t ~is_fin:true ~opcode:kind ~payload ~off ~len
+type mode =
+  [ `Client of unit -> int32
+  | `Server
+  ]
 
-let send_bytes ?mask t ~kind payload ~off ~len =
-  Websocket.Frame.schedule_serialize_bytes ?mask t ~is_fin:true ~opcode:kind ~payload ~off ~len
+type t =
+  { faraday : Faraday.t
+  ; mode : mode
+  ; mutable when_ready_to_write : unit -> unit
+  }
+
+let default_ready_to_write = Sys.opaque_identity (fun () -> ())
+
+let create mode =
+  { faraday = Faraday.create 0x1000
+  ; mode
+  ; when_ready_to_write = default_ready_to_write;
+  }
+
+let mask t =
+  match t.mode with
+  | `Client m -> Some (m ())
+  | `Server -> None
+
+let ready_to_write t =
+  let callback = t.when_ready_to_write in
+  t.when_ready_to_write <- default_ready_to_write;
+  callback ()
+
+let schedule t ~kind payload ~off ~len =
+  let mask = mask t in
+  Websocket.Frame.schedule_serialize t.faraday ?mask ~is_fin:true ~opcode:(kind :> Websocket.Opcode.t) ~payload ~off ~len;
+  ready_to_write t
+
+let send_bytes t ~kind payload ~off ~len =
+  let mask = mask t in
+  Websocket.Frame.schedule_serialize_bytes t.faraday ?mask ~is_fin:true ~opcode:(kind :> Websocket.Opcode.t) ~payload ~off ~len;
+  ready_to_write t
 
 let send_ping t =
-  Websocket.Frame.serialize_control t ~opcode:`Ping
+  Websocket.Frame.serialize_control t.faraday ~opcode:`Ping;
+  ready_to_write t
 
 let send_pong t =
-  Websocket.Frame.serialize_control t ~opcode:`Pong
+  Websocket.Frame.serialize_control t.faraday ~opcode:`Pong;
+  ready_to_write t
 
-let flushed t f = Faraday.flush t f
+let flushed t f = Faraday.flush t.faraday f
 
 let close t =
-  Websocket.Frame.serialize_control t ~opcode:`Connection_close;
-  Faraday.close t
-;;
+  Websocket.Frame.serialize_control t.faraday ~opcode:`Connection_close;
+  Faraday.close t.faraday;
+  ready_to_write t
+
+let next t =
+  match Faraday.operation t.faraday with
+  | `Close         -> `Close 0 (* XXX(andreas): should track unwritten bytes *)
+  | `Yield         -> `Yield
+  | `Writev iovecs -> `Write iovecs
+
+let report_result t result =
+  match result with
+  | `Closed -> close t
+  | `Ok len -> Faraday.shift t.faraday len
+
+let is_closed t =
+  Faraday.is_closed t.faraday
+
+let when_ready_to_write t callback =
+  if not (t.when_ready_to_write == default_ready_to_write)
+  then failwith "Wsd.when_ready_to_write: only one callback can be registered at a time"
+  else if is_closed t
+  then callback ()
+  else t.when_ready_to_write <- callback

lib/wsd.mli

@@ -1,17 +1,26 @@
+module IOVec = Httpaf.IOVec
+
+type mode =
+  [ `Client of unit -> int32
+  | `Server
+  ]
+
 type t 
 
-val schedule
-  :  ?mask:int32
+val create
+  : mode
   -> t
+
+val schedule
+  :  t
   -> kind:[ `Text | `Binary ]
   -> Bigstring.t
   -> off:int
   -> len:int
   -> unit
 
 val send_bytes
-  :  ?mask:int32
-  -> t
+  :  t
   -> kind:[ `Text | `Binary ]
   -> Bytes.t
   -> off:int
@@ -23,3 +32,10 @@ val send_pong : t -> unit
 
 val flushed : t -> (unit -> unit) -> unit
 val close   : t -> unit
+
+val next : t -> [ `Write of Bigstring.t IOVec.t list | `Yield | `Close of int ]
+val report_result : t -> [`Ok of int | `Closed ] -> unit
+
+val is_closed : t -> bool
+
+val when_ready_to_write : t -> (unit -> unit) -> unit