thread-local: expose set/remove

src/thread_local.ml

@@ -40,7 +40,7 @@ let remove_ref_ self key : unit =
     Thread.yield ()
   done
 
-let set_ self key (r : _ ref) : unit =
+let set_ref_ self key (r : _ ref) : unit =
   while
     let m = A.get self in
     let m' = Key_map_.add key r m in
@@ -59,9 +59,18 @@ let get_or_create_ref_ (self : _ t) key ~v : _ ref * _ option =
     r, Some old
   with Not_found ->
     let r = ref v in
-    set_ self key r;
+    set_ref_ self key r;
     r, None
 
+let set (self : _ t) v : unit =
+  let key = get_key_ () in
+  let _, _ = get_or_create_ref_ self key ~v in
+  ()
+
+let remove (self : _ t) : unit =
+  let key = get_key_ () in
+  remove_ref_ self key
+
 let get_or_create ~create (self : 'a t) : 'a =
   let key = get_key_ () in
   try
@@ -71,7 +80,7 @@ let get_or_create ~create (self : 'a t) : 'a =
     Gc.finalise (fun _ -> remove_ref_ self key) (Thread.self ());
     let v = create () in
     let r = ref v in
-    set_ self key r;
+    set_ref_ self key r;
     v
 
 let with_ self v f =

src/thread_local.mli

@@ -15,6 +15,10 @@ val get_exn : 'a t -> 'a
 (** Like {!get} but fails with an exception
     @raise Not_found if no value was found *)
 
+val set : 'a t -> 'a -> unit
+
+val remove : _ t -> unit
+
 val get_or_create : create:(unit -> 'a) -> 'a t -> 'a
 
 val with_ : 'a t -> 'a -> ('a option -> 'b) -> 'b