<commit_before>(*
Copyright 2015 Open vStorage NV

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

     http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*)

open Prelude
open Lwt
open Slice
open Asd_protocol
open Protocol

class client fd (ic, oc) id =
  let read_response deserializer =
    Llio.input_string ic >>= fun res_s ->
    let res_buf = Llio.make_buffer res_s 0 in
    match Llio.int_from res_buf with
    | 0 ->
      Lwt.return (deserializer res_buf)
    | err ->
      let open Error in
      let err' = deserialize' err res_buf in
      Lwt_log.debug_f "Exception in asd_client %s: %s" id (show err') >>= fun () ->
      lwt_fail err'
  in
  object(self)
    method private query : type req res. (req, res) query -> req -> res Lwt.t =
      fun command req ->
        let descr = code_to_description (command_to_code (Wrap_query command)) in
        Lwt_log.debug_f
          "asd_client %s: %s"
          id descr >>= fun () ->
        Alba_statistics.Statistics.with_timing_lwt
          (fun () ->
             let buf = Buffer.create 20 in
             Llio.int_to buf (command_to_code (Wrap_query command));
             query_request_serializer command buf req;
             Lwt_extra2.llio_output_and_flush oc (Buffer.contents buf) >>= fun () ->
             read_response (query_response_deserializer command)) >>= fun (t, r) ->
        Lwt_log.debug_f "asd_client %s: %s took %f" id descr t >>= fun () ->
        Lwt.return r

    method private update : type req res. (req, res) update -> req -> res Lwt.t =
      fun command req ->
        let descr = code_to_description (command_to_code (Wrap_update command)) in
        Lwt_log.debug_f
          "asd_client %s: %s"
          id descr >>= fun () ->
        Alba_statistics.Statistics.with_timing_lwt
          (fun () ->
             let buf = Buffer.create 20 in
             Llio.int_to buf (command_to_code (Wrap_update command));
             update_request_serializer command buf req;
             Lwt_extra2.llio_output_and_flush oc (Buffer.contents buf) >>= fun () ->
             read_response (update_response_deserializer command)) >>= fun (t, r) ->
        Lwt_log.debug_f "asd_client %s: %s took %f" id descr t >>= fun () ->
        Lwt.return r

    val mutable supports_multiget2 = None
    method multi_get keys =
      match supports_multiget2 with
      | None ->
         (* try multiget2, if it succeeds the asd supports it *)
         Lwt.catch
           (fun () ->
            self # multi_get2 keys >>= fun res ->
            supports_multiget2 <- Some true;
            Lwt.return res)
           (function
             | Error.Exn Error.Unknown_operation ->
                supports_multiget2 <- Some false;
                self # query MultiGet keys
             | exn ->
                Lwt.fail exn)
      | Some true ->
         self # multi_get2 keys
      | Some false ->
         self # query MultiGet keys

    method multi_get2 keys =
      self # query MultiGet2 keys >>= fun res ->
      Lwt_list.map_s
        (let open Value in
         function
          | None -> Lwt.return_none
          | Some (blob, cs) ->
             match blob with
             | Direct s -> Lwt.return (Some (s, cs))
             | Later size ->
                let target = Bytes.create size in

                let buffered = Lwt_io.buffered ic in
                (if buffered > 0
                 then Lwt_io.read_into ic target 0 buffered
                 else Lwt.return 0) >>= fun read ->
                assert (read = buffered);
                assert (0 = Lwt_io.buffered ic);

                Lwt_extra2.read_all
                  fd target
                  read (size - read)
                >>= fun read' ->
                assert (read + read' = size);
                Lwt.return (Some (Slice.wrap_bytes target, cs)))
        res

    method multi_get_string keys =
      self # multi_get (List.map Slice.wrap_string keys) >>= fun res ->
      Lwt.return
        (List.map
           (Option.map (fun (slice, cs) -> Slice.get_string_unsafe slice, cs))
           res)

    method get key =
      self # multi_get [ key ] >>= fun res ->
      List.hd_exn res |>
      Lwt.return

    method get_string key =
      self # multi_get_string [ key ] >>= fun res ->
      List.hd_exn res |>
      Lwt.return

    method set key value assertable ?(cs = Checksum.Checksum.NoChecksum) () =
      let u = Update.set key value cs assertable in
      self # update Apply ([], [u]) >>= fun _ ->
      Lwt.return ()

    method set_string ?(cs = Checksum.Checksum.NoChecksum)
             key value assertable
      =<commit_msg>asd client multiget2 bugfix for the case size <= buffered<commit_after>(*
Copyright 2015 Open vStorage NV

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

     http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*)

open Prelude
open Lwt
open Slice
open Asd_protocol
open Protocol

class client fd (ic, oc) id =
  let read_response deserializer =
    Llio.input_string ic >>= fun res_s ->
    let res_buf = Llio.make_buffer res_s 0 in
    match Llio.int_from res_buf with
    | 0 ->
      Lwt.return (deserializer res_buf)
    | err ->
      let open Error in
      let err' = deserialize' err res_buf in
      Lwt_log.debug_f "Exception in asd_client %s: %s" id (show err') >>= fun () ->
      lwt_fail err'
  in
  object(self)
    method private query : type req res. (req, res) query -> req -> res Lwt.t =
      fun command req ->
        let descr = code_to_description (command_to_code (Wrap_query command)) in
        Lwt_log.debug_f
          "asd_client %s: %s"
          id descr >>= fun () ->
        Alba_statistics.Statistics.with_timing_lwt
          (fun () ->
             let buf = Buffer.create 20 in
             Llio.int_to buf (command_to_code (Wrap_query command));
             query_request_serializer command buf req;
             Lwt_extra2.llio_output_and_flush oc (Buffer.contents buf) >>= fun () ->
             read_response (query_response_deserializer command)) >>= fun (t, r) ->
        Lwt_log.debug_f "asd_client %s: %s took %f" id descr t >>= fun () ->
        Lwt.return r

    method private update : type req res. (req, res) update -> req -> res Lwt.t =
      fun command req ->
        let descr = code_to_description (command_to_code (Wrap_update command)) in
        Lwt_log.debug_f
          "asd_client %s: %s"
          id descr >>= fun () ->
        Alba_statistics.Statistics.with_timing_lwt
          (fun () ->
             let buf = Buffer.create 20 in
             Llio.int_to buf (command_to_code (Wrap_update command));
             update_request_serializer command buf req;
             Lwt_extra2.llio_output_and_flush oc (Buffer.contents buf) >>= fun () ->
             read_response (update_response_deserializer command)) >>= fun (t, r) ->
        Lwt_log.debug_f "asd_client %s: %s took %f" id descr t >>= fun () ->
        Lwt.return r

    val mutable supports_multiget2 = None
    method multi_get keys =
      match supports_multiget2 with
      | None ->
         (* try multiget2, if it succeeds the asd supports it *)
         Lwt.catch
           (fun () ->
            self # multi_get2 keys >>= fun res ->
            supports_multiget2 <- Some true;
            Lwt.return res)
           (function
             | Error.Exn Error.Unknown_operation ->
                supports_multiget2 <- Some false;
                self # query MultiGet keys
             | exn ->
                Lwt.fail exn)
      | Some true ->
         self # multi_get2 keys
      | Some false ->
         self # query MultiGet keys

    method multi_get2 keys =
      self # query MultiGet2 keys >>= fun res ->
      Lwt_list.map_s
        (let open Value in
         function
          | None -> Lwt.return_none
          | Some (blob, cs) ->
             match blob with
             | Direct s -> Lwt.return (Some (s, cs))
             | Later size ->
                let target = Bytes.create size in

                let buffered = Lwt_io.buffered ic in
                (if size <= buffered
                 then
                   begin
                     Lwt_io.read_into ic target 0 size >>= fun read ->
                     assert (read = size);
                     Lwt.return ()
                   end
                 else
                   begin
                     (if buffered > 0
                      then Lwt_io.read_into ic target 0 buffered
                      else Lwt.return 0) >>= fun read ->
                     assert (read = buffered);
                     assert (0 = Lwt_io.buffered ic);

                     Lwt_extra2.read_all
                       fd target
                       read (size - read)
                     >>= fun read' ->
                     assert (read + read' = size);
                     Lwt.return ()
                   end) >>= fun () ->

                Lwt.return (Some (Slice.wrap_bytes target, cs)))
        res

    method multi_get_string keys =
      self # multi_get (List.map Slice.wrap_string keys) >>= fun res ->
      Lwt.return
        (List.map
           (Option.map (fun (slice, cs) -> Slice.get_string_unsafe slice, cs))
           res)

    method get key =
      self # multi_get [ key ] >>= fun res ->
      List.hd_exn res |>
      Lwt.return

    method get_string key =
      self # multi_get_string [ key ] >>= fun res ->
      List.hd_exn res |>
      Lwt.return

    method set key value assertable ?(cs = Checksum.Checksum.NoChecksum) () =
      let u = Update.set key value cs assertable in
      self # update Apply ([], [u]) >>= fun _ ->
      Lwt.return ()

    method set_string ?(cs = Checksum.Checksum.NoChecksum)
             key value assertable
      =
{'fraction_non_alphanumeric': 0.08963558920965452, 'fraction_numerical': 0.0058684335068622815, 'mean_word_length': 1.7783328950828294, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 165, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3675471', 'n_tokens_mistral': 3170, 'n_tokens_neox': 3064, 'n_words': 1006}