refactor(local_storage): use a local reference

this allows to not modify the map as much. We know that a given
reference will only ever be accessed by its thread (because it's the
only one to use that key). the map (and atomic ref) is only modified
when `set` is called on a key that is not used at all in this TLS
variable.

src/threads/CCLocal_storage.ml

@@ -27,63 +27,78 @@ end)
 
 [@@@endif]
 
-type 'a t = 'a Key_map_.t A.t
+type 'a t = 'a option ref Key_map_.t A.t
 
 let create () : _ t = A.make Key_map_.empty
 
 let[@inline] n_entries self = Key_map_.cardinal (A.get self)
 
-let get (self: _ t ) : _ option =
-  let m = A.get self in
-  let key = get_key_ () in
-  Key_map_.get key m
-
 let get_exn self =
   let m = A.get self in
   let key = get_key_ () in
-  Key_map_.find key m
+  let r = Key_map_.find key m in
+  match !r with
+  | Some x -> x
+  | None -> raise Not_found
 
-let get_or ~default self =
+let[@inline] get self =
+  try Some (get_exn self)
+  with Not_found -> None
+
+let[@inline] get_or ~default self =
   try get_exn self
   with Not_found -> default
 
-let set (self: _ t ) v : unit =
-  let key = get_key_ () in
-  while
-    let m = A.get self in
-    let m' = Key_map_.add key v m in
-    not (A.compare_and_set self m m')
-  do () done
-
-let set_get (self: _ t ) v : _ option =
-  let key = get_key_ () in
-  let rec loop () =
-    let m = A.get self in
-    let m' = Key_map_.add key v m in
-    if A.compare_and_set self m m' then Key_map_.get key m
-    else loop()
-  in loop ()
-
-let remove self =
-  let key = get_key_ () in
+(* remove reference for the key *)
+let[@inline] remove_ref_ self key =
   while
     let m = A.get self in
     let m' = Key_map_.remove key m in
     not (A.compare_and_set self m m')
-  do () done
+  do() done
 
-let[@inline] set_opt_ self v =
-  match v with
-  | None -> remove self
-  | Some v' -> set self v'
+type get_or_create_status =
+  | Created
+  | Reused
+
+(* get or associate a reference to [key], and return it.
+   Also return a function to remove the reference if we just created it. *)
+let get_or_create_ref_ (self:_ t ) key : _ ref * get_or_create_status =
+  try Key_map_.find key (A.get self), Reused
+  with Not_found ->
+    let r = ref None in
+    while
+      let m = A.get self in
+      let m' = Key_map_.add key r m in
+      not (A.compare_and_set self m m')
+    do () done;
+    r, Created
+
+let set (self: _ t ) v : unit =
+  let key = get_key_ () in
+  let r, _ = get_or_create_ref_ self key in
+  r := Some v
+
+let set_get (self: _ t ) v : _ option =
+  let key = get_key_ () in
+  let r, _ = get_or_create_ref_ self key in
+  let old = !r in
+  r := Some v;
+  old
 
 let with_ self x f =
-  let old = set_get self x in
+  let key = get_key_ () in
+  let r, status = get_or_create_ref_ self key in
+  let old = !r in
+  r := Some x;
+
   try
-    let r = f() in
-    set_opt_ self old;
-    r
+    let res = f () in
+    r := old;
+    if status == Created then remove_ref_ self key;
+    res
   with e ->
-    let bt = Printexc.get_raw_backtrace () in
-    set_opt_ self old;
-    Printexc.raise_with_backtrace e bt
+    r := old;
+    if status == Created then remove_ref_ self key;
+    raise e
+

src/threads/CCLocal_storage.mli

@@ -29,9 +29,6 @@ val set : 'a t -> 'a -> unit
 val set_get : 'a t -> 'a -> 'a option
 (** Set content for this thread, and return the old value. *)
 
-val remove : 'a t -> unit
-(** Remove value *)
-
 val n_entries : _ t -> int
 (** Number of entries in the map currently.