Per-thread allocation counters and limits

This reverts commit f0fcc41d755876a1b02d1c7c79f57515059f6417.

New changes: now works on 32-bit platforms too.  I added some basic
support for 64-bit subtraction and comparison operations to the x86
NCG.

13 files changed, 142 insertions, 52 deletions

diff --git a/rts/Capability.c b/rts/Capability.c
index 289eeb2c5b..21f63f39d9 100644
--- a/rts/Capability.c
+++ b/rts/Capability.c
@@ -297,6 +297,10 @@ initCapability( Capability *cap, nat i )
     cap->r.rCCCS = NULL;
 #endif
 
+    // 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;
+
     traceCapCreate(cap);
     traceCapsetAssignCap(CAPSET_OSPROCESS_DEFAULT, i);
     traceCapsetAssignCap(CAPSET_CLOCKDOMAIN_DEFAULT, i);
diff --git a/rts/HeapStackCheck.cmm b/rts/HeapStackCheck.cmm
index 0659fed89f..a1fb5d446d 100644
--- a/rts/HeapStackCheck.cmm
+++ b/rts/HeapStackCheck.cmm
@@ -100,7 +100,9 @@ stg_gc_noregs
             CurrentNursery = bdescr_link(CurrentNursery);
             OPEN_NURSERY();
             if (Capability_context_switch(MyCapability()) != 0 :: CInt ||
-                Capability_interrupt(MyCapability())      != 0 :: CInt) {
+                Capability_interrupt(MyCapability())      != 0 :: CInt ||
+                (StgTSO_alloc_limit(CurrentTSO) `lt` (0::I64) &&
+                 (TO_W_(StgTSO_flags(CurrentTSO)) & TSO_ALLOC_LIMIT) != 0)) {
                 ret = ThreadYielding;
                 goto sched;
             } else {
diff --git a/rts/Linker.c b/rts/Linker.c
index 7d029c62ac..2c74a0dd35 100644
--- a/rts/Linker.c
+++ b/rts/Linker.c
@@ -1264,6 +1264,10 @@ typedef struct _RtsSymbolVal {
       SymI_HasProto(rtsSupportsBoundThreads)                            \
       SymI_HasProto(rts_isProfiled)                                     \
       SymI_HasProto(rts_isDynamic)                                      \
+      SymI_HasProto(rts_getThreadAllocationCounter)                     \
+      SymI_HasProto(rts_setThreadAllocationCounter)                     \
+      SymI_HasProto(rts_enableThreadAllocationLimit)                    \
+      SymI_HasProto(rts_disableThreadAllocationLimit)                   \
       SymI_HasProto(setProgArgv)                                        \
       SymI_HasProto(startupHaskell)                                     \
       SymI_HasProto(shutdownHaskell)                                    \
diff --git a/rts/Prelude.h b/rts/Prelude.h
index 0c54148ba2..614c255af5 100644
--- a/rts/Prelude.h
+++ b/rts/Prelude.h
@@ -37,6 +37,7 @@ extern StgClosure ZCMain_main_closure;
 
 PRELUDE_CLOSURE(base_GHCziIOziException_stackOverflow_closure);
 PRELUDE_CLOSURE(base_GHCziIOziException_heapOverflow_closure);
+PRELUDE_CLOSURE(base_GHCziIOziException_allocationLimitExceeded_closure);
 PRELUDE_CLOSURE(base_GHCziIOziException_blockedIndefinitelyOnThrowTo_closure);
 PRELUDE_CLOSURE(base_GHCziIOziException_blockedIndefinitelyOnMVar_closure);
 PRELUDE_CLOSURE(base_GHCziIOziException_blockedIndefinitelyOnSTM_closure);
@@ -101,6 +102,7 @@ PRELUDE_INFO(base_GHCziStable_StablePtr_con_info);
 
 #define stackOverflow_closure     DLL_IMPORT_DATA_REF(base_GHCziIOziException_stackOverflow_closure)
 #define heapOverflow_closure      DLL_IMPORT_DATA_REF(base_GHCziIOziException_heapOverflow_closure)
+#define allocationLimitExceeded_closure DLL_IMPORT_DATA_REF(base_GHCziIOziException_allocationLimitExceeded_closure)
 #define blockedIndefinitelyOnMVar_closure DLL_IMPORT_DATA_REF(base_GHCziIOziException_blockedIndefinitelyOnMVar_closure)
 #define blockedIndefinitelyOnSTM_closure DLL_IMPORT_DATA_REF(base_GHCziIOziException_blockedIndefinitelyOnSTM_closure)
 #define nonTermination_closure    DLL_IMPORT_DATA_REF(base_ControlziExceptionziBase_nonTermination_closure)
diff --git a/rts/RaiseAsync.c b/rts/RaiseAsync.c
index 10585c89fa..3b206ffa7e 100644
--- a/rts/RaiseAsync.c
+++ b/rts/RaiseAsync.c
@@ -89,6 +89,60 @@ suspendComputation (Capability *cap, StgTSO *tso, StgUpdateFrame *stop_here)
 }
 
 /* -----------------------------------------------------------------------------
+   throwToSelf
+
+   Useful for throwing an async exception in a thread from the
+   runtime.  It handles unlocking the throwto message returned by
+   throwTo().
+
+   Note [Throw to self when masked]
+   
+   When a StackOverflow occurs when the thread is masked, we want to
+   defer the exception to when the thread becomes unmasked/hits an
+   interruptible point.  We already have a mechanism for doing this,
+   the blocked_exceptions list, but the use here is a bit unusual,
+   because an exception is normally only added to this list upon
+   an asynchronous 'throwTo' call (with all of the relevant
+   multithreaded nonsense). Morally, a stack overflow should be an
+   asynchronous exception sent by a thread to itself, and it should
+   have the same semantics.  But there are a few key differences:
+   
+   - If you actually tried to send an asynchronous exception to
+     yourself using throwTo, the exception would actually immediately
+     be delivered.  This is because throwTo itself is considered an
+     interruptible point, so the exception is always deliverable. Thus,
+     ordinarily, we never end up with a message to onesself in the
+     blocked_exceptions queue.
+   
+   - In the case of a StackOverflow, we don't actually care about the
+     wakeup semantics; when an exception is delivered, the thread that
+     originally threw the exception should be woken up, since throwTo
+     blocks until the exception is successfully thrown.  Fortunately,
+     it is harmless to wakeup a thread that doesn't actually need waking
+     up, e.g. ourselves.
+   
+   - No synchronization is necessary, because we own the TSO and the
+     capability.  You can observe this by tracing through the execution
+     of throwTo.  We skip synchronizing the message and inter-capability
+     communication.
+   
+   We think this doesn't break any invariants, but do be careful!
+   -------------------------------------------------------------------------- */
+
+void
+throwToSelf (Capability *cap, StgTSO *tso, StgClosure *exception)
+{
+    MessageThrowTo *m;
+
+    m = throwTo(cap, tso, tso, exception);
+
+    if (m != NULL) {
+        // throwTo leaves it locked
+        unlockClosure((StgClosure*)m, &stg_MSG_THROWTO_info);
+    }
+}
+
+/* -----------------------------------------------------------------------------
    throwTo
 
    This function may be used to throw an exception from one thread to
diff --git a/rts/RaiseAsync.h b/rts/RaiseAsync.h
index e2763d0cb8..6bfed8d6ca 100644
--- a/rts/RaiseAsync.h
+++ b/rts/RaiseAsync.h
@@ -28,6 +28,10 @@ void throwToSingleThreaded_ (Capability *cap,
                              StgClosure *exception,
                              rtsBool stop_at_atomically);
 
+void throwToSelf (Capability *cap,
+                  StgTSO *tso,
+                  StgClosure *exception);
+
 void suspendComputation (Capability *cap,
                          StgTSO *tso,
                          StgUpdateFrame *stop_here);
diff --git a/rts/RtsFlags.c b/rts/RtsFlags.c
index 44c05cec3b..82e90e5b0e 100644
--- a/rts/RtsFlags.c
+++ b/rts/RtsFlags.c
@@ -137,6 +137,7 @@ void initRtsFlagsDefaults(void)
 #else
     RtsFlags.GcFlags.heapBase           = 0;   /* means don't care */
 #endif
+    RtsFlags.GcFlags.allocLimitGrace    = (100*1024) / BLOCK_SIZE;
 
 #ifdef DEBUG
     RtsFlags.DebugFlags.scheduler       = rtsFalse;
@@ -403,6 +404,8 @@ usage_text[] = {
 "            +PAPI_EVENT   - collect papi preset event PAPI_EVENT",
 "            #NATIVE_EVENT - collect native event NATIVE_EVENT (in hex)",
 #endif
+"  -xq       The allocation limit given to a thread after it receives",
+"            an AllocationLimitExceeded exception. (default: 100k)",
 "",
 "RTS options may also be specified using the GHCRTS environment variable.",
 "",
@@ -1361,6 +1364,13 @@ error = rtsTrue;
 
                   /* The option prefix '-xx' is reserved for future extension.  KSW 1999-11. */
 
+                case 'q':
+                  OPTION_UNSAFE;
+                  RtsFlags.GcFlags.allocLimitGrace
+                      = decodeSize(rts_argv[arg], 3, BLOCK_SIZE, HS_INT_MAX)
+                          / BLOCK_SIZE;
+                  break;
+
                   default:
                     OPTION_SAFE;
                     errorBelch("unknown RTS option: %s",rts_argv[arg]);
diff --git a/rts/RtsStartup.c b/rts/RtsStartup.c
index 32bed5af8f..b8201e1651 100644
--- a/rts/RtsStartup.c
+++ b/rts/RtsStartup.c
@@ -214,6 +214,7 @@ hs_init_ghc(int *argc, char **argv[], RtsConfig rts_config)
     getStablePtr((StgPtr)blockedIndefinitelyOnMVar_closure);
     getStablePtr((StgPtr)nonTermination_closure);
     getStablePtr((StgPtr)blockedIndefinitelyOnSTM_closure);
+    getStablePtr((StgPtr)allocationLimitExceeded_closure);
     getStablePtr((StgPtr)nestedAtomically_closure);
 
     getStablePtr((StgPtr)runSparks_closure);
diff --git a/rts/Schedule.c b/rts/Schedule.c
index b11270832d..c2260f0282 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 */
 }
 
diff --git a/rts/Threads.c b/rts/Threads.c
index 76e844a3f6..90efd9ce4e 100644
--- a/rts/Threads.c
+++ b/rts/Threads.c
@@ -110,6 +110,8 @@ createThread(Capability *cap, W_ size)
     tso->stackobj       = stack;
     tso->tot_stack_size = stack->stack_size;
 
+    tso->alloc_limit = 0;
+
     tso->trec = NO_TREC;
 
 #ifdef PROFILING
@@ -164,6 +166,31 @@ rts_getThreadId(StgPtr tso)
   return ((StgTSO *)tso)->id;
 }
 
+/* ---------------------------------------------------------------------------
+ * Getting & setting the thread allocation limit
+ * ------------------------------------------------------------------------ */
+HsInt64 rts_getThreadAllocationCounter(StgPtr tso)
+{
+    // NB. doesn't take into account allocation in the current nursery
+    // block, so it might be off by up to 4k.
+    return ((StgTSO *)tso)->alloc_limit;
+}
+
+void rts_setThreadAllocationCounter(StgPtr tso, HsInt64 i)
+{
+    ((StgTSO *)tso)->alloc_limit = i;
+}
+
+void rts_enableThreadAllocationLimit(StgPtr tso)
+{
+    ((StgTSO *)tso)->flags |= TSO_ALLOC_LIMIT;
+}
+
+void rts_disableThreadAllocationLimit(StgPtr tso)
+{
+    ((StgTSO *)tso)->flags &= ~TSO_ALLOC_LIMIT;
+}
+
 /* -----------------------------------------------------------------------------
    Remove a thread from a queue.
    Fails fatally if the TSO is not on the queue.
@@ -524,21 +551,8 @@ threadStackOverflow (Capability *cap, StgTSO *tso)
                                  stg_min(tso->stackobj->stack + tso->stackobj->stack_size,
                                          tso->stackobj->sp+64)));
 
-        if (tso->flags & TSO_BLOCKEX) {
-            // NB. StackOverflow exceptions must be deferred if the thread is
-            // inside Control.Exception.mask.  See bug #767 and bug #8303.
-            // This implementation is a minor hack, see Note [Throw to self when masked]
-            MessageThrowTo *msg = (MessageThrowTo*)allocate(cap, sizeofW(MessageThrowTo));
-            SET_HDR(msg, &stg_MSG_THROWTO_info, CCS_SYSTEM);
-            msg->source = tso;
-            msg->target = tso;
-            msg->exception = (StgClosure *)stackOverflow_closure;
-            blockedThrowTo(cap, tso, msg);
-        } else {
-            // Send this thread the StackOverflow exception
-            throwToSingleThreaded(cap, tso, (StgClosure *)stackOverflow_closure);
-            return;
-        }
+        // Note [Throw to self when masked], also #767 and #8303.
+        throwToSelf(cap, tso, (StgClosure *)stackOverflow_closure);
     }
 
 
@@ -669,39 +683,6 @@ threadStackOverflow (Capability *cap, StgTSO *tso)
     // IF_DEBUG(scheduler,printTSO(new_tso));
 }
 
-/* Note [Throw to self when masked]
- *
- * When a StackOverflow occurs when the thread is masked, we want to
- * defer the exception to when the thread becomes unmasked/hits an
- * interruptible point.  We already have a mechanism for doing this,
- * the blocked_exceptions list, but the use here is a bit unusual,
- * because an exception is normally only added to this list upon
- * an asynchronous 'throwTo' call (with all of the relevant
- * multithreaded nonsense). Morally, a stack overflow should be an
- * asynchronous exception sent by a thread to itself, and it should
- * have the same semantics.  But there are a few key differences:
- *
- * - If you actually tried to send an asynchronous exception to
- *   yourself using throwTo, the exception would actually immediately
- *   be delivered.  This is because throwTo itself is considered an
- *   interruptible point, so the exception is always deliverable. Thus,
- *   ordinarily, we never end up with a message to onesself in the
- *   blocked_exceptions queue.
- *
- * - In the case of a StackOverflow, we don't actually care about the
- *   wakeup semantics; when an exception is delivered, the thread that
- *   originally threw the exception should be woken up, since throwTo
- *   blocks until the exception is successfully thrown.  Fortunately,
- *   it is harmless to wakeup a thread that doesn't actually need waking
- *   up, e.g. ourselves.
- *
- * - No synchronization is necessary, because we own the TSO and the
- *   capability.  You can observe this by tracing through the execution
- *   of throwTo.  We skip synchronizing the message and inter-capability
- *   communication.
- *
- * We think this doesn't break any invariants, but do be careful!
- */
 
 
 /* ---------------------------------------------------------------------------
diff --git a/rts/package.conf.in b/rts/package.conf.in
index 82d2870cde..ce44a09651 100644
--- a/rts/package.conf.in
+++ b/rts/package.conf.in
@@ -99,6 +99,7 @@ ld-options:
          , "-Wl,-u,_base_ControlziExceptionziBase_nonTermination_closure"
          , "-Wl,-u,_base_GHCziIOziException_blockedIndefinitelyOnMVar_closure"
          , "-Wl,-u,_base_GHCziIOziException_blockedIndefinitelyOnSTM_closure"
+         , "-Wl,-u,_base_GHCziIOziException_allocationLimitExceeded_closure"
          , "-Wl,-u,_base_ControlziExceptionziBase_nestedAtomically_closure"
          , "-Wl,-u,_base_GHCziEventziThread_blockedOnBadFD_closure"
          , "-Wl,-u,_base_GHCziWeak_runFinalizzerBatch_closure"
@@ -140,6 +141,7 @@ ld-options:
          , "-Wl,-u,base_ControlziExceptionziBase_nonTermination_closure"
          , "-Wl,-u,base_GHCziIOziException_blockedIndefinitelyOnMVar_closure"
          , "-Wl,-u,base_GHCziIOziException_blockedIndefinitelyOnSTM_closure"
+         , "-Wl,-u,base_GHCziIOziException_allocationLimitExceeded_closure"
          , "-Wl,-u,base_ControlziExceptionziBase_nestedAtomically_closure"
          , "-Wl,-u,base_GHCziEventziThread_blockedOnBadFD_closure"
          , "-Wl,-u,base_GHCziWeak_runFinalizzerBatch_closure"
diff --git a/rts/sm/Storage.c b/rts/sm/Storage.c
index 379d9da769..afb171b568 100644
--- a/rts/sm/Storage.c
+++ b/rts/sm/Storage.c
@@ -684,7 +684,10 @@ StgPtr allocate (Capability *cap, W_ n)
 
     TICK_ALLOC_HEAP_NOCTR(WDS(n));
     CCS_ALLOC(cap->r.rCCCS,n);
-    
+    if (cap->r.rCurrentTSO != NULL) {
+        cap->r.rCurrentTSO->alloc_limit -= n*sizeof(W_);
+    }
+
     if (n >= LARGE_OBJECT_THRESHOLD/sizeof(W_)) {
         // The largest number of words such that
         // the computation of req_blocks will not overflow.
@@ -829,6 +832,9 @@ allocatePinned (Capability *cap, W_ n)
 
     TICK_ALLOC_HEAP_NOCTR(WDS(n));
     CCS_ALLOC(cap->r.rCCCS,n);
+    if (cap->r.rCurrentTSO != NULL) {
+        cap->r.rCurrentTSO->alloc_limit -= n*sizeof(W_);
+    }
 
     bd = cap->pinned_object_block;
     
diff --git a/rts/win32/libHSbase.def b/rts/win32/libHSbase.def
index 8140528c70..2091e85c9c 100644
--- a/rts/win32/libHSbase.def
+++ b/rts/win32/libHSbase.def
@@ -32,11 +32,12 @@ EXPORTS
 
 	base_GHCziTopHandler_flushStdHandles_closure
 
-	base_GHCziWeak_runFinalizzerBatch_closure
+        base_GHCziWeak_runFinalizzerBatch_closure
 	base_GHCziPack_unpackCString_closure
 	base_GHCziIOziException_blockedIndefinitelyOnMVar_closure
 	base_GHCziIOziException_blockedIndefinitelyOnSTM_closure
-	base_GHCziIOziException_stackOverflow_closure
+        base_GHCziIOziException_allocationLimitExceeded_closure
+        base_GHCziIOziException_stackOverflow_closure
 
 	base_ControlziExceptionziBase_nonTermination_closure
 	base_ControlziExceptionziBase_nestedAtomically_closure