Merge pull request #1 from ocaml-multicore/better_teardown

Clean up better at teardown.

Author: Sudha247

lib/task.ml

@@ -8,9 +8,13 @@ type task_msg =
   Task : 'a task * 'a promise -> task_msg
 | Quit : task_msg
 
-type pool =
-  {domains : unit Domain.t array;
-   task_chan : task_msg Multi_channel.t}
+type pool_data = {
+  domains : unit Domain.t array;
+  task_chan : task_msg Multi_channel.t;
+  name: string option
+}
+
+type pool = pool_data option Atomic.t
 
 let do_task f p =
   try
@@ -24,7 +28,7 @@ let do_task f p =
 
 let named_pools = Hashtbl.create 8
 
-let named_pool_mutex = Mutex.create ()
+let named_pools_mutex = Mutex.create ()
 
 let setup_pool ?name ~num_additional_domains () =
   if num_additional_domains < 0 then
@@ -40,27 +44,34 @@ let setup_pool ?name ~num_additional_domains () =
         worker ()
   in
   let domains = Array.init num_additional_domains (fun _ -> Domain.spawn worker) in
-  let p = {domains; task_chan} in
-  let _ = match name with
-    | Some x ->
-      Mutex.lock named_pool_mutex;
-      Hashtbl.add named_pools x p;
-      Mutex.unlock named_pool_mutex
+  let p = Atomic.make (Some {domains; task_chan; name}) in
+  begin match name with
     | None -> ()
-  in
+    | Some x ->
+        Mutex.lock named_pools_mutex;
+        Hashtbl.add named_pools x p;
+        Mutex.unlock named_pools_mutex
+  end;
   p
 
+let get_pool_data p =
+  match Atomic.get p with
+  | None -> raise (Invalid_argument "pool already torn down")
+  | Some p -> p
+
 let async pool task =
+  let pd = get_pool_data pool in
   let p = Atomic.make None in
-  Multi_channel.send pool.task_chan (Task(task,p));
+  Multi_channel.send pd.task_chan (Task(task,p));
   p
 
 let rec await pool promise =
+  let pd = get_pool_data pool in
   match Atomic.get promise with
   | None ->
       begin
         try
-          match Multi_channel.recv_poll pool.task_chan with
+          match Multi_channel.recv_poll pd.task_chan with
           | Task (t, p) -> do_task t p
           | Quit -> raise TasksActive
         with
@@ -71,24 +82,37 @@ let rec await pool promise =
   | Some (Error e) -> raise e
 
 let teardown_pool pool =
-  for _i=1 to Array.length pool.domains do
-    Multi_channel.send pool.task_chan Quit
+  let pd = get_pool_data pool in
+  for _i=1 to Array.length pd.domains do
+    Multi_channel.send pd.task_chan Quit
   done;
   Multi_channel.clear_local_state ();
-  Array.iter Domain.join pool.domains
+  Array.iter Domain.join pd.domains;
+  (* Remove the pool from the table *)
+  begin match pd.name with
+  | None -> ()
+  | Some n ->
+      Mutex.lock named_pools_mutex;
+      Hashtbl.remove named_pools n;
+      Mutex.unlock named_pools_mutex
+  end;
+  Atomic.set pool None
 
 let lookup_pool name =
-  Mutex.lock named_pool_mutex;
+  Mutex.lock named_pools_mutex;
   let p = Hashtbl.find_opt named_pools name in
-  Mutex.unlock named_pool_mutex;
+  Mutex.unlock named_pools_mutex;
   p
 
-let get_num_domains pool = (Array.length pool.domains + 1)
+let get_num_domains pool =
+  let pd = get_pool_data pool in
+  Array.length pd.domains + 1
 
 let parallel_for_reduce ?(chunk_size=0) ~start ~finish ~body pool reduce_fun init =
+  let pd = get_pool_data pool in
   let chunk_size = if chunk_size > 0 then chunk_size
       else begin
-        let n_domains = (Array.length pool.domains) + 1 in
+        let n_domains = (Array.length pd.domains) + 1 in
         let n_tasks = finish - start + 1 in
         if n_domains = 1 then n_tasks
         else max 1 (n_tasks/(8*n_domains))
@@ -112,9 +136,10 @@ let parallel_for_reduce ?(chunk_size=0) ~start ~finish ~body pool reduce_fun ini
   reduce_fun init (work start finish)
 
 let parallel_for ?(chunk_size=0) ~start ~finish ~body pool =
+  let pd = get_pool_data pool in
   let chunk_size = if chunk_size > 0 then chunk_size
       else begin
-        let n_domains = (Array.length pool.domains) + 1 in
+        let n_domains = (Array.length pd.domains) + 1 in
         let n_tasks = finish - start + 1 in
         if n_domains = 1 then n_tasks
         else max 1 (n_tasks/(8*n_domains))
@@ -133,7 +158,7 @@ let parallel_for ?(chunk_size=0) ~start ~finish ~body pool =
   work pool body start finish
 
 let parallel_scan pool op elements =
-
+  let pd = get_pool_data pool in
   let scan_part op elements prefix_sum start finish =
     assert (Array.length elements > (finish - start));
     for i = (start + 1) to finish do
@@ -147,7 +172,7 @@ let parallel_scan pool op elements =
     done
   in
   let n = Array.length elements in
-  let p = (Array.length pool.domains) + 1 in
+  let p = (Array.length pd.domains) + 1 in
   let prefix_s = Array.copy elements in
 
   parallel_for pool ~chunk_size:1 ~start:0 ~finish:(p - 1)

test/test_task.ml

@@ -57,7 +57,15 @@ let () =
   prefix_sum p2 ();
   Task.teardown_pool pool1;
   Task.teardown_pool pool2;
+
+  try
+    sum_sequence pool2 0 0 ();
+    assert false
+  with Invalid_argument _ -> ();
+
+  assert (Task.lookup_pool "pool1" = None);
+
   try
     let _ = Task.setup_pool ~num_additional_domains:(-1) () in ()
   with Invalid_argument _ -> ();
-  print_endline "ok"
+  print_endline "ok"