Fix memory leak

dscheck.opam

@@ -11,6 +11,7 @@ depends: [
   "dune" {>= "2.9"}
   "containers"
   "oseq"
+  "alcotest" {>= "1.6.0"}
   "odoc" {with-doc}
 ]
 build: [

src/tracedAtomic.ml

@@ -48,6 +48,7 @@ type process_data = {
   mutable next_repr: int option;
   mutable resume_func : (unit, unit) handler -> unit;
   mutable finished : bool;
+  mutable discontinue_f : unit -> unit;
 }
 
 let every_func = ref (fun () -> ())
@@ -82,18 +83,33 @@ let incr r = ignore (fetch_and_add r 1)
 let decr r = ignore (fetch_and_add r (-1))
 
 (* Tracing infrastructure *)
+
+exception Terminated_early
+
+let discontinue k () = 
+  discontinue_with k Terminated_early ({
+    retc = (fun _ -> ()); 
+    exnc = (function Terminated_early -> () | e -> raise e); 
+    effc = (fun (type a) (_ : a Effect.t) -> None);
+  })
+;;
+
 let processes = CCVector.create ()
 
-let update_process_data process_id f op repr =
+let update_process_data process_id f op repr k =
   let process_rec = CCVector.get processes process_id in
   process_rec.resume_func <- f;
   process_rec.next_repr <- repr;
-  process_rec.next_op <- op
+  process_rec.next_op <- op;
+  process_rec.discontinue_f <- discontinue k
+;;
 
 let finish_process process_id =
   let process_rec = CCVector.get processes process_id in
   process_rec.finished <- true;
+  process_rec.discontinue_f <- (fun () -> ());
   finished_processes := !finished_processes + 1
+;;
 
 let handler current_process_id runner =
   {
@@ -112,44 +128,44 @@ let handler current_process_id runner =
                 let i = !atomics_counter in
                 let m = (Atomic.make v, i) in
                 atomics_counter := !atomics_counter + 1;
-                update_process_data current_process_id (fun h -> continue_with k m h) Make (Some i);
+                update_process_data current_process_id (fun h -> continue_with k m h) Make (Some i) k;
                 runner ())
          | Get (v,i) ->
            Some
              (fun (k : (a, _) continuation) ->
-                update_process_data current_process_id (fun h -> continue_with k (Atomic.get v) h) Get (Some i);
+                update_process_data current_process_id (fun h -> continue_with k (Atomic.get v) h) Get (Some i) k;
                 runner ())
          | Set ((r,i), v) ->
            Some
              (fun (k : (a, _) continuation) ->
-                update_process_data current_process_id (fun h -> continue_with k (Atomic.set r v) h) Set (Some i);
+                update_process_data current_process_id (fun h -> continue_with k (Atomic.set r v) h) Set (Some i) k;
                 runner ())
          | Exchange ((a,i), b) ->
            Some
              (fun (k : (a, _) continuation) ->
-                update_process_data current_process_id (fun h -> continue_with k (Atomic.exchange a b) h) Exchange (Some i);
+                update_process_data current_process_id (fun h -> continue_with k (Atomic.exchange a b) h) Exchange (Some i) k;
                 runner ())
          | CompareAndSwap ((x,i), s, v) ->
            Some
              (fun (k : (a, _) continuation) ->
                 update_process_data current_process_id (fun h ->
-                    continue_with k (Atomic.compare_and_set x s v) h) CompareAndSwap (Some i);
+                    continue_with k (Atomic.compare_and_set x s v) h) CompareAndSwap (Some i) k;
                 runner ())
          | FetchAndAdd ((v,i), x) ->
            Some
              (fun (k : (a, _) continuation) ->
                 update_process_data current_process_id (fun h ->
-                    continue_with k (Atomic.fetch_and_add v x) h) FetchAndAdd (Some i);
+                    continue_with k (Atomic.fetch_and_add v x) h) FetchAndAdd (Some i) k;
                 runner ())
          | _ ->
            None);
   }
 
 let spawn f =
-  let fiber_f h =
-    continue_with (fiber f) () h in
+  let fiber_f = fiber f in
+  let resume_func = continue_with fiber_f () in
   CCVector.push processes
-    { next_op = Start; next_repr = None; resume_func = fiber_f; finished = false }
+    { next_op = Start; next_repr = None; resume_func; finished = false; discontinue_f = discontinue fiber_f}
 
 let rec last_element l =
   match l with
@@ -208,6 +224,7 @@ let do_run init_func init_schedule =
                         |> Seq.filter (fun (_,proc) -> not proc.finished)
                         |> Seq.map (fun (id,_) -> id)
                         |> IntSet.of_seq in
+  CCVector.iter (fun proc -> proc.discontinue_f ()) processes;
   CCVector.clear processes;
   atomics_counter := 1;
   match last_element init_schedule with

tests/dune

@@ -7,3 +7,8 @@
  (name test_naive_counter)
  (libraries dscheck alcotest)
  (modules test_naive_counter))
+
+(test
+ (name test_michael_scott_queue)
+ (libraries dscheck alcotest)
+ (modules test_michael_scott_queue michael_scott_queue))

tests/michael_scott_queue.ml

@@ -0,0 +1,62 @@
+(*
+ * Copyright (c) 2015, Théo Laurent <theo.laurent@ens.fr>
+ * Copyright (c) 2015, KC Sivaramakrishnan <sk826@cl.cam.ac.uk>
+ *
+ * Permission to use, copy, modify, and/or distribute this software for any
+ * purpose with or without fee is hereby granted, provided that the above
+ * copyright notice and this permission notice appear in all copies.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
+ * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+ * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
+ * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+ * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
+ * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
+ * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+ *)
+
+(* Michael-Scott queue copied over from [lockfree](github.com/ocaml-multicore/lockfree) *)
+module Atomic = Dscheck.TracedAtomic
+
+type 'a node = Nil | Next of 'a * 'a node Atomic.t
+type 'a t = { head : 'a node Atomic.t; tail : 'a node Atomic.t }
+
+let create () =
+  let head = Next (Obj.magic (), Atomic.make Nil) in
+  { head = Atomic.make head; tail = Atomic.make head }
+
+let is_empty q =
+  match Atomic.get q.head with
+  | Nil -> failwith "MSQueue.is_empty: impossible"
+  | Next (_, x) -> ( match Atomic.get x with Nil -> true | _ -> false)
+
+let pop q =
+  let rec loop () =
+    let s = Atomic.get q.head in
+    let nhead =
+      match s with
+      | Nil -> failwith "MSQueue.pop: impossible"
+      | Next (_, x) -> Atomic.get x
+    in
+    match nhead with
+    | Nil -> None
+    | Next (v, _) when Atomic.compare_and_set q.head s nhead -> Some v
+    | _ -> loop ()
+  in
+  loop ()
+
+let push q v =
+  let rec find_tail_and_enq curr_end node =
+    if Atomic.compare_and_set curr_end Nil node then ()
+    else
+      match Atomic.get curr_end with
+      | Nil -> find_tail_and_enq curr_end node
+      | Next (_, n) -> find_tail_and_enq n node
+  in
+  let newnode = Next (v, Atomic.make Nil) in
+  let tail = Atomic.get q.tail in
+  match tail with
+  | Nil -> failwith "HW_MSQueue.push: impossible"
+  | Next (_, n) ->
+      find_tail_and_enq n newnode;
+      ignore (Atomic.compare_and_set q.tail tail newnode)

tests/test_list.ml

@@ -50,8 +50,8 @@ let test_list_naive domains () =
   | exception _ -> ()
   | _ -> failwith "expected failure"
 
-let test_list_safe () = Atomic.trace (create_test add_node_safe 3)
-(* 4 takes over 10 Gb of RAM *)
+let test_list_safe () = 
+  Atomic.trace (create_test add_node_safe 3)
 
 let () =
   let open Alcotest in
@@ -60,7 +60,7 @@ let () =
       ( "list",
         [
           test_case "naive-1-domain" `Quick (test_list_naive_single_domain);
-          test_case "naive-2-domains" `Quick (test_list_naive 8);
+          test_case "naive-2-domains" `Quick (test_list_naive 2);
           test_case "naive-8-domains" `Quick (test_list_naive 8);
           test_case "safe" `Quick test_list_safe;
         ] );

tests/test_michael_scott_queue.ml

@@ -0,0 +1,39 @@
+module Atomic = Dscheck.TracedAtomic
+
+let drain queue =
+  let remaining = ref 0 in
+  while not (Michael_scott_queue.is_empty queue) do
+    remaining := !remaining + 1;
+    assert (Option.is_some (Michael_scott_queue.pop queue))
+  done;
+  !remaining
+
+let producer_consumer () =
+  Atomic.trace (fun () ->
+      let queue = Michael_scott_queue.create () in
+      let items_total = 4 in
+
+      Atomic.spawn (fun () ->
+          for _ = 1 to items_total do
+            Michael_scott_queue.push queue 0
+          done);
+
+      (* consumer *)
+      let popped = ref 0 in
+      Atomic.spawn (fun () ->
+          for _ = 1 to items_total do
+            match Michael_scott_queue.pop queue with
+            | None -> ()
+            | Some _ -> popped := !popped + 1
+          done);
+
+      (* checks*)
+      Atomic.final (fun () ->
+          Atomic.check (fun () ->
+              let remaining = drain queue in
+              !popped + remaining = items_total)))
+
+let () =
+  let open Alcotest in
+  run "michael_scott_queue_dscheck"
+    [ ("basic", [ test_case "1-producer-1-consumer" `Slow producer_consumer ]) ]