Add hs_thread_done() (#8124)

See documentation for details.

4 files changed, 78 insertions, 11 deletions

diff --git a/rts/HsFFI.c b/rts/HsFFI.c
index 856536f5aa..8fae246111 100644
--- a/rts/HsFFI.c
+++ b/rts/HsFFI.c
@@ -11,6 +11,7 @@
 #include "Rts.h"
 
 #include "Stable.h"
+#include "Task.h"
 
 // hs_init and hs_exit are defined in RtsStartup.c
 
@@ -59,3 +60,9 @@ hs_free_fun_ptr(HsFunPtr fp)
     /* I simply *love* all these similar names... */
     freeHaskellFunctionPtr(fp);
 }
+
+void
+hs_thread_done(void)
+{
+    freeMyTask();
+}
diff --git a/rts/RtsAPI.c b/rts/RtsAPI.c
index 725bfeb0b5..f01a0efee8 100644
--- a/rts/RtsAPI.c
+++ b/rts/RtsAPI.c
@@ -614,3 +614,9 @@ rts_unlock (Capability *cap)
       traceTaskDelete(task);
     }
 }
+
+void rts_done (void)
+{
+    freeMyTask();
+}
+
diff --git a/rts/Task.c b/rts/Task.c
index a044bc3672..12c22c4b02 100644
--- a/rts/Task.c
+++ b/rts/Task.c
@@ -134,6 +134,44 @@ allocTask (void)
     }
 }
 
+void freeMyTask (void)
+{
+    Task *task;
+
+    task = myTask();
+
+    if (task == NULL) return;
+
+    if (!task->stopped) {
+        errorBelch(
+            "freeMyTask() called, but the Task is not stopped; ignoring");
+        return;
+    }
+
+    if (task->worker) {
+        errorBelch("freeMyTask() called on a worker; ignoring");
+        return;
+    }
+
+    ACQUIRE_LOCK(&all_tasks_mutex);
+
+    if (task->all_prev) {
+        task->all_prev->all_next = task->all_next;
+    } else {
+        all_tasks = task->all_next;
+    }
+    if (task->all_next) {
+        task->all_next->all_prev = task->all_prev;
+    }
+
+    taskCount--;
+
+    RELEASE_LOCK(&all_tasks_mutex);
+
+    freeTask(task);
+    setMyTask(NULL);
+}
+
 static void
 freeTask (Task *task)
 {
@@ -219,7 +257,7 @@ newInCall (Task *task)
         task->spare_incalls = incall->next;
         task->n_spare_incalls--;
     } else {
-        incall = stgMallocBytes((sizeof(InCall)), "newBoundTask");
+        incall = stgMallocBytes((sizeof(InCall)), "newInCall");
     }
 
     incall->tso = NULL;
diff --git a/rts/Task.h b/rts/Task.h
index 4e0e13e93c..cf70256326 100644
--- a/rts/Task.h
+++ b/rts/Task.h
@@ -37,12 +37,20 @@
    Ownership of Task
    -----------------
 
-   The OS thread named in the Task structure has exclusive access to
-   the structure, as long as it is the running_task of its Capability.
-   That is, if (task->cap->running_task == task), then task->id owns
-   the Task.  Otherwise the Task is owned by the owner of the parent
-   data structure on which it is sleeping; for example, if the task is
-   sleeping on spare_workers field of a Capability, then the owner of the
+   Task ownership is a little tricky.  The default situation is that
+   the Task is an OS-thread-local structure that is owned by the OS
+   thread named in task->id.  An OS thread not currently executing
+   Haskell code might call newBoundTask() at any time, which assumes
+   that it has access to the Task for the current OS thread.
+
+   The all_next and all_prev fields of a Task are owned by
+   all_tasks_mutex, which must also be taken if we want to create or
+   free a Task.
+
+   For an OS thread in Haskell, if (task->cap->running_task != task),
+   then the Task is owned by the owner of the parent data structure on
+   which it is sleeping; for example, if the task is sleeping on
+   spare_workers field of a Capability, then the owner of the
    Capability has access to the Task.
 
    When a task is migrated from sleeping on one Capability to another,
@@ -147,7 +155,7 @@ typedef struct Task_ {
     // on spare_workers.
     struct Task_ *next;
 
-    // Links tasks on the all_tasks list
+    // Links tasks on the all_tasks list; need ACQUIRE_LOCK(&all_tasks_mutex)
     struct Task_ *all_next;
     struct Task_ *all_prev;
 
@@ -169,16 +177,24 @@ extern Task *all_tasks;
 void initTaskManager (void);
 nat  freeTaskManager (void);
 
-// Create a new Task for a bound thread
-// Requires: sched_mutex.
+// Create a new Task for a bound thread.  This Task must be released
+// by calling boundTaskExiting.  The Task is cached in
+// thread-local storage and will remain even after boundTaskExiting()
+// has been called; to free the memory, see freeMyTask().
 //
 Task *newBoundTask (void);
 
 // The current task is a bound task that is exiting.
-// Requires: sched_mutex.
 //
 void boundTaskExiting (Task *task);
 
+// Free a Task if one was previously allocated by newBoundTask().
+// This is not necessary unless the thread that called newBoundTask()
+// will be exiting, or if this thread has finished calling Haskell
+// functions.
+//
+void freeMyTask(void);
+
 // Notify the task manager that a task has stopped.  This is used
 // mainly for stats-gathering purposes.
 // Requires: sched_mutex.