Simple resource pooling with Async

async_pool.ml

open Core.Std
open Async.Std

type 'a t = {
  capacity : int;
  generate : unit -> 'a Deferred.t;
  validate : 'a -> bool Deferred.t;
  mutable in_use : int;
  free : 'a Queue.t;
  waiters : 'a Ivar.t Queue.t;
}

let create ?validate ~generate capacity = {
  capacity;
  generate;
  validate = Option.value validate ~default:(const (return true));
  in_use   = 0;
  free     = Queue.create ();
  waiters  = Queue.create ();
}

let acquire t =
  let work t =
    match Queue.dequeue t.free with
    | None ->
      if t.in_use < t.capacity then begin
        (* initiate new member *)
        t.in_use <- t.in_use + 1;
        Monitor.try_with ~run:`Now ~rest:`Raise t.generate >>|
        function
        | Ok member -> member
        | Error e ->
          t.in_use <- t.in_use - 1;
          raise e
      end else begin
        let i = Ivar.create () in
        Queue.enqueue t.waiters i;
        Ivar.read i
      end
    | Some member ->
      return member
  in
  let rec loop t =
    work t
    >>= fun member -> t.validate member
    >>= fun is_valid ->
    if is_valid
    then return member
    else begin
      t.in_use <- t.in_use - 1;
      loop t
    end
  in
  loop t

let release t member =
  match Queue.dequeue t.waiters with
  | None ->
    Queue.enqueue t.free member
  | Some i ->
    Ivar.fill i member