Wrapper around Lwt threads that allows carrying an environment around, transparently.

lwt_env.ml

(* ocamlfind ocamlc -c env.ml -package lwt.unix *)

(*
   Wrapper around Lwt threads that allows
   carrying an environment around, transparently.

   Such an environment would typically consist of a request ID
   that we can use in logs to identify all messages relating to
   the same request received by the server.

   Limitations:
   - was never used
   - will not work with system threads as is (current_env needs to be local
     to the current system thread)
   - would have to be functorized to accept an environment of any type
*)

type env = (string, string) Hashtbl.t option
  (* The type of the environment. Can be anything, should be
     a parameter to a functor. *)

let default_env = None
  (* The default environment to return when outside of a 'bind' context.
     Could be anything of type 'env'. *)

module type Env = sig
  type 'a t
    (* The type of a wrapped lwt thread. *)

  val return : 'a -> 'a t
  val bind : 'a t -> ('a -> 'b t) -> 'b t
    (* Replacements for 'Lwt.bind' and 'Lwt.return' *)

  val get : unit -> env
    (* Return the environment inherited from the lwt thread
       we are binding to. *)

  val construct : 'a Lwt.t -> env -> 'a t
    (* Use this in the server loop to create an environment for the
       thread representing the request. *)

  val deconstruct : 'a t -> 'a Lwt.t * env
    (* Recover lwt thread and environment associated with it. *)
end

module Env : Env = struct
  type 'a t = 'a Lwt.t * env

  let current_env = ref None
    (* Not systhread-safe. Would need thread-local storage. *)

  let get () =
    match !current_env with
    | None -> default_env
    | Some env -> env

  let return x =
    (Lwt.return x, get ())

  let bind (at, env) f =
    let new_env = ref env in
    let bt =
      Lwt.bind at (fun a ->
        current_env := Some env;
        let bt, env = f a in
        new_env := env;
        current_env := None;
        bt
      )
    in
    (bt, !new_env)

  let construct lwt env = (lwt, env)
  let deconstruct x = x
end