From 839eb3fcdf33c1b0d116451fb6a04079c4f18d62 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Mon, 9 Sep 2024 14:44:23 -0400
Subject: [PATCH] feat tef-tldrs: expose a subscriber

---
 src/tef-tldrs/dune                | 15 ++++++++++-----
 src/tef-tldrs/trace_tef_tldrs.ml  | 19 +++++++++++++------
 src/tef-tldrs/trace_tef_tldrs.mli |  8 +++++++-
 3 files changed, 30 insertions(+), 12 deletions(-)

diff --git a/src/tef-tldrs/dune b/src/tef-tldrs/dune
index 9f99684..7bd7fab 100644
--- a/src/tef-tldrs/dune
+++ b/src/tef-tldrs/dune
@@ -1,6 +1,11 @@
-
 (library
-  (name trace_tef_tldrs)
-  (public_name trace-tef.tldrs)
-  (synopsis "Multiprocess tracing using the `tldrs` daemon")
-  (libraries trace.core trace.private.util trace-tef unix threads))
+ (name trace_tef_tldrs)
+ (public_name trace-tef.tldrs)
+ (synopsis "Multiprocess tracing using the `tldrs` daemon")
+ (libraries
+  trace.core
+  trace.private.util
+  trace-subscriber
+  trace-tef
+  unix
+  threads))
diff --git a/src/tef-tldrs/trace_tef_tldrs.ml b/src/tef-tldrs/trace_tef_tldrs.ml
index ae09334..90b4e9f 100644
--- a/src/tef-tldrs/trace_tef_tldrs.ml
+++ b/src/tef-tldrs/trace_tef_tldrs.ml
@@ -84,7 +84,7 @@ let find_role ~out () : role =
       | Some path -> Some (write_to_file path)
       | None -> None))
 
-let collector_ (client : as_client) : collector =
+let subscriber_ (client : as_client) : Trace_subscriber.t =
   (* connect to unix socket *)
   let sock = Unix.socket Unix.PF_UNIX Unix.SOCK_STREAM 0 in
   (try Unix.connect sock (Unix.ADDR_UNIX client.socket)
@@ -105,24 +105,31 @@ let collector_ (client : as_client) : collector =
   in
 
   fpf out "OPEN %s\n%!" client.trace_id;
-  Trace_tef.Internal_.collector_jsonl ~finally ~out:(`Output out) ()
+  Trace_tef.Private_.subscriber_jsonl ~finally ~out:(`Output out) ()
+
+let subscriber ~out () =
+  let role = find_role ~out () in
+  match role with
+  | None -> assert false
+  | Some c -> subscriber_ c
 
 let collector ~out () : collector =
   let role = find_role ~out () in
   match role with
   | None -> assert false
-  | Some c -> collector_ c
+  | Some c -> subscriber_ c |> Trace_subscriber.collector
 
 let setup ?(out = `Env) () =
   let role = find_role ~out () in
   match role with
   | None -> ()
-  | Some c -> Trace_core.setup_collector @@ collector_ c
+  | Some c ->
+    Trace_core.setup_collector @@ Trace_subscriber.collector @@ subscriber_ c
 
 let with_setup ?out () f =
   setup ?out ();
   Fun.protect ~finally:Trace_core.shutdown f
 
-module Internal_ = struct
-  include Trace_tef.Internal_
+module Private_ = struct
+  include Trace_tef.Private_
 end
diff --git a/src/tef-tldrs/trace_tef_tldrs.mli b/src/tef-tldrs/trace_tef_tldrs.mli
index dc67648..e9e8132 100644
--- a/src/tef-tldrs/trace_tef_tldrs.mli
+++ b/src/tef-tldrs/trace_tef_tldrs.mli
@@ -1,7 +1,13 @@
+(** Emit traces by talking to the {{: https://github.com/imandra-ai/tldrs} tldrs} daemon *)
+
 val collector : out:[ `File of string ] -> unit -> Trace_core.collector
 (** Make a collector that writes into the given output.
     See {!setup} for more details. *)
 
+val subscriber : out:[ `File of string ] -> unit -> Trace_subscriber.t
+(** Make a subscriber that writes into the given output.
+    @since NEXT_RELEASE *)
+
 type output = [ `File of string ]
 (** Output for tracing.
     - [`File "foo"] will enable tracing and print events into file
@@ -30,7 +36,7 @@ val with_setup : ?out:[ output | `Env ] -> unit -> (unit -> 'a) -> 'a
 
 (**/**)
 
-module Internal_ : sig
+module Private_ : sig
   val mock_all_ : unit -> unit
   (** use fake, deterministic timestamps, TID, PID *)