1 files changed, 22 insertions, 0 deletions
diff --git a/rts/Schedule.c b/rts/Schedule.c
index b11270832d..e9b0289599 100644
--- a/rts/Schedule.c
+++ b/rts/Schedule.c
@@ -481,6 +481,10 @@ run_thread:
     // happened.  So find the new location:
     t = cap->r.rCurrentTSO;
 
+    // cap->r.rCurrentTSO is charged for calls to allocate(), so we
+    // don't want it set when not running a Haskell thread.
+    cap->r.rCurrentTSO = NULL;
+
     // And save the current errno in this thread.
     // XXX: possibly bogus for SMP because this thread might already
     // be running again, see code below.
@@ -1078,6 +1082,21 @@ schedulePostRunThread (Capability *cap, StgTSO *t)
         }
     }
 
+    //
+    // If the current thread's allocation limit has run out, send it
+    // the AllocationLimitExceeded exception.
+
+    if (t->alloc_limit < 0 && (t->flags & TSO_ALLOC_LIMIT)) {
+        // Use a throwToSelf rather than a throwToSingleThreaded, because
+        // it correctly handles the case where the thread is currently
+        // inside mask.  Also the thread might be blocked (e.g. on an
+        // MVar), and throwToSingleThreaded doesn't unblock it
+        // correctly in that case.
+        throwToSelf(cap, t, allocationLimitExceeded_closure);
+        t->alloc_limit = (StgInt64)RtsFlags.GcFlags.allocLimitGrace
+            * BLOCK_SIZE;
+    }
+
   /* some statistics gathering in the parallel case */
 }
 
@@ -2214,6 +2233,9 @@ suspendThread (StgRegTable *reg, rtsBool interruptible)
   task->incall->suspended_tso = tso;
   task->incall->suspended_cap = cap;
 
+  // Otherwise allocate() will write to invalid memory.
+  cap->r.rCurrentTSO = NULL
+
   ACQUIRE_LOCK(&cap->lock);
 
   suspendTask(cap,task);