266 lines
7.7 KiB
OCaml
266 lines
7.7 KiB
OCaml
(**************************************************************************)
|
|
(* *)
|
|
(* Copyright (c) 2014 - 2016. *)
|
|
(* Dynamic Ledger Solutions, Inc. <contact@tezos.com> *)
|
|
(* *)
|
|
(* All rights reserved. No warranty, explicit or implicit, provided. *)
|
|
(* *)
|
|
(**************************************************************************)
|
|
|
|
module LC = Lwt_condition
|
|
|
|
open Lwt.Infix
|
|
open Logging.Core
|
|
|
|
let may f = function
|
|
| None -> Lwt.return_unit
|
|
| Some x -> f x
|
|
|
|
let never_ending = fst (Lwt.wait ())
|
|
|
|
(* A non exception-based cancelation mechanism. Builds a [cancelation]
|
|
thread to bind / pick on, awoken when a cancelation is requested by
|
|
[cancel ()]. [on_cancel cb] registers a callback to be called at
|
|
cancelation. [cancel ()] finishes when all calbacks have completed
|
|
(sequentially), instantly when called more than once. *)
|
|
let canceler ()
|
|
: (unit -> unit Lwt.t) *
|
|
(unit -> unit Lwt.t) *
|
|
((unit -> unit Lwt.t) -> unit) =
|
|
let cancelation = LC.create () in
|
|
let cancelation_complete = LC.create () in
|
|
let cancel_hook = ref (fun () -> Lwt.return ()) in
|
|
let canceling = ref false and canceled = ref false in
|
|
let cancel () =
|
|
if !canceled then
|
|
Lwt.return ()
|
|
else if !canceling then
|
|
LC.wait cancelation_complete
|
|
else begin
|
|
canceling := true ;
|
|
LC.broadcast cancelation () ;
|
|
!cancel_hook () >>= fun () ->
|
|
canceled := true ;
|
|
LC.broadcast cancelation_complete () ;
|
|
Lwt.return ()
|
|
end
|
|
in
|
|
let on_cancel cb =
|
|
let hook = !cancel_hook in
|
|
cancel_hook := (fun () -> hook () >>= cb) ;
|
|
in
|
|
let cancelation () =
|
|
if !canceling then Lwt.return ()
|
|
else LC.wait cancelation
|
|
in
|
|
cancelation, cancel, on_cancel
|
|
|
|
type trigger =
|
|
| Absent
|
|
| Present
|
|
| Waiting of unit Lwt.u
|
|
|
|
let trigger () : (unit -> unit) * (unit -> unit Lwt.t) =
|
|
let state = ref Absent in
|
|
let trigger () =
|
|
match !state with
|
|
| Absent -> state := Present
|
|
| Present -> ()
|
|
| Waiting u ->
|
|
state := Absent;
|
|
Lwt.wakeup u ()
|
|
in
|
|
let wait () =
|
|
match !state with
|
|
| Absent ->
|
|
let waiter, u = Lwt.wait () in
|
|
state := Waiting u;
|
|
waiter
|
|
| Present ->
|
|
state := Absent;
|
|
Lwt.return_unit
|
|
| Waiting u ->
|
|
Lwt.waiter_of_wakener u
|
|
in
|
|
trigger, wait
|
|
|
|
type 'a queue =
|
|
| Absent
|
|
| Present of 'a list ref
|
|
| Waiting of 'a list Lwt.u
|
|
|
|
let queue () : ('a -> unit) * (unit -> 'a list Lwt.t) =
|
|
let state = ref Absent in
|
|
let queue v =
|
|
match !state with
|
|
| Absent -> state := Present (ref [v])
|
|
| Present r -> r := v :: !r
|
|
| Waiting u ->
|
|
state := Absent;
|
|
Lwt.wakeup u [v]
|
|
in
|
|
let wait () =
|
|
match !state with
|
|
| Absent ->
|
|
let waiter, u = Lwt.wait () in
|
|
state := Waiting u;
|
|
waiter
|
|
| Present r ->
|
|
state := Absent;
|
|
Lwt.return (List.rev !r)
|
|
| Waiting u ->
|
|
Lwt.waiter_of_wakener u
|
|
in
|
|
queue, wait
|
|
|
|
(* A worker launcher, takes a cancel callback to call upon *)
|
|
let worker ?(safe=false) name ~run ~cancel =
|
|
let stop = LC.create () in
|
|
let fail e =
|
|
log_error "%s worker failed with %s" name (Printexc.to_string e) ;
|
|
cancel () >>= fun () ->
|
|
if safe then Lwt.return_unit else Lwt.fail e
|
|
in
|
|
let waiter = LC.wait stop in
|
|
log_info "%s worker started" name ;
|
|
Lwt.async
|
|
(fun () ->
|
|
Lwt.catch run fail >>= fun () ->
|
|
LC.signal stop ();
|
|
Lwt.return ()) ;
|
|
waiter >>= fun () ->
|
|
log_info "%s worker ended" name ;
|
|
Lwt.return ()
|
|
|
|
|
|
let rec chop k l =
|
|
if k = 0 then l else begin
|
|
match l with
|
|
| _::t -> chop (k-1) t
|
|
| _ -> assert false
|
|
end
|
|
let stable_sort cmp l =
|
|
let rec rev_merge l1 l2 accu =
|
|
match l1, l2 with
|
|
| [], l2 -> Lwt.return (List.rev_append l2 accu)
|
|
| l1, [] -> Lwt.return (List.rev_append l1 accu)
|
|
| h1::t1, h2::t2 ->
|
|
cmp h1 h2 >>= function
|
|
| x when x <= 0 -> rev_merge t1 l2 (h1::accu)
|
|
| _ -> rev_merge l1 t2 (h2::accu)
|
|
in
|
|
let rec rev_merge_rev l1 l2 accu =
|
|
match l1, l2 with
|
|
| [], l2 -> Lwt.return (List.rev_append l2 accu)
|
|
| l1, [] -> Lwt.return (List.rev_append l1 accu)
|
|
| h1::t1, h2::t2 ->
|
|
cmp h1 h2 >>= function
|
|
| x when x > 0 -> rev_merge_rev t1 l2 (h1::accu)
|
|
| _ -> rev_merge_rev l1 t2 (h2::accu)
|
|
in
|
|
let rec sort n l =
|
|
match n, l with
|
|
| 2, x1 :: x2 :: _ -> begin
|
|
cmp x1 x2 >|= function
|
|
| x when x <= 0 -> [x1; x2]
|
|
| _ -> [x2; x1]
|
|
end
|
|
| 3, x1 :: x2 :: x3 :: _ -> begin
|
|
cmp x1 x2 >>= function
|
|
| x when x <= 0 -> begin
|
|
cmp x2 x3 >>= function
|
|
| x when x <= 0 -> Lwt.return [x1; x2; x3]
|
|
| _ -> cmp x1 x3 >|= function
|
|
| x when x <= 0 -> [x1; x3; x2]
|
|
| _ -> [x3; x1; x2]
|
|
end
|
|
| _ -> begin
|
|
cmp x1 x3 >>= function
|
|
| x when x <= 0 -> Lwt.return [x2; x1; x3]
|
|
| _ -> cmp x2 x3 >|= function
|
|
| x when x <= 0 -> [x2; x3; x1]
|
|
| _ -> [x3; x2; x1]
|
|
end
|
|
end
|
|
| n, l ->
|
|
let n1 = n asr 1 in
|
|
let n2 = n - n1 in
|
|
let l2 = chop n1 l in
|
|
rev_sort n1 l >>= fun s1 ->
|
|
rev_sort n2 l2 >>= fun s2 ->
|
|
rev_merge_rev s1 s2 []
|
|
and rev_sort n l =
|
|
match n, l with
|
|
| 2, x1 :: x2 :: _ -> begin
|
|
cmp x1 x2 >|= function
|
|
| x when x > 0 -> [x1; x2]
|
|
| _ -> [x2; x1]
|
|
end
|
|
| 3, x1 :: x2 :: x3 :: _ -> begin
|
|
cmp x1 x2 >>= function
|
|
| x when x > 0 -> begin
|
|
cmp x2 x3 >>= function
|
|
| x when x > 0 -> Lwt.return [x1; x2; x3]
|
|
| _ ->
|
|
cmp x1 x3 >|= function
|
|
| x when x > 0 -> [x1; x3; x2]
|
|
| _ -> [x3; x1; x2]
|
|
end
|
|
| _ -> begin
|
|
cmp x1 x3 >>= function
|
|
| x when x > 0 -> Lwt.return [x2; x1; x3]
|
|
| _ ->
|
|
cmp x2 x3 >|= function
|
|
| x when x > 0 -> [x2; x3; x1]
|
|
| _ -> [x3; x2; x1]
|
|
end
|
|
end
|
|
| n, l ->
|
|
let n1 = n asr 1 in
|
|
let n2 = n - n1 in
|
|
let l2 = chop n1 l in
|
|
sort n1 l >>= fun s1 ->
|
|
sort n2 l2 >>= fun s2 ->
|
|
rev_merge s1 s2 []
|
|
in
|
|
let len = List.length l in
|
|
if len < 2 then Lwt.return l else sort len l
|
|
|
|
let sort = stable_sort
|
|
|
|
let rec read_bytes ?(pos = 0) ?len fd buf =
|
|
let len = match len with None -> Bytes.length buf - pos | Some l -> l in
|
|
let rec inner pos len =
|
|
if len = 0 then
|
|
Lwt.return_unit
|
|
else
|
|
Lwt_unix.read fd buf pos len >>= function
|
|
| 0 -> Lwt.fail End_of_file (* other endpoint cleanly closed its connection *)
|
|
| nb_read -> inner (pos + nb_read) (len - nb_read)
|
|
in
|
|
inner pos len
|
|
|
|
let read_mbytes ?(pos=0) ?len fd buf =
|
|
let len = match len with None -> MBytes.length buf - pos | Some l -> l in
|
|
let rec inner pos len =
|
|
if len = 0 then
|
|
Lwt.return_unit
|
|
else
|
|
Lwt_bytes.read fd buf pos len >>= function
|
|
| 0 -> Lwt.fail End_of_file (* other endpoint cleanly closed its connection *)
|
|
| nb_read -> inner (pos + nb_read) (len - nb_read)
|
|
in
|
|
inner pos len
|
|
|
|
let write_mbytes ?(pos=0) ?len descr buf =
|
|
let len = match len with None -> MBytes.length buf - pos | Some l -> l in
|
|
let rec inner pos len =
|
|
if len = 0 then
|
|
Lwt.return_unit
|
|
else
|
|
Lwt_bytes.write descr buf pos len >>= function
|
|
| 0 -> Lwt.fail End_of_file (* other endpoint cleanly closed its connection *)
|
|
| nb_written -> inner (pos + nb_written) (len - nb_written) in
|
|
inner pos len
|
|
|