<commit_before>(*
 * Copyright (C) 2011-2012 <NAME> <<EMAIL>>
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version,
 * with the special exception on linking described in file LICENSE.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *)

open Printf
open Lwt
open Obs_data_model
open Obs_protocol
open Obs_request

let data_protocol_version = (0, 0, 0)

type replication_wait = Await_reception | Await_commit

module TY : sig
  type ks_id = private int

  val ks_id_of_int : int -> ks_id
  val new_unique_id : unit -> ks_id
end =
struct
  type ks_id = int

  let new_unique_id =
    let n = ref 0 in
      (fun () -> incr n; !n)

  let ks_id_of_int n = n
end

open TY

module Make
  (D : sig
     include Obs_data_model.S
     include Obs_data_model.BACKUP_SUPPORT with type backup_cursor := backup_cursor

     module Replication : Obs_replication.REPLICATION_SERVER
       with type db := db and type raw_dump := Raw_dump.raw_dump

     val use_thread_pool : db -> bool -> unit
   end)
  (P : PAYLOAD) =
struct
  open Request

  type request_id = string

  exception Abort_exn of request_id
  exception Commit_exn of request_id

  module H = Hashtbl.Make(struct
                              type t = ks_id
                              let hash (n : ks_id) = (n :> int)
                              let equal a b = (a == b)
                            end)

  module Notif_queue : sig
    type 'a t
    val empty : 'a t
    val push : 'a -> 'a t -> 'a t
    val iter : ('a -> unit) -> 'a t  -> unit
  end =
  struct
    type 'a t = 'a list

    let empty = []

    (* TODO: further "compression" by keeping set of previous notifications? *)
    let push x = function
        y :: _ as l when y = x -> l
      | l -> x :: l

    let iter f t = List.iter f (List.rev t)
  end

  let num_clients = ref 0

  let dummy_auto_yield () = return ()

  let auto_yielder = ref dummy_auto_yield

  type client_id = int
  type topic = string
  type keyspace_name = string

  type subscription_descriptor = { subs_ks : keyspace_name; subs_topic : topic }

  type subs_stream = string Lwt_stream.t * (string option -> unit)
  type mutable_string_set = (string, unit) Hashtbl.t

  type t =
      {
        db : D.db;
        subscriptions : (subscription_descriptor, subs_stream H.t) Hashtbl.t;
        raw_dumps : (Int64.t, D.Raw_dump.raw_dump) Hashtbl.t;
        mutable raw_dump_seqno : Int64.t;
        replication_wait : replication_wait;
      }

  type client_state =
      {
        id : client_id;
        keyspaces : (keyspace * req_handler Stack.t) H.t;
        ich : Lwt_io.input_channel;
        och : Lwt_io.output_channel;
        server : t;
        mutable in_buf : string;
        out_buf : Obs_bytea.t;
        debug : bool;
        mutable pending_reqs : int;
        wait_for_pending_reqs : unit Lwt_condition.t;
        async_req_region : Lwt_util.region;
        signal_error : Request.request option Lwt.t * Request.request option Lwt.u;
      }

  and keyspace =
    { ks_unique_id : ks_id;
      ks_tx_key : tx_data Lwt.key;
      ks_ks : D.keyspace;
      ks_subs_stream : subs_stream;
      ks_subscriptions : mutable_string_set;
  }

  and req_handler = request_id:request_id -> Request.request -> unit Lwt.t

  and tx_data =
      {
        mutable tx_notifications : string Notif_queue.t;
      }

  let find_or_add find add h k =
    try find h k with Not_found -> add h k

  let read_exactly c n =
    let s = String.create n in
      Lwt_io.read_into_exactly c.ich s 0 n >>
      return s

  (* the chunk of the region needed to perform a given request --- allows to
   * impose diff limits per type *)
  let request_slot_cost = function
      Load _ -> 100
    | _ -> 1

  let with_raw_dump c id default f =
    try_lwt
      let raw_dump = Hashtbl.find c.server.raw_dumps id in
        f raw_dump
    with Not_found -> return default

  let rec pp_list pp fmt = function
      [] -> ()
    | x :: tl -> Format.fprintf fmt "%a,@ " pp x;
                 pp_list pp fmt tl

  let pp_request_id fmt id =
    if Obs_protocol.is_sync_req id then
      Format.fprintf fmt "<sync>"
    else
      Format.fprintf fmt "%s" Cryptokit.(transform_string (Hexa.encode ()) id)

  let pp_slice fmt (lastkey, kds) =
    let open Obs_data_model in
    let pp_col fmt col =
      Format.fprintf fmt
        "{ @[name: %S;@ data: %S;@ timestamp: %s@]}"
        col.name col.data
        (match col.timestamp with
             No_timestamp -> "auto"
           | Timestamp x -> Int64.to_string x) in
    let pp_kd fmt kd =
      Format.fprintf fmt
        "{ @[key: %S;@ last_column: %S;@ columns: @[%a@] }@]"
        kd.key kd.last_column (pp_list pp_col) kd.columns
    in
      Format.fprintf fmt "(@[%s,@ %a)@]"
        (match lastkey with None -> "None" | Some x -> sprintf "Some %S" x)
        (pp_list pp_kd) kds

  let maybe_pp c pp ~request_id x =
    if c.debug then begin
      Format.eprintf
        "Sending response for %a to %d:@\n @[%a@]@."
        pp_request_id request_id c.id pp x;
    end;
    return x

  let rec service c =
    !auto_yielder () >>
    lwt request_id, len, crc =
      match_lwt read_header c.ich with
          Header x -> return x
        | Corrupted_header ->
            (* we can't even trust the request_id, so all that's left is
             * dropping the connection *)
            raise_lwt (Error Corrupted_frame)
    in
      match_lwt
        Lwt.choose [fst c.signal_error; read_request c ~request_id len crc]
      with
          None -> service c
        | Some r ->
            ignore begin
              c.pending_reqs <- c.pending_reqs + 1;
              try_lwt
                Lwt_util.run_in_region c.async_req_region
                  (request_slot_cost r)
                  (fun () ->
                     begin try_lwt
                       relay_to_handler c ~request_id r
                     with
                       | End_of_file
                       | Unix.Unix_error((Unix.ECONNRESET | Unix.EPIPE), _, _) ->
                           (* catch exns that indicate that the connection
                            * is gone, and signal End_of_file *)
                           begin try
                             Lwt.wakeup_exn (snd c.signal_error) End_of_file
                           with _ -> ()
                           end;
                           return ()
                       | e ->
                           Format.eprintf
                             "Internal error %s@\n\
                              request:@\n\
                              %a@."
                             (Printexc.to_string e)
                             Request.pp r;
                           try_lwt
                             P.internal_error c.och ~request_id ()
                           with _ -> return ()
                     end)
              finally
                c.pending_reqs <- c.pending_reqs - 1;
                if c.pending_reqs = 0 then
                  Lwt_condition.broadcast c.wait_for_pending_reqs ();
                return ()
            end;
            (* here we block if the allowed number of async reqs is reached
             * until one of them is done *)
            Lwt_util.run_in_region c.async_req_region 1 (fun () -> return ()) >>
            service c

  and relay_to_handler c ~request_id = function
    | List_tables { List_tables.keyspace; _ }
    | Table_size_on_disk { Table_size_on_disk.keyspace; _ }
    | Key_range_size_on_disk { Key_range_size_on_disk.keyspace; _ }
    | Begin { Begin.keyspace; _ }
    | Commit { Commit.keyspace; _ }
    | Abort { Abort.keyspace; _ }
    | Lock { Lock.keyspace; _ }
    | Exist_keys { Exist_keys.keyspace; _ }
    | Get_keys { Get_keys.keyspace; _ }
    | Count_keys { Count_keys.keyspace; _ }
    | Get_slice { Get_slice.keyspace; _ }
    | Get_slice_values { Get_slice_values.keyspace; _ }
    | Get_slice_values_timestamps { Get_slice_values_timestamps.keyspace; _ }
    | Get_columns { Get_columns.keyspace; _ }
    | Get_column_values { Get_column_values.keyspace; _ }
    | Get_column { Get_column.keyspace; _ }
    | Put_columns { Put_columns.keyspace; _ }
    | Delete_columns { Delete_columns.keyspace; _ }
    | Delete_key { Delete_key.keyspace; _ }
    | Dump { Dump.keyspace; _ }
    | Load { Load.keyspace; _ }
    | Stats { Stats.keyspace; _ }
    | Listen { Listen.keyspace; _ }
    | Unlisten { Unlisten.keyspace; _ }
    | Notify { Notify.keyspace; _ }
    | Await { Await.keyspace; _ }
    | Release_keyspace { Release_keyspace.keyspace; _ } as r ->
        begin try
          let handlers = snd (H.find c.keyspaces (ks_id_of_int keyspace)) in
            if Stack.is_empty handlers then
              respond c ~request_id r
            else begin
              let handler = Stack.top handlers in
                handler ~request_id r
            end
        with Not_found ->
          (* will usually respond with a unknown_keyspace error *)
          respond c ~request_id r
        end
    | Register_keyspace _ | Get_keyspace _ | List_keyspaces _
    | Trigger_raw_dump _ | Raw_dump_release _ | Raw_dump_list_files _
    | Raw_dump_file_digest _ as r ->
        respond c ~request_id r

  and read_request c ~request_id len crc =
    if String.length c.in_buf < len then c.in_buf <- String.create len;
    Lwt_io.read_into_exactly c.ich c.in_buf 0 len >>
    lwt crc2 = read_exactly c 4 in
    let crc2' = Obs_crc32c.substring c.in_buf 0 len in
    let gb n = Char.code c.in_buf.[n] in
    let format_id = gb 0 + (gb 1 lsl 8) + (gb 2 lsl 16) + (gb 3 lsl 24) in
      Obs_crc32c.xor crc2 crc;
      if crc2 <> crc2' then begin
        P.bad_request c.och ~request_id () >>
        return None
      end else begin
        match Obs_protocol_payload.Obs_request_serialization.of_format_id format_id with
            `Extprot -> begin
              let m =
                try Some (Extprot.Conv.deserialize Request.read ~offset:4 c.in_buf)
                with _ -> None
              in match m with
                  None -> P.bad_request c.och ~request_id () >> return None
                | Some _ as x -> return x
            end
          | `Raw -> P.bad_request c.och ~request_id () >> return None
          | `Unknown -> P.unknown_serialization c.och ~request_id () >> return None
      end

  and respond ?buf c ~request_id r =
    if c.debug then
      Format.eprintf "Got request %a from %d@\n @[%a@]@."
        pp_request_id request_id c.id Request.pp r;
    match r with
      Register_keyspace { Register_keyspace.name } ->
        lwt ks_ks = D.register_keyspace c.server.db name in
        let ks_unique_id = ks_id_of_int (D.keyspace_id ks_ks) in
        let ks =
          { ks_unique_id; ks_ks; ks_tx_key = <KEY> ();
            ks_subs_stream = Lwt_stream.create ();
            ks_subscriptions = Hashtbl.create 13;
          }
        in
          H.add c.keyspaces ks_unique_id (ks, Stack.create ());
          P.return_keyspace ?buf c.och ~request_id (ks_unique_id :> int)
    | Get_keyspace { Get_keyspace.name; } -> begin
        match_lwt D.get_keyspace c.server.db name with
            None -> P.return_keyspace_maybe ?buf c.och ~request_id None
          | Some ks_ks ->
              let ks_unique_id = ks_id_of_int (D.keyspace_id ks_ks) in
              let ks =
                { ks_ks; ks_unique_id; ks_tx_key = <KEY> ();
                  ks_subs_stream = Lwt_stream.create ();
                  ks_subscriptions = Hashtbl.create 13;
                }
              in
                H.add c.keyspaces ks_unique_id (ks, Stack.create ());
                P.return_keyspace_maybe ?buf c.och ~request_id
                  (Some ks_unique_id :> int option)
      end
    | Release_keyspace { Release_keyspace.keyspace } -> begin
        try_lwt
          let keyspace = fst (H.find c.keyspaces (ks_id_of_int keyspace)) in
            terminate_keyspace_subs c.server keyspace;
            return ()
        with Not_found -> return ()
        finally
          P.return_ok ?buf c.och ~request_id ()
      end
    | List_keyspaces _ ->
        D.list_keyspaces c.server.db >>=
          P.return_keyspace_list ?buf c.och ~request_id
    | List_tables { List_tables.keyspace } ->
        with_keyspace c keyspace ~request_id
          (fun ks ->
             (D.list_tables ks.ks_ks :> string list Lwt.t) >>=
             P.return_table_list ?buf c.och ~request_id)
    | Table_size_on_disk { Table_size_on_disk.keyspace; table; } ->
        with_keyspace c keyspace ~request_id
          (fun ks ->
             D.table_size_on_disk ks.ks_ks table >>=
             P.return_table_size_on_disk ?buf c.och ~request_id)
    | Key_range_size_on_disk { Key_range_size_on_disk.keyspace; table; range; } ->
        with_keyspace c keyspace ~request_id
          (fun ks ->
             D.key_range_size_on_disk ks.ks_ks table
               ?first:range.first ?up_to:range.up_to >>=
             P.return_key_range_size_on_disk ?buf c.och ~request_id)
    | Begin { Begin.keyspace; tx_type; } ->
        with_keyspace c keyspace ~request_id
          (fun ks ->
             let is_outermost, parent_tx_data, tx_data =
               match Lwt.get ks.ks_tx_key with
                   None ->
                     let d = { tx_notifications = Notif_queue.empty } in
                       (true, d, d)
                 | Some d -> (false, d, { tx_notifications = d.tx_notifications; }) in
             let transaction_f = match tx_type with
                 Tx_type.Repeatable_read -> D.repeatable_read_transaction
               | Tx_type.Read_committed -> D.read_committed_transaction in
             let handlers = snd (H.find c.keyspaces (ks_id_of_int keyspace)) in
               try_lwt
                 (* we'll put the request_id from the request that forced the
                  * commit (carried by Commit_exn) in finishing_reqid *)
                 let finishing_reqid = ref "" in
                 lwt () =
                   transaction_f ks.ks_ks
                     (fun _ ->
                        Lwt.with_value ks.ks_tx_key (Some tx_data)
                          (fun () ->
                             try_lwt
                               (* install handler *)
                               let req_stream, pushf = Lwt_stream.create () in
                               let handler ~request_id r =
                                 (* we keep a ref to the stream because
                                  * otherwise it can be collected! pushf
                                  * only has a weak ref to it *)
                                 ignore req_stream;
                                 pushf (Some (`Req (request_id, r)));
                                 return ()
                               in
                                 Stack.push handler handlers;
                                 P.return_ok ?buf c.och ~request_id () >>
                                 let rec handle_reqs () =
                                   match_lwt Lwt_stream.get req_stream with
                                       (Some `Req (request_id, r)) ->
                                         (* we respond in a background thread
                                          * so as to process reqs in parallel
                                          * but need to capture exns (in
                                          * particular, Abort_exn and
                                          * Commit_exn) *)
                                         ignore begin
                                           try_lwt
                                             respond c ~request_id r
                                           with e ->
                                             pushf (Some (`Exn e));
                                             return ()
                                         end;
                                         handle_reqs ()
                                     | Some (`Exn e) ->
                                         raise_lwt e
                                     | None -> return ()
                                 in handle_reqs () >>
                                    return ()
                             with Commit_exn req_id ->
                               finishing_reqid := req_id;
                               return ()
                             finally
                               let _ : req_handler = Stack.pop handlers in
                                 return ())) >>
                     P.return_ok ?buf c.och ~request_id:!finishing_reqid ()
                 in
                   (* we deliver notifications _after_ actual commit
                    * (otherwise, new data wouldn't be found if another client
                    *  performs a query right away) *)
                   begin if is_outermost then
                     Notif_queue.iter (notify c.server ks.ks_ks) tx_data.tx_notifications
                   else parent_tx_data.tx_notifications <- tx_data.tx_notifications
                   end;
                   return ()
               with Abort_exn request_id ->
                   P.return_ok ?buf c.och ~request_id ())
    | Commit { Commit.keyspace; _ } ->
        (* only commit if we're inside a tx *)
        with_keyspace c keyspace ~request_id
          (fun ks ->
             begin match Lwt.get ks.ks_tx_key with
                 None -> P.return_ok ?buf c.och ~request_id ()
               | Some _ -> raise_lwt (Commit_exn request_id)
             end)
    | Abort { Abort.keyspace; _ } ->
        (* only abort if we're inside a tx *)
        with_keyspace c keyspace ~request_id
          (fun ks ->
             begin match Lwt.get ks.ks_tx_key with
                 None -> P.return_ok ?buf c.och ~request_id ()
               | Some _ -> raise_lwt (Abort_exn request_id)
             end)
    | Lock { Lock.keyspace; names; shared; } -><commit_msg>Obs_protocol_server refactoring.<commit_after>(*
 * Copyright (C) 2011-2012 <NAME> <<EMAIL>>
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version,
 * with the special exception on linking described in file LICENSE.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *)

open Printf
open Lwt
open Obs_data_model
open Obs_protocol
open Obs_request

let data_protocol_version = (0, 0, 0)

type replication_wait = Await_reception | Await_commit

module TY : sig
  type ks_id = private int

  val ks_id_of_int : int -> ks_id
  val new_unique_id : unit -> ks_id
end =
struct
  type ks_id = int

  let new_unique_id =
    let n = ref 0 in
      (fun () -> incr n; !n)

  let ks_id_of_int n = n
end

open TY

module Make
  (D : sig
     include Obs_data_model.S
     include Obs_data_model.BACKUP_SUPPORT with type backup_cursor := backup_cursor

     module Replication : Obs_replication.REPLICATION_SERVER
       with type db := db and type raw_dump := Raw_dump.raw_dump

     val use_thread_pool : db -> bool -> unit
   end)
  (P : PAYLOAD) =
struct
  open Request

  type request_id = string

  exception Abort_exn of request_id
  exception Commit_exn of request_id

  module H = Hashtbl.Make(struct
                              type t = ks_id
                              let hash (n : ks_id) = (n :> int)
                              let equal a b = (a == b)
                            end)

  module Notif_queue : sig
    type 'a t
    val empty : 'a t
    val push : 'a -> 'a t -> 'a t
    val iter : ('a -> unit) -> 'a t  -> unit
  end =
  struct
    type 'a t = 'a list

    let empty = []

    (* TODO: further "compression" by keeping set of previous notifications? *)
    let push x = function
        y :: _ as l when y = x -> l
      | l -> x :: l

    let iter f t = List.iter f (List.rev t)
  end

  let num_clients = ref 0

  let dummy_auto_yield () = return ()

  let auto_yielder = ref dummy_auto_yield

  type client_id = int
  type topic = string
  type keyspace_name = string

  type subscription_descriptor = { subs_ks : keyspace_name; subs_topic : topic }

  type subs_stream = string Lwt_stream.t * (string option -> unit)
  type mutable_string_set = (string, unit) Hashtbl.t

  type t =
      {
        db : D.db;
        subscriptions : (subscription_descriptor, subs_stream H.t) Hashtbl.t;
        raw_dumps : (Int64.t, D.Raw_dump.raw_dump) Hashtbl.t;
        mutable raw_dump_seqno : Int64.t;
        replication_wait : replication_wait;
      }

  type client_state =
      {
        id : client_id;
        keyspaces : (keyspace * req_handler Stack.t) H.t;
        ich : Lwt_io.input_channel;
        och : Lwt_io.output_channel;
        server : t;
        mutable in_buf : string;
        out_buf : Obs_bytea.t;
        debug : bool;
        mutable pending_reqs : int;
        wait_for_pending_reqs : unit Lwt_condition.t;
        async_req_region : Lwt_util.region;
        signal_error : Request.request option Lwt.t * Request.request option Lwt.u;
      }

  and keyspace =
    { ks_unique_id : ks_id;
      ks_tx_key : tx_data Lwt.key;
      ks_ks : D.keyspace;
      ks_subs_stream : subs_stream;
      ks_subscriptions : mutable_string_set;
  }

  and req_handler = request_id:request_id -> Request.request -> unit Lwt.t

  and tx_data =
      {
        mutable tx_notifications : string Notif_queue.t;
      }

  let find_or_add find add h k =
    try find h k with Not_found -> add h k

  let read_exactly c n =
    let s = String.create n in
      Lwt_io.read_into_exactly c.ich s 0 n >>
      return s

  (* the chunk of the region needed to perform a given request --- allows to
   * impose diff limits per type *)
  let request_slot_cost = function
      Load _ -> 100
    | _ -> 1

  let with_raw_dump c id default f =
    try_lwt
      let raw_dump = Hashtbl.find c.server.raw_dumps id in
        f raw_dump
    with Not_found -> return default

  let rec pp_list pp fmt = function
      [] -> ()
    | x :: tl -> Format.fprintf fmt "%a,@ " pp x;
                 pp_list pp fmt tl

  let pp_request_id fmt id =
    if Obs_protocol.is_sync_req id then
      Format.fprintf fmt "<sync>"
    else
      Format.fprintf fmt "%s" Cryptokit.(transform_string (Hexa.encode ()) id)

  let pp_slice fmt (lastkey, kds) =
    let open Obs_data_model in
    let pp_col fmt col =
      Format.fprintf fmt
        "{ @[name: %S;@ data: %S;@ timestamp: %s@]}"
        col.name col.data
        (match col.timestamp with
             No_timestamp -> "auto"
           | Timestamp x -> Int64.to_string x) in
    let pp_kd fmt kd =
      Format.fprintf fmt
        "{ @[key: %S;@ last_column: %S;@ columns: @[%a@] }@]"
        kd.key kd.last_column (pp_list pp_col) kd.columns
    in
      Format.fprintf fmt "(@[%s,@ %a)@]"
        (match lastkey with None -> "None" | Some x -> sprintf "Some %S" x)
        (pp_list pp_kd) kds

  let maybe_pp c pp ~request_id x =
    if c.debug then begin
      Format.eprintf
        "Sending response for %a to %d:@\n @[%a@]@."
        pp_request_id request_id c.id pp x;
    end;
    return x

  let rec service c =
    !auto_yielder () >>
    lwt request_id, len, crc =
      match_lwt read_header c.ich with
          Header x -> return x
        | Corrupted_header ->
            (* we can't even trust the request_id, so all that's left is
             * dropping the connection *)
            raise_lwt (Error Corrupted_frame)
    in
      match_lwt
        Lwt.choose [fst c.signal_error; read_request c ~request_id len crc]
      with
          None -> service c
        | Some r ->
            ignore begin
              c.pending_reqs <- c.pending_reqs + 1;
              try_lwt
                Lwt_util.run_in_region c.async_req_region
                  (request_slot_cost r)
                  (fun () ->
                     begin try_lwt
                       relay_to_handler c ~request_id r
                     with
                       | End_of_file
                       | Unix.Unix_error((Unix.ECONNRESET | Unix.EPIPE), _, _) ->
                           (* catch exns that indicate that the connection
                            * is gone, and signal End_of_file *)
                           begin try
                             Lwt.wakeup_exn (snd c.signal_error) End_of_file
                           with _ -> ()
                           end;
                           return ()
                       | e ->
                           Format.eprintf
                             "Internal error %s@\n\
                              request:@\n\
                              %a@."
                             (Printexc.to_string e)
                             Request.pp r;
                           try_lwt
                             P.internal_error c.och ~request_id ()
                           with _ -> return ()
                     end)
              finally
                c.pending_reqs <- c.pending_reqs - 1;
                if c.pending_reqs = 0 then
                  Lwt_condition.broadcast c.wait_for_pending_reqs ();
                return ()
            end;
            (* here we block if the allowed number of async reqs is reached
             * until one of them is done *)
            Lwt_util.run_in_region c.async_req_region 1 (fun () -> return ()) >>
            service c

  and relay_to_handler c ~request_id = function
    | List_tables { List_tables.keyspace; _ }
    | Table_size_on_disk { Table_size_on_disk.keyspace; _ }
    | Key_range_size_on_disk { Key_range_size_on_disk.keyspace; _ }
    | Begin { Begin.keyspace; _ }
    | Commit { Commit.keyspace; _ }
    | Abort { Abort.keyspace; _ }
    | Lock { Lock.keyspace; _ }
    | Exist_keys { Exist_keys.keyspace; _ }
    | Get_keys { Get_keys.keyspace; _ }
    | Count_keys { Count_keys.keyspace; _ }
    | Get_slice { Get_slice.keyspace; _ }
    | Get_slice_values { Get_slice_values.keyspace; _ }
    | Get_slice_values_timestamps { Get_slice_values_timestamps.keyspace; _ }
    | Get_columns { Get_columns.keyspace; _ }
    | Get_column_values { Get_column_values.keyspace; _ }
    | Get_column { Get_column.keyspace; _ }
    | Put_columns { Put_columns.keyspace; _ }
    | Delete_columns { Delete_columns.keyspace; _ }
    | Delete_key { Delete_key.keyspace; _ }
    | Dump { Dump.keyspace; _ }
    | Load { Load.keyspace; _ }
    | Stats { Stats.keyspace; _ }
    | Listen { Listen.keyspace; _ }
    | Unlisten { Unlisten.keyspace; _ }
    | Notify { Notify.keyspace; _ }
    | Await { Await.keyspace; _ }
    | Release_keyspace { Release_keyspace.keyspace; _ } as r ->
        begin try
          let handlers = snd (H.find c.keyspaces (ks_id_of_int keyspace)) in
            if Stack.is_empty handlers then
              respond c ~request_id r
            else begin
              let handler = Stack.top handlers in
                handler ~request_id r
            end
        with Not_found ->
          (* will usually respond with a unknown_keyspace error *)
          respond c ~request_id r
        end
    | Register_keyspace _ | Get_keyspace _ | List_keyspaces _
    | Trigger_raw_dump _ | Raw_dump_release _ | Raw_dump_list_files _
    | Raw_dump_file_digest _ as r ->
        respond c ~request_id r

  and read_request c ~request_id len crc =
    if String.length c.in_buf < len then c.in_buf <- String.create len;
    Lwt_io.read_into_exactly c.ich c.in_buf 0 len >>
    lwt crc2 = read_exactly c 4 in
    let crc2' = Obs_crc32c.substring c.in_buf 0 len in
    let gb n = Char.code c.in_buf.[n] in
    let format_id = gb 0 + (gb 1 lsl 8) + (gb 2 lsl 16) + (gb 3 lsl 24) in
      Obs_crc32c.xor crc2 crc;
      if crc2 <> crc2' then begin
        P.bad_request c.och ~request_id () >>
        return None
      end else begin
        match Obs_protocol_payload.Obs_request_serialization.of_format_id format_id with
            `Extprot -> begin
              let m =
                try Some (Extprot.Conv.deserialize Request.read ~offset:4 c.in_buf)
                with _ -> None
              in match m with
                  None -> P.bad_request c.och ~request_id () >> return None
                | Some _ as x -> return x
            end
          | `Raw -> P.bad_request c.och ~request_id () >> return None
          | `Unknown -> P.unknown_serialization c.och ~request_id () >> return None
      end

  and respond ?buf c ~request_id r =
    if c.debug then
      Format.eprintf "Got request %a from %d@\n @[%a@]@."
        pp_request_id request_id c.id Request.pp r;
    match r with
      Register_keyspace { Register_keyspace.name } ->
        lwt ks_ks = D.register_keyspace c.server.db name in
        let ks_unique_id = ks_id_of_int (D.keyspace_id ks_ks) in
        let ks =
          { ks_unique_id; ks_ks; ks_tx_key = <KEY> ();
            ks_subs_stream = Lwt_stream.create ();
            ks_subscriptions = Hashtbl.create 13;
          }
        in
          H.add c.keyspaces ks_unique_id (ks, Stack.create ());
          P.return_keyspace ?buf c.och ~request_id (ks_unique_id :> int)
    | Get_keyspace { Get_keyspace.name; } -> begin
        match_lwt D.get_keyspace c.server.db name with
            None -> P.return_keyspace_maybe ?buf c.och ~request_id None
          | Some ks_ks ->
              let ks_unique_id = ks_id_of_int (D.keyspace_id ks_ks) in
              let ks =
                { ks_ks; ks_unique_id; ks_tx_key = <KEY> ();
                  ks_subs_stream = Lwt_stream.create ();
                  ks_subscriptions = Hashtbl.create 13;
                }
              in
                H.add c.keyspaces ks_unique_id (ks, Stack.create ());
                P.return_keyspace_maybe ?buf c.och ~request_id
                  (Some ks_unique_id :> int option)
      end
    | Release_keyspace { Release_keyspace.keyspace } -> begin
        try_lwt
          let keyspace = fst (H.find c.keyspaces (ks_id_of_int keyspace)) in
            terminate_keyspace_subs c.server keyspace;
            return ()
        with Not_found -> return ()
        finally
          P.return_ok ?buf c.och ~request_id ()
      end
    | List_keyspaces _ ->
        D.list_keyspaces c.server.db >>=
          P.return_keyspace_list ?buf c.och ~request_id
    | List_tables { List_tables.keyspace } ->
        with_keyspace c keyspace ~request_id
          (fun ks ->
             (D.list_tables ks.ks_ks :> string list Lwt.t) >>=
             P.return_table_list ?buf c.och ~request_id)
    | Table_size_on_disk { Table_size_on_disk.keyspace; table; } ->
        with_keyspace c keyspace ~request_id
          (fun ks ->
             D.table_size_on_disk ks.ks_ks table >>=
             P.return_table_size_on_disk ?buf c.och ~request_id)
    | Key_range_size_on_disk { Key_range_size_on_disk.keyspace; table; range; } ->
        with_keyspace c keyspace ~request_id
          (fun ks ->
             D.key_range_size_on_disk ks.ks_ks table
               ?first:range.first ?up_to:range.up_to >>=
             P.return_key_range_size_on_disk ?buf c.och ~request_id)
    | Begin { Begin.keyspace; tx_type; } ->
        with_keyspace c keyspace ~request_id
          (fun ks ->
             let is_outermost, parent_tx_data, tx_data =
               match Lwt.get ks.ks_tx_key with
                   None ->
                     let d = { tx_notifications = Notif_queue.empty } in
                       (true, d, d)
                 | Some d -> (false, d, { tx_notifications = d.tx_notifications; }) in
             let transaction_f = match tx_type with
                 Tx_type.Repeatable_read -> D.repeatable_read_transaction
               | Tx_type.Read_committed -> D.read_committed_transaction in
             let handlers = snd (H.find c.keyspaces (ks_id_of_int keyspace)) in
               try_lwt
                 (* we get the request_id from the request that forced the
                  * commit (carried by Commit_exn) *)
                 lwt commit_reqid =
                   transaction_f ks.ks_ks
                     (fun _ ->
                        try_lwt
                          (* install handler *)
                          let req_stream, pushf = Lwt_stream.create () in
                          let handler ~request_id r =
                            (* we keep a ref to the stream because
                             * otherwise it can be collected! pushf
                             * only has a weak ref to it *)
                            ignore req_stream;
                            pushf (Some (`Req (request_id, r)));
                            return ()
                          in
                            Stack.push handler handlers;
                            P.return_ok ?buf c.och ~request_id () >>
                            let rec handle_reqs () =
                              match_lwt Lwt_stream.get req_stream with
                                  (Some `Req (request_id, r)) ->
                                    (* we respond in a background thread so as
                                     * to process reqs in parallel but need to
                                     * capture exns (in particular, Abort_exn
                                     * and Commit_exn) *)
                                    ignore begin
                                      try_lwt
                                        Lwt.with_value
                                          ks.ks_tx_key (Some tx_data)
                                          (fun () -> respond c ~request_id r)
                                      with e ->
                                        pushf (Some (`Exn e));
                                        return ()
                                    end;
                                    handle_reqs ()
                                | Some (`Exn (Commit_exn req_id)) -> return req_id
                                | Some (`Exn e) -> raise_lwt e
                                | None -> fail (Failure "TX request stream canceled")
                            in handle_reqs ()
                        finally
                          let _ : req_handler = Stack.pop handlers in
                                 return ())
                 in
                   P.return_ok ?buf c.och ~request_id:commit_reqid () >>
                   (* we deliver notifications _after_ actual commit
                    * (otherwise, new data wouldn't be found if another client
                    *  performs a query right away) *)
                   let () =
                     begin if is_outermost then
                       Notif_queue.iter (notify c.server ks.ks_ks)
                         tx_data.tx_notifications
                     else parent_tx_data.tx_notifications <- tx_data.tx_notifications
                     end
                   in return ()
               with Abort_exn request_id -> P.return_ok ?buf c.och ~request_id ())
    | Commit { Commit.keyspace; _ } ->
        (* only commit if we're inside a tx *)
        with_keyspace c keyspace ~request_id
          (fun ks ->
             begin match Lwt.get ks.ks_tx_key with
                 None -> P.return_ok ?buf c.och ~request_id ()
               | Some _ -> raise_lwt (Commit_exn request_id)
             end)
    | Abort { Abort.keyspace; _ } ->
        (* only abort if we're inside a tx *)
        with_keyspace c keyspace ~request_id
          (fun ks ->
             begin match Lwt.get ks.ks_tx_key with
                 None -> P.return_ok ?buf c.och ~request_id ()
               | Some _ -> raise_lwt (Abort_exn request_id)
             end)
    | Lock { Lock.keyspace; names; shared; } ->
{'fraction_non_alphanumeric': 0.07888151452507888, 'fraction_numerical': 0.003808073115003808, 'mean_word_length': 1.5107559926244623, 'pattern_counts': {'":': 0, '<': 27, '<?xml version=': 0, '>': 243, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 82, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6119718', 'n_tokens_mistral': 11202, 'n_tokens_neox': 10823, 'n_words': 3056}