+waitForReturnCapability (Capability **pCap, Task *task)
+{
+#if !defined(THREADED_RTS)
+
+ MainCapability.running_task = task;
+ task->cap = &MainCapability;
+ *pCap = &MainCapability;
+
+#else
+ Capability *cap = *pCap;
+
+ if (cap == NULL) {
+ // Try last_free_capability first
+ cap = last_free_capability;
+ if (!cap->running_task) {
+ nat i;
+ // otherwise, search for a free capability
+ for (i = 0; i < n_capabilities; i++) {
+ cap = &capabilities[i];
+ if (!cap->running_task) {
+ break;
+ }
+ }
+ // Can't find a free one, use last_free_capability.
+ cap = last_free_capability;
+ }
+
+ // record the Capability as the one this Task is now assocated with.
+ task->cap = cap;
+
+ } else {
+ ASSERT(task->cap == cap);
+ }
+
+ ACQUIRE_LOCK(&cap->lock);
+
+ IF_DEBUG(scheduler,
+ sched_belch("returning; I want capability %d", cap->no));
+
+ if (!cap->running_task) {
+ // It's free; just grab it
+ cap->running_task = task;
+ RELEASE_LOCK(&cap->lock);
+ } else {
+ newReturningTask(cap,task);
+ RELEASE_LOCK(&cap->lock);
+
+ for (;;) {
+ ACQUIRE_LOCK(&task->lock);
+ // task->lock held, cap->lock not held
+ if (!task->wakeup) waitCondition(&task->cond, &task->lock);
+ cap = task->cap;
+ task->wakeup = rtsFalse;
+ RELEASE_LOCK(&task->lock);
+
+ // now check whether we should wake up...
+ ACQUIRE_LOCK(&cap->lock);
+ if (cap->running_task == NULL) {
+ if (cap->returning_tasks_hd != task) {
+ giveCapabilityToTask(cap,cap->returning_tasks_hd);
+ RELEASE_LOCK(&cap->lock);
+ continue;
+ }
+ cap->running_task = task;
+ popReturningTask(cap);
+ RELEASE_LOCK(&cap->lock);
+ break;
+ }
+ RELEASE_LOCK(&cap->lock);
+ }
+
+ }
+
+ ASSERT_FULL_CAPABILITY_INVARIANTS(cap,task);
+
+ IF_DEBUG(scheduler,
+ sched_belch("returning; got capability %d", cap->no));
+
+ *pCap = cap;
+#endif
+}
+
+#if defined(THREADED_RTS)
+/* ----------------------------------------------------------------------------
+ * yieldCapability
+ * ------------------------------------------------------------------------- */
+
+void
+yieldCapability (Capability** pCap, Task *task)
+{
+ Capability *cap = *pCap;
+
+ // The fast path has no locking, if we don't enter this while loop
+
+ while ( cap->returning_tasks_hd != NULL || !anyWorkForMe(cap,task) ) {
+ IF_DEBUG(scheduler, sched_belch("giving up capability %d", cap->no));
+
+ // We must now release the capability and wait to be woken up
+ // again.
+ task->wakeup = rtsFalse;
+ releaseCapabilityAndQueueWorker(cap);
+
+ for (;;) {
+ ACQUIRE_LOCK(&task->lock);
+ // task->lock held, cap->lock not held
+ if (!task->wakeup) waitCondition(&task->cond, &task->lock);
+ cap = task->cap;
+ task->wakeup = rtsFalse;
+ RELEASE_LOCK(&task->lock);
+
+ IF_DEBUG(scheduler, sched_belch("woken up on capability %d", cap->no));
+ ACQUIRE_LOCK(&cap->lock);
+ if (cap->running_task != NULL) {
+ IF_DEBUG(scheduler, sched_belch("capability %d is owned by another task", cap->no));
+ RELEASE_LOCK(&cap->lock);
+ continue;
+ }
+
+ if (task->tso == NULL) {
+ ASSERT(cap->spare_workers != NULL);
+ // if we're not at the front of the queue, release it
+ // again. This is unlikely to happen.
+ if (cap->spare_workers != task) {
+ giveCapabilityToTask(cap,cap->spare_workers);
+ RELEASE_LOCK(&cap->lock);
+ continue;
+ }
+ cap->spare_workers = task->next;
+ task->next = NULL;
+ }
+ cap->running_task = task;
+ RELEASE_LOCK(&cap->lock);
+ break;
+ }
+
+ IF_DEBUG(scheduler, sched_belch("got capability %d", cap->no));
+ ASSERT(cap->running_task == task);
+ }
+
+ *pCap = cap;
+
+ ASSERT_FULL_CAPABILITY_INVARIANTS(cap,task);
+
+ return;
+}
+
+/* ----------------------------------------------------------------------------
+ * Wake up a thread on a Capability.
+ *
+ * This is used when the current Task is running on a Capability and
+ * wishes to wake up a thread on a different Capability.
+ * ------------------------------------------------------------------------- */
+
+void
+wakeupThreadOnCapability (Capability *cap, StgTSO *tso)
+{
+ ASSERT(tso->cap == cap);
+ ASSERT(tso->bound ? tso->bound->cap == cap : 1);
+
+ ACQUIRE_LOCK(&cap->lock);
+ if (cap->running_task == NULL) {
+ // nobody is running this Capability, we can add our thread
+ // directly onto the run queue and start up a Task to run it.
+ appendToRunQueue(cap,tso);
+
+ // start it up
+ cap->running_task = myTask(); // precond for releaseCapability_()
+ releaseCapability_(cap);
+ } else {
+ appendToWakeupQueue(cap,tso);
+ // someone is running on this Capability, so it cannot be
+ // freed without first checking the wakeup queue (see
+ // releaseCapability_).
+ }
+ RELEASE_LOCK(&cap->lock);
+}
+
+/* ----------------------------------------------------------------------------
+ * prodCapabilities
+ *
+ * Used to indicate that the interrupted flag is now set, or some
+ * other global condition that might require waking up a Task on each
+ * Capability.
+ * ------------------------------------------------------------------------- */
+
+static void
+prodCapabilities(rtsBool all)
+{
+ nat i;
+ Capability *cap;
+ Task *task;
+
+ for (i=0; i < n_capabilities; i++) {
+ cap = &capabilities[i];
+ ACQUIRE_LOCK(&cap->lock);
+ if (!cap->running_task) {
+ if (cap->spare_workers) {
+ task = cap->spare_workers;
+ ASSERT(!task->stopped);
+ giveCapabilityToTask(cap,task);
+ if (!all) {
+ RELEASE_LOCK(&cap->lock);
+ return;
+ }
+ }
+ }
+ RELEASE_LOCK(&cap->lock);
+ }
+ return;
+}
+
+void
+prodAllCapabilities (void)
+{
+ prodCapabilities(rtsTrue);
+}
+
+/* ----------------------------------------------------------------------------
+ * prodOneCapability
+ *
+ * Like prodAllCapabilities, but we only require a single Task to wake
+ * up in order to service some global event, such as checking for
+ * deadlock after some idle time has passed.
+ * ------------------------------------------------------------------------- */
+
+void
+prodOneCapability (void)
+{
+ prodCapabilities(rtsFalse);
+}
+
+/* ----------------------------------------------------------------------------
+ * shutdownCapability
+ *
+ * At shutdown time, we want to let everything exit as cleanly as
+ * possible. For each capability, we let its run queue drain, and
+ * allow the workers to stop.
+ *
+ * This function should be called when interrupted and
+ * shutting_down_scheduler = rtsTrue, thus any worker that wakes up
+ * will exit the scheduler and call taskStop(), and any bound thread
+ * that wakes up will return to its caller. Runnable threads are
+ * killed.
+ *
+ * ------------------------------------------------------------------------- */
+
+void
+shutdownCapability (Capability *cap, Task *task)
+{
+ nat i;
+
+ ASSERT(sched_state == SCHED_SHUTTING_DOWN);
+
+ task->cap = cap;
+
+ for (i = 0; i < 50; i++) {
+ IF_DEBUG(scheduler, sched_belch("shutting down capability %d, attempt %d", cap->no, i));
+ ACQUIRE_LOCK(&cap->lock);
+ if (cap->running_task) {
+ RELEASE_LOCK(&cap->lock);
+ IF_DEBUG(scheduler, sched_belch("not owner, yielding"));
+ yieldThread();
+ continue;
+ }
+ cap->running_task = task;
+ if (!emptyRunQueue(cap) || cap->spare_workers) {
+ IF_DEBUG(scheduler, sched_belch("runnable threads or workers still alive, yielding"));
+ releaseCapability_(cap); // this will wake up a worker
+ RELEASE_LOCK(&cap->lock);
+ yieldThread();
+ continue;
+ }
+ IF_DEBUG(scheduler, sched_belch("capability %d is stopped.", cap->no));
+ RELEASE_LOCK(&cap->lock);
+ break;
+ }
+ // we now have the Capability, its run queue and spare workers
+ // list are both empty.
+}
+
+/* ----------------------------------------------------------------------------
+ * tryGrabCapability
+ *
+ * Attempt to gain control of a Capability if it is free.
+ *
+ * ------------------------------------------------------------------------- */
+
+rtsBool
+tryGrabCapability (Capability *cap, Task *task)
+{
+ if (cap->running_task != NULL) return rtsFalse;
+ ACQUIRE_LOCK(&cap->lock);
+ if (cap->running_task != NULL) {
+ RELEASE_LOCK(&cap->lock);
+ return rtsFalse;
+ }
+ task->cap = cap;
+ cap->running_task = task;
+ RELEASE_LOCK(&cap->lock);
+ return rtsTrue;
+}
+
+
+#endif /* THREADED_RTS */
+