Source file pipe_transport.ml

open Core
open Async_kernel
open Int.Replace_polymorphic_compare

include struct
  open Transport
  module Header = Header
  module Handler_result = Handler_result
  module Send_result = Send_result
  module Reader = Reader
  module Writer = Writer
end

let rec pipe_read t f =
  match%bind Pipe.values_available t with
  | `Eof -> f `Eof
  | `Ok ->
    (match Pipe.read_now' t with
     | `Nothing_available -> pipe_read t f
     | (`Eof | `Ok _) as res -> f res)
;;

module type DATA = sig
  type t [@@deriving sexp_of]

  val length : t -> int
  val add_to_bigbuffer : Bigbuffer.t -> t -> unit
  val to_bigstring : t -> Bigstring.t
  val of_bigstring : Bigstring.t -> t
end

module Pipe_and_buffer = struct
  type 'a t =
    { pipe : 'a Pipe.Reader.t
    ; buffer : Bigbuffer.t
    }
  [@@deriving sexp_of]

  let create pipe = { pipe; buffer = Bigbuffer.create Header.length }
end

module Pipe_reader (Data : DATA) = struct
  type t = Data.t Pipe_and_buffer.t [@@deriving sexp_of]

  let close (t : t) =
    Pipe.close_read t.pipe;
    Deferred.unit
  ;;

  let is_closed (t : t) = Pipe.is_closed t.pipe

  let read_forever (t : t) ~on_message ~on_end_of_batch : (_, _) Deferred.Result.t =
    let buffer = t.buffer in
    (* Dequeue until enough data ([need]) is available and try read messages *)
    let rec process_queue ~need queue =
      match Queue.dequeue queue with
      | None ->
        (* Queue is empty - we need to wait for more data *)
        return (Second need)
      | Some data ->
        let buff_len = Bigbuffer.length buffer in
        let data_len = Data.length data in
        if buff_len = 0 && data_len >= need
        then
          (* We have enough data and buffer is empty.
             This is the common case where every bin_prot message is sent in its own
             websocket frame *)
          process_data ~pos:0 ~data:(Data.to_bigstring data) ~length:data_len queue
        else (
          (* Slow path. We have to accumulate [Bigstring.t]s *)
          Data.add_to_bigbuffer buffer data;
          if Bigbuffer.length buffer >= need
          then (
            (* There is enough data in the buffer *)
            let data = Bigbuffer.volatile_contents buffer in
            let length = Bigbuffer.length buffer in
            process_data ~pos:0 ~data ~length queue)
          else
            (* We need more data - keep processing the queue *)
            process_queue ~need queue)
    (* Read as many messages as possible from the current data *)
    and process_data ~pos ~data ~length queue =
      if length >= Header.length
      then (
        let payload_len = Header.unsafe_get_payload_length data ~pos in
        let total_len = Header.length + payload_len in
        if length >= total_len
        then (
          match on_message data ~pos:(pos + Header.length) ~len:payload_len with
          | Handler_result.Stop x ->
            let pos = pos + total_len in
            let length = length - total_len in
            (* Make sure that all data we've read (and not "consumed") from the pipe
               is kept in the buffer.
               That is the remaining bits in [data] and all contents in [queue].
            *)
            set_buffer ~pos ~length ~data;
            Queue.iter queue ~f:(fun data -> Data.add_to_bigbuffer buffer data);
            return (First x)
          | (Continue | Wait _) as continue_or_wait ->
            let%bind () =
              match continue_or_wait with
              | Handler_result.Stop _ ->
                (* Impossible: this branch does not contain [Stop _] *)
                assert false
              | Continue -> Deferred.unit
              | Wait d -> d
            in
            let pos = pos + total_len in
            let length = length - total_len in
            assert (length >= 0);
            if length > 0
            then
              (* Process the rest of the data.
                 More than one bin_prot message was sent inside the websocket frame *)
              process_data ~pos ~data ~length queue
            else
              (* length = 0 *)
              (* All data was read - keep processing the queue *)
              set_buffer_and_process_queue ~need:Header.length ~pos ~length:0 ~queue ~data)
        else set_buffer_and_process_queue ~need:total_len ~pos ~length ~queue ~data)
      else set_buffer_and_process_queue ~need:Header.length ~pos ~length ~queue ~data
    and set_buffer ~pos ~length ~data =
      if length = 0
      then Bigbuffer.clear buffer
      else (
        let data = Bigstring.sub data ~pos ~len:length in
        Bigbuffer.clear buffer;
        Bigbuffer.add_bigstring buffer data)
    and set_buffer_and_process_queue ~need ~pos ~length ~queue ~data =
      set_buffer ~pos ~length ~data;
      process_queue ~need queue
    in
    let rec wait_for_read ~need =
      pipe_read t.pipe (function
        | `Eof -> return (Error `Eof)
        | `Ok queue -> process_queue ~need queue >>= end_of_batch_and_continue)
    and end_of_batch_and_continue next =
      on_end_of_batch ();
      match next with
      | First result -> return (Ok result)
      | Second need -> wait_for_read ~need
    in
    let start_processing_existing_data ~data ~length =
      process_data ~pos:0 ~data ~length (Queue.create ()) >>= end_of_batch_and_continue
    in
    (* We either start by
       - processing remaining bits in the buffer
       - or waiting for more data in the pipe *)
    let need = Header.length in
    let length = Bigbuffer.length buffer in
    if Int.( < ) length need
    then wait_for_read ~need
    else start_processing_existing_data ~length ~data:(Bigbuffer.volatile_contents buffer)
  ;;

  module For_testing = struct
    type t = Data.t

    let of_bigstring = Data.of_bigstring
  end
end

module Pipe_and_monitor = struct
  type 'a t =
    { pipe : 'a Pipe.Writer.t
    ; monitor : Monitor.t
    }
  [@@deriving sexp_of]

  let create pipe = { pipe; monitor = Monitor.create () }
end

(* We don't perform any buffering here.
   A message is consider to be flushed as soon as it enters the pipe. *)
module Pipe_writer (Data : DATA) = struct
  type t = Data.t Pipe_and_monitor.t [@@deriving sexp_of]

  let close (t : t) =
    Pipe.close t.pipe;
    Deferred.unit
  ;;

  let is_closed (t : t) = Pipe.is_closed t.pipe
  let monitor (t : t) = t.monitor

  (* Because we don't maintain any buffer, there are no pending writes *)
  let bytes_to_write (_ : t) = 0
  let stopped (t : t) = Pipe.closed t.pipe

  (* We consider that a message is flushed as soon as it reaches the underlining
     transport. *)
  let flushed (_ : t) = Deferred.unit
  let ready_to_write = flushed
  let sent_result x : _ Send_result.t = Sent x

  let check_closed (t : t) f =
    if not (Pipe.is_closed t.pipe) then f () else Send_result.Closed
  ;;

  let send_bin_prot t writer x =
    check_closed t (fun () ->
      let buf = Bin_prot.Utils.bin_dump ~header:true writer x in
      Pipe.write_without_pushback t.pipe (Data.of_bigstring buf);
      sent_result ())
  ;;

  let send_bin_prot_and_bigstring
        t
        (writer : _ Bin_prot.Type_class.writer)
        x
        ~buf
        ~pos
        ~len:payload_size
    =
    check_closed t (fun () ->
      (* Write the size header manually and concatenate the two *)
      let data_size = writer.size x in
      let data = Bigstring.create (data_size + Header.length + payload_size) in
      Header.unsafe_set_payload_length data ~pos:0 (data_size + payload_size);
      let dst_pos = writer.write data ~pos:Header.length x in
      Bigstring.blit ~src:buf ~src_pos:pos ~dst:data ~dst_pos ~len:payload_size;
      Pipe.write_without_pushback t.pipe (Data.of_bigstring data);
      sent_result ())
  ;;

  let send_bin_prot_and_bigstring_non_copying t writer x ~buf ~pos ~len =
    match send_bin_prot_and_bigstring t writer x ~buf ~pos ~len with
    | Sent () -> sent_result Deferred.unit
    | (Closed | Message_too_big _) as r -> r
  ;;
end

module Bigstring_data = struct
  type t = Bigstring.t [@@deriving sexp_of]

  let length = Bigstring.length
  let add_to_bigbuffer = Bigbuffer.add_bigstring
  let to_bigstring x = x
  let of_bigstring x = x
end

module String_data = struct
  type t = String.t [@@deriving sexp_of]

  let length = String.length
  let add_to_bigbuffer = Bigbuffer.add_string
  let to_bigstring x = Bigstring.of_string x
  let of_bigstring x = Bigstring.to_string x
end

module Bigstring_pipe_reader = Pipe_reader (Bigstring_data)
module Bigstring_pipe_writer = Pipe_writer (Bigstring_data)
module String_pipe_reader = Pipe_reader (String_data)
module String_pipe_writer = Pipe_writer (String_data)

module Kind = struct
  type 'a t =
    | String : string t
    | Bigstring : bigstring t

  let string = String
  let bigstring = Bigstring
end

let make_reader (type a) (x : a Kind.t) (reader : a Pipe.Reader.t) =
  let reader = Pipe_and_buffer.create reader in
  match x with
  | Kind.String -> Reader.pack (module String_pipe_reader) reader
  | Kind.Bigstring -> Reader.pack (module Bigstring_pipe_reader) reader
;;

let make_writer (type a) (x : a Kind.t) (writer : a Pipe.Writer.t) =
  let writer = Pipe_and_monitor.create writer in
  match x with
  | Kind.String -> Writer.pack (module String_pipe_writer) writer
  | Kind.Bigstring -> Writer.pack (module Bigstring_pipe_writer) writer
;;

let create kind reader writer =
  { Transport.reader = make_reader kind reader; writer = make_writer kind writer }
;;

(* Testing *)

module type Transport_reader = sig
  module For_testing : sig
    type t

    val of_bigstring : Bigstring.t -> t
  end

  include Transport.Reader.S with type t = For_testing.t Pipe_and_buffer.t
end

module Test_reader (Transport_reader : Transport_reader) = struct
  open Transport_reader.For_testing

  let default_message = Bigstring.of_string "FOOBAR123"

  let create_message_buf (message : Bigstring.t) : Bigstring.t =
    let message_len = Bigstring.length message in
    let buf = Bigstring.create (Header.length + message_len) in
    Header.unsafe_set_payload_length buf ~pos:0 message_len;
    Bigstring.blit
      ~src:message
      ~dst:buf
      ~src_pos:0
      ~dst_pos:Header.length
      ~len:message_len;
    buf
  ;;

  let create_message () = create_message_buf default_message

  let on_message ?(message = default_message) on_message_count buffer ~pos ~len =
    on_message_count := !on_message_count + 1;
    [%test_result: Bigstring.t] (Bigstring.sub buffer ~pos ~len) ~expect:message;
    Handler_result.Continue
  ;;

  let on_end_of_batch on_end_of_batch_count () = incr on_end_of_batch_count
  let wait () = Async_kernel_scheduler.Expert.run_cycles_until_no_jobs_remain ()

  let run_test run verify =
    let reader, writer = Pipe.create () in
    let reader = Pipe_and_buffer.create reader in
    let on_message_count = ref 0 in
    let on_end_of_batch_count = ref 0 in
    let on_message = on_message on_message_count in
    let on_end_of_batch = on_end_of_batch on_end_of_batch_count in
    ignore (Transport_reader.read_forever reader ~on_message ~on_end_of_batch);
    wait ();
    run (fun x -> Pipe.write_without_pushback writer (of_bigstring x));
    wait ();
    verify ~on_message_count ~on_end_of_batch_count
  ;;

  let verify ~on_message_count ~on_end_of_batch_count message_count batch_count =
    [%test_result: int] !on_message_count ~expect:message_count;
    [%test_result: int] !on_end_of_batch_count ~expect:batch_count
  ;;

  let write_1_msg_1_chunk write = write (create_message ())

  (* Not enough data, we need to accumulate *)
  let write_1_msg_3_chunk write =
    let buf = create_message_buf (Bigstring.of_string "") in
    let message_raw = default_message in
    Header.unsafe_set_payload_length buf ~pos:0 (Bigstring.length message_raw);
    write buf;
    wait ();
    let msg_len = Bigstring.length message_raw in
    let split_at = msg_len / 2 in
    let part1 = Bigstring.sub message_raw ~pos:0 ~len:split_at in
    let part2 = Bigstring.sub message_raw ~pos:split_at ~len:(msg_len - split_at) in
    write part1;
    wait ();
    write part2
  ;;

  (* Too much data, we need to loop withing the same batch *)
  let write_2_msg_1_chunk write =
    let one_msg = create_message_buf default_message in
    let two_msg = Bigstring.concat [ one_msg; one_msg ] in
    write two_msg
  ;;

  let write_2_msg_one_byte_at_a_time_34_chunk write =
    let one_msg = create_message_buf default_message in
    let two_msg = Bigstring.concat [ one_msg; one_msg ] in
    for i = 0 to Bigstring.length two_msg - 1 do
      write (Bigstring.sub two_msg ~pos:i ~len:1);
      wait ()
    done
  ;;

  let%test_unit _ = run_test write_1_msg_1_chunk (verify 1 1)

  let%test_unit _ =
    let f write =
      write_1_msg_1_chunk write;
      write_1_msg_1_chunk write
    in
    run_test f (verify 2 1)
  ;;

  let%test_unit _ = run_test write_1_msg_3_chunk (verify 1 3)
  let%test_unit _ = run_test write_2_msg_1_chunk (verify 2 1)
  let%test_unit _ = run_test write_2_msg_one_byte_at_a_time_34_chunk (verify 2 34)

  let%test_unit _ =
    let f write =
      write (create_message ());
      wait ();
      write_1_msg_3_chunk write;
      wait ();
      write_2_msg_1_chunk write;
      wait ();
      write_2_msg_one_byte_at_a_time_34_chunk write;
      wait ();
      write (create_message ())
    in
    run_test f (verify 7 40)
  ;;

  let%test_unit "can call read_forever multiple times" =
    let reader, writer = Pipe.create () in
    let reader = Pipe_and_buffer.create reader in
    let on_end_of_batch () = () in
    let write x = Pipe.write_without_pushback writer (of_bigstring x) in
    let num = ref 0 in
    let stop_after_on_message buffer ~pos ~len =
      incr num;
      [%test_result: Bigstring.t] (Bigstring.sub buffer ~pos ~len) ~expect:default_message;
      Handler_result.Stop ()
    in
    write_2_msg_1_chunk write;
    Pipe.close writer;
    don't_wait_for
      (let%bind x =
         Transport_reader.read_forever
           reader
           ~on_message:stop_after_on_message
           ~on_end_of_batch
       in
       assert (Result.is_ok x);
       let%bind x =
         Transport_reader.read_forever
           reader
           ~on_message:stop_after_on_message
           ~on_end_of_batch
       in
       assert (Result.is_ok x);
       Deferred.unit);
    wait ();
    [%test_result: Int.t] ~expect:2 !num
  ;;
end

let%test_module "Test_reader_string" = (module Test_reader (String_pipe_reader))
let%test_module "Test_reader_bigstring" = (module Test_reader (Bigstring_pipe_reader))

module Bench_reader (Transport_reader : Transport_reader) = struct
  open Transport_reader.For_testing

  let default_message = Bigstring.of_string "FOOBAR123"

  let create_message_buf (message : Bigstring.t) : Bigstring.t =
    let message_len = Bigstring.length message in
    let buf = Bigstring.create (Header.length + message_len) in
    Header.unsafe_set_payload_length buf ~pos:0 message_len;
    Bigstring.blit
      ~src:message
      ~dst:buf
      ~src_pos:0
      ~dst_pos:Header.length
      ~len:message_len;
    buf
  ;;

  let create_message () = create_message_buf default_message
  let one_message = create_message ()
  let two_messages = Bigstring.concat [ one_message; one_message ]

  let one_message_p1, one_message_p2 =
    let msg_len = Bigstring.length one_message in
    let split_at = msg_len / 2 in
    let part1 = Bigstring.sub one_message ~pos:0 ~len:split_at in
    let part2 = Bigstring.sub one_message ~pos:split_at ~len:(msg_len - split_at) in
    part1, part2
  ;;

  let on_message _buffer ~pos:_ ~len:_ = Handler_result.Continue
  let on_end_of_batch () = ()
  let wait () = Async_kernel_scheduler.Expert.run_cycles_until_no_jobs_remain ()

  let run_test run =
    let reader, writer = Pipe.create () in
    let reader = Pipe_and_buffer.create reader in
    let on_message = on_message in
    let on_end_of_batch = on_end_of_batch in
    ignore (Transport_reader.read_forever reader ~on_message ~on_end_of_batch);
    wait ();
    for _ = 0 to 100 do
      run (fun x -> Pipe.write_without_pushback writer (of_bigstring x))
    done;
    wait ()
  ;;

  let write_1_msg_1_chunk write = write one_message

  (* Not enough data, we need to accumulate *)
  let buf = create_message_buf (Bigstring.of_string "")

  let write_1_msg_3_chunk write =
    Header.unsafe_set_payload_length buf ~pos:0 (Bigstring.length one_message);
    write buf;
    write one_message_p1;
    write one_message_p2
  ;;

  (* Too much data, we need to loop withing the same batch *)
  let write_2_msg_1_chunk write = write two_messages

  let write_2_msg_one_byte_at_a_time_34_chunk write =
    for i = 0 to Bigstring.length two_messages - 1 do
      write (Bigstring.sub two_messages ~pos:i ~len:1)
    done
  ;;

  let%bench "write_1_msg_1_chunk" = run_test write_1_msg_1_chunk

  let%bench "write_2_msg_1_chunk" =
    let f write =
      write_1_msg_1_chunk write;
      write_1_msg_1_chunk write
    in
    run_test f
  ;;

  let%bench "write_1_msg_3_chunk" = run_test write_1_msg_3_chunk
  let%bench "write_2_msg_1_chunk" = run_test write_2_msg_1_chunk
  let%bench "write_2_msg_one_byte" = run_test write_2_msg_one_byte_at_a_time_34_chunk

  let%bench "multi" =
    let f write =
      write (create_message ());
      write_1_msg_3_chunk write;
      write_2_msg_1_chunk write;
      write_2_msg_one_byte_at_a_time_34_chunk write;
      write (create_message ())
    in
    run_test f
  ;;
end

let%bench_module "Test_reader_string" = (module Bench_reader (String_pipe_reader))
let%bench_module "Test_reader_bigstring" = (module Bench_reader (Bigstring_pipe_reader))