GitLab

Teo Camarasu pushed to branch wip/T22859 at Glasgow Haskell Compiler / GHC

Commits:

d3d7e58c

by Teo Camarasu at 2025-06-10T16:09:47+01:00

Implement user-defined allocation limit handlers

Resolves #22859

24 changed files:

compiler/GHC/Builtin/primops.txt.pp
compiler/GHC/StgToCmm/Prim.hs
compiler/GHC/StgToJS/Prim.hs
libraries/ghc-experimental/ghc-experimental.cabal.in
+ libraries/ghc-experimental/src/System/Mem/Experimental.hs
libraries/ghc-internal/ghc-internal.cabal.in
+ libraries/ghc-internal/src/GHC/Internal/AllocationLimitHandler.hs
rts/Prelude.h
rts/PrimOps.cmm
rts/RtsStartup.c
rts/RtsSymbols.c
rts/Schedule.c
rts/external-symbols.list.in
rts/include/rts/storage/GC.h
rts/include/rts/storage/TSO.h
rts/include/stg/MiscClosures.h
testsuite/tests/interface-stability/base-exports.stdout
testsuite/tests/interface-stability/base-exports.stdout-javascript-unknown-ghcjs
testsuite/tests/interface-stability/base-exports.stdout-mingw32
testsuite/tests/interface-stability/base-exports.stdout-ws-32
testsuite/tests/interface-stability/ghc-experimental-exports.stdout
+ testsuite/tests/rts/T22859.hs
+ testsuite/tests/rts/T22859.stderr
testsuite/tests/rts/all.T

Changes:

compiler/GHC/Builtin/primops.txt.pp

@@ -4065,6 +4065,15 @@ primop  SetThreadAllocationCounter "setThreadAllocationCounter#" GenPrimOp
     effect = ReadWriteEffect
     out_of_line      = True
 +primop  SetOtherThreadAllocationCounter "setOtherThreadAllocationCounter#" GenPrimOp
 +   Int64# -> ThreadId# -> State# RealWorld -> State# RealWorld
 +   { Sets the allocation counter for the another thread to the given value.
 +     This doesn't take allocations into the current nursery chunk into account.
 +     Therefore it is only accurate if the other thread is not currently running. }
 +   with
 +   effect = ReadWriteEffect
 +   out_of_line      = True
++
  primtype StackSnapshot#
     { Haskell representation of a @StgStack*@ that was created (cloned)
       with a function in "GHC.Stack.CloneStack". Please check the

compiler/GHC/StgToCmm/Prim.hs

@@ -1775,6 +1775,7 @@ emitPrimOp cfg primop =
    TraceEventBinaryOp -> alwaysExternal
    TraceMarkerOp -> alwaysExternal
    SetThreadAllocationCounter -> alwaysExternal
 +  SetOtherThreadAllocationCounter -> alwaysExternal
    KeepAliveOp -> alwaysExternal
   where

compiler/GHC/StgToJS/Prim.hs

@@ -1173,6 +1173,7 @@ genPrim prof bound ty op = case op of
    WhereFromOp                       -> unhandledPrimop op -- should be easily implementable with o.f.n
    SetThreadAllocationCounter        -> unhandledPrimop op
 +  SetOtherThreadAllocationCounter   -> unhandledPrimop op
  ------------------------------- Vector -----------------------------------------
  -- For now, vectors are unsupported on the JS backend. Simply put, they do not

libraries/ghc-experimental/ghc-experimental.cabal.in

@@ -38,6 +38,7 @@ library
        GHC.RTS.Flags.Experimental
        GHC.Stats.Experimental
        Prelude.Experimental
 +      System.Mem.Experimental
      if arch(wasm32)
          exposed-modules:  GHC.Wasm.Prim
      other-extensions:

libraries/ghc-experimental/src/System/Mem/Experimental.hs

 +module System.Mem.Experimental
 +  ( setGlobalAllocationLimitHandler
 +  , AllocationLimitKillBehaviour(..)
 +  , getAllocationCounterFor
 +  , setAllocationCounterFor
 +  , enableAllocationLimitFor
 +  , disableAllocationLimitFor
 +  )
 +  where
 +import GHC.Internal.AllocationLimitHandler

libraries/ghc-internal/ghc-internal.cabal.in

@@ -122,6 +122,7 @@ Library
          rts == 1.0.*
      exposed-modules:
 +        GHC.Internal.AllocationLimitHandler
          GHC.Internal.ClosureTypes
          GHC.Internal.Control.Arrow
          GHC.Internal.Control.Category

libraries/ghc-internal/src/GHC/Internal/AllocationLimitHandler.hs

 +{-# LANGUAGE MagicHash #-}
 +{-# LANGUAGE UnboxedTuples #-}
 +{-# LANGUAGE UnliftedFFITypes #-}
 +{-# LANGUAGE GHCForeignImportPrim #-}
 +{-# OPTIONS_HADDOCK not-home #-}
 +module GHC.Internal.AllocationLimitHandler
 +  ( runAllocationLimitHandler
 +  , setGlobalAllocationLimitHandler
 +  , AllocationLimitKillBehaviour(..)
 +  , getAllocationCounterFor
 +  , setAllocationCounterFor
 +  , enableAllocationLimitFor
 +  , disableAllocationLimitFor
 +  )
 +  where
 +import GHC.Internal.Base
 +import GHC.Internal.Conc.Sync (ThreadId(..))
 +import GHC.Internal.Data.IORef (IORef, readIORef, writeIORef, newIORef)
 +import GHC.Internal.Foreign.C.Types
 +import GHC.Internal.IO (unsafePerformIO)
 +import GHC.Internal.Int (Int64(..))
++
++
 +{-# NOINLINE allocationLimitHandler #-}
 +allocationLimitHandler :: IORef (ThreadId -> IO ())
 +allocationLimitHandler = unsafePerformIO (newIORef defaultHandler)
++
 +defaultHandler :: ThreadId -> IO ()
 +defaultHandler _ = pure ()
++
 +foreign import ccall "setAllocLimitKill" setAllocLimitKill :: CBool -> CBool -> IO ()
++
 +runAllocationLimitHandler :: ThreadId# -> IO ()
 +runAllocationLimitHandler tid = do
 +  hook <- getAllocationLimitHandler
 +  hook $ ThreadId tid
++
 +getAllocationLimitHandler :: IO (ThreadId -> IO ())
 +getAllocationLimitHandler = readIORef allocationLimitHandler
++
 +data AllocationLimitKillBehaviour =
 +  KillOnAllocationLimit
 +  -- ^ Throw a @AllocationLimitExceeded@ async exception to the thread when the
 +  -- allocation limit is exceeded.
 +  | DontKillOnAllocationLimit
 +  -- ^ Do not throw an exception when the allocation limit is exceeded.
++
 +-- | Define the behaviour for handling allocation limits.
 +-- By default we throw a @AllocationLimitExceeded@ async exception to the thread.
 +-- This can be controlled using @AllocationLimitKillBehaviour@.
 +--
 +-- We can also run a user-specified handler, which can be done in addition to
 +-- or in place of the exception.
 +-- This allows for instance logging on the allocation limit being exceeded,
 +-- or dynamically determining whether to terminate the thread.
 +-- The handler is not guaranteed to run before the thread is terminated or restarted.
 +--
 +-- Note: that if you don't terminate the thread, then the allocation limit gets
 +-- removed.
 +-- If you wish to keep the allocation limit you will have to reset it using
 +-- @setAllocationCounter@ and @enableAllocationLimit@.
 +setGlobalAllocationLimitHandler :: AllocationLimitKillBehaviour -> Maybe (ThreadId -> IO ()) -> IO ()
 +setGlobalAllocationLimitHandler killBehaviour mHandler = do
 +  shouldRunHandler <- case mHandler of
 +    Just hook -> do
 +      writeIORef allocationLimitHandler hook
 +      pure 1
 +    Nothing -> do
 +      writeIORef allocationLimitHandler defaultHandler
 +      pure 0
 +  let shouldKill =
 +        case killBehaviour of
 +          KillOnAllocationLimit -> 1
 +          DontKillOnAllocationLimit -> 0
 +  setAllocLimitKill shouldKill shouldRunHandler
++
 +-- | Retrieves the allocation counter for the another thread.
 +foreign import prim "stg_getOtherThreadAllocationCounterzh" getOtherThreadAllocationCounter#
 +  :: ThreadId#
 +  -> State# RealWorld
 +  -> (# State# RealWorld, Int64# #)
++
 +-- | Get the allocation counter for a different thread.
 +-- Note this doesn't take the current nursery chunk into account.
 +-- If the thread is running then it may underestimate allocations by the size of a nursery thread.
 +getAllocationCounterFor :: ThreadId -> IO Int64
 +getAllocationCounterFor (ThreadId t#) = IO $ \s ->
 +  case getOtherThreadAllocationCounter# t# s of (# s', i# #)  -> (# s', I64# i# #)
++
 +-- | Set the allocation counter for a different thread.
 +-- This can be combined with 'enableAllocationLimitFor' to enable allocation limits for another thread.
 +-- Note this doesn't take the current nursery chunk into account.
 +-- If the thread is running then it may overestimate allocations by the size of a nursery thread,
 +-- and trigger the limit sooner than expected.
 +setAllocationCounterFor :: Int64 -> ThreadId -> IO ()
 +setAllocationCounterFor (I64# i#) (ThreadId t#) = IO $ \s ->
 +  case setOtherThreadAllocationCounter# i# t# s of s' -> (# s', () #)
++
++
 +-- | Enable allocation limit processing the thread @t@.
 +enableAllocationLimitFor :: ThreadId -> IO ()
 +enableAllocationLimitFor (ThreadId t) = do
 +  rts_enableThreadAllocationLimit t
++
 +-- | Disable allocation limit processing the thread @t@.
 +disableAllocationLimitFor :: ThreadId -> IO ()
 +disableAllocationLimitFor (ThreadId t) = do
 +  rts_disableThreadAllocationLimit t
++
 +foreign import ccall unsafe "rts_enableThreadAllocationLimit"
 +  rts_enableThreadAllocationLimit :: ThreadId# -> IO ()
++
 +foreign import ccall unsafe "rts_disableThreadAllocationLimit"
 +  rts_disableThreadAllocationLimit :: ThreadId# -> IO ()

rts/Prelude.h

@@ -67,6 +67,7 @@ PRELUDE_CLOSURE(ghczminternal_GHCziInternalziEventziWindows_processRemoteComplet
  PRELUDE_CLOSURE(ghczminternal_GHCziInternalziTopHandler_flushStdHandles_closure);
  PRELUDE_CLOSURE(ghczminternal_GHCziInternalziTopHandler_runMainIO_closure);
 +PRELUDE_CLOSURE(ghczminternal_GHCziInternalziAllocationLimitHandler_runAllocationLimitHandler_closure);
  PRELUDE_INFO(ghczminternal_GHCziInternalziCString_unpackCStringzh_info);
  PRELUDE_INFO(ghczminternal_GHCziInternalziTypes_Czh_con_info);
@@ -102,6 +103,7 @@ PRELUDE_INFO(ghczminternal_GHCziInternalziStable_StablePtr_con_info);
  #if defined(mingw32_HOST_OS)
  #define processRemoteCompletion_closure DLL_IMPORT_DATA_REF(ghczminternal_GHCziInternalziEventziWindows_processRemoteCompletion_closure)
  #endif
 +#define runAllocationLimitHandler_closure DLL_IMPORT_DATA_REF(ghczminternal_GHCziInternalziAllocationLimitHandler_runAllocationLimitHandler_closure)
  #define flushStdHandles_closure   DLL_IMPORT_DATA_REF(ghczminternal_GHCziInternalziTopHandler_flushStdHandles_closure)
  #define runMainIO_closure   DLL_IMPORT_DATA_REF(ghczminternal_GHCziInternalziTopHandler_runMainIO_closure)

rts/PrimOps.cmm

@@ -2889,6 +2889,11 @@ stg_getThreadAllocationCounterzh ()
      return (StgTSO_alloc_limit(CurrentTSO) - TO_I64(offset));
+ }
 +stg_getOtherThreadAllocationCounterzh ( gcptr t )
 +{
 +    return (StgTSO_alloc_limit(t));
 +}
++
  stg_setThreadAllocationCounterzh ( I64 counter )
+ {
      // Allocation in the current block will be subtracted by
@@ -2901,6 +2906,12 @@ stg_setThreadAllocationCounterzh ( I64 counter )
      return ();
+ }
 +stg_setOtherThreadAllocationCounterzh ( I64 counter, gcptr t )
 +{
 +    StgTSO_alloc_limit(t) = counter;
 +    return ();
 +}
++
  #define KEEP_ALIVE_FRAME_FIELDS(w_,p_,info_ptr,p1,p2,c)   \
    w_ info_ptr,                                            \

rts/RtsStartup.c

@@ -224,6 +224,7 @@ static void initBuiltinGcRoots(void)
       * GHC.Core.Make.mkExceptionId.
       */
      getStablePtr((StgPtr)absentSumFieldError_closure);
 +    getStablePtr((StgPtr)runAllocationLimitHandler_closure);
+ }
  void

rts/RtsSymbols.c

@@ -916,7 +916,9 @@ extern char **environ;
        SymI_HasDataProto(stg_traceMarkerzh)                                  \
        SymI_HasDataProto(stg_traceBinaryEventzh)                             \
        SymI_HasDataProto(stg_getThreadAllocationCounterzh)                   \
 +      SymI_HasDataProto(stg_getOtherThreadAllocationCounterzh)              \
        SymI_HasDataProto(stg_setThreadAllocationCounterzh)                   \
 +      SymI_HasDataProto(stg_setOtherThreadAllocationCounterzh)              \
        SymI_HasProto(getMonotonicNSec)                                   \
        SymI_HasProto(lockFile)                                           \
        SymI_HasProto(unlockFile)                                         \

rts/Schedule.c

@@ -41,6 +41,7 @@
  #include "Threads.h"
  #include "Timer.h"
  #include "ThreadPaused.h"
 +#include "ThreadLabels.h"
  #include "Messages.h"
  #include "StablePtr.h"
  #include "StableName.h"
@@ -94,6 +95,10 @@ StgWord recent_activity = ACTIVITY_YES;
   */
  StgWord sched_state = SCHED_RUNNING;
++
 +bool allocLimitKill = true;
 +bool allocLimitRunHook = false;
++
  /*
   * This mutex protects most of the global scheduler data in
   * the THREADED_RTS runtime.
@@ -1125,19 +1130,36 @@ schedulePostRunThread (Capability *cap, StgTSO *t)
+         }
+     }
 -    //
 -    // If the current thread's allocation limit has run out, send it
 -    // the AllocationLimitExceeded exception.
 +    // Handle the current thread's allocation limit running out,
      if (PK_Int64((W_*)&(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);
 -        ASSIGN_Int64((W_*)&(t->alloc_limit),
 -                     (StgInt64)RtsFlags.GcFlags.allocLimitGrace * BLOCK_SIZE);
 +        if(allocLimitKill) {
 +          // Throw the AllocationLimitExceeded exception.
 +          // 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);
 +          ASSIGN_Int64((W_*)&(t->alloc_limit),
 +                      (StgInt64)RtsFlags.GcFlags.allocLimitGrace * BLOCK_SIZE);
 +        } else {
 +          // If we aren't killing the thread, we must disable the limit
 +          // otherwise we will immediatelly retrigger it.
 +          // User defined handlers should re-enable it if wanted.
 +          t->flags = t->flags & ~TSO_ALLOC_LIMIT;
 +        }
++
 +        if(allocLimitRunHook)
 +        {
 +          // Create a thread to run the allocation limit handler.
 +          StgClosure* c = rts_apply(cap, runAllocationLimitHandler_closure, (StgClosure*)t);
 +          StgTSO* hookThread = createIOThread(cap, RtsFlags.GcFlags.initialStkSize, c);
 +          setThreadLabel(cap, hookThread, "allocation limit handler thread");
 +          // Schedule the handler to be run immediatelly.
 +          pushOnRunQueue(cap, hookThread);
 +        }
++
+     }
    /* some statistics gathering in the parallel case */
@@ -3342,3 +3364,9 @@ resurrectThreads (StgTSO *threads)
+         }
+     }
+ }
++
 +void setAllocLimitKill(bool shouldKill, bool shouldHook)
 +{
 +   allocLimitKill = shouldKill;
 +   allocLimitRunHook = shouldHook;
 +}

rts/external-symbols.list.in

@@ -43,6 +43,7 @@ ghczminternal_GHCziInternalziTypes_Izh_con_info
  ghczminternal_GHCziInternalziTypes_Fzh_con_info
  ghczminternal_GHCziInternalziTypes_Dzh_con_info
  ghczminternal_GHCziInternalziTypes_Wzh_con_info
 +ghczminternal_GHCziInternalziAllocationLimitHandler_runAllocationLimitHandler_closure
  ghczminternal_GHCziInternalziPtr_Ptr_con_info
  ghczminternal_GHCziInternalziPtr_FunPtr_con_info
  ghczminternal_GHCziInternalziInt_I8zh_con_info

rts/include/rts/storage/GC.h

@@ -209,6 +209,10 @@ void flushExec(W_ len, AdjustorExecutable exec_addr);
  // Used by GC checks in external .cmm code:
  extern W_ large_alloc_lim;
 +// Should triggering an allocation limit kill the thread
 +// and should we run a user-defined hook when it is triggered.
 +void setAllocLimitKill(bool, bool);
++
  /* -----------------------------------------------------------------------------
     Performing Garbage Collection
     -------------------------------------------------------------------------- */

rts/include/rts/storage/TSO.h

@@ -157,9 +157,10 @@ typedef struct StgTSO_ {
      /*
       * The allocation limit for this thread, which is updated as the
       * thread allocates.  If the value drops below zero, and
 -     * TSO_ALLOC_LIMIT is set in flags, we raise an exception in the
 -     * thread, and give the thread a little more space to handle the
 -     * exception before we raise the exception again.
 +     * TSO_ALLOC_LIMIT is set in flags, then a handler is triggerd.
 +     * Either we raise an exception in the thread, and give the thread
 +     * a little more space to handle the exception before we raise the
 +     * exception again; or we run a user defined handler.
+      *
       * This is an integer, because we might update it in a place where
       * it isn't convenient to raise the exception, so we want it to

rts/include/stg/MiscClosures.h

@@ -604,7 +604,9 @@ RTS_FUN_DECL(stg_traceEventzh);
  RTS_FUN_DECL(stg_traceBinaryEventzh);
  RTS_FUN_DECL(stg_traceMarkerzh);
  RTS_FUN_DECL(stg_getThreadAllocationCounterzh);
 +RTS_FUN_DECL(stg_getOtherThreadAllocationCounterzh);
  RTS_FUN_DECL(stg_setThreadAllocationCounterzh);
 +RTS_FUN_DECL(stg_setOtherThreadAllocationCounterzh);
  RTS_FUN_DECL(stg_castWord64ToDoublezh);
  RTS_FUN_DECL(stg_castDoubleToWord64zh);

testsuite/tests/interface-stability/base-exports.stdout

@@ -4607,6 +4607,7 @@ module GHC.Base where
    sequence :: forall (m :: * -> *) a. Monad m => [m a] -> m [a]
    setAddrRange# :: Addr# -> Int# -> Int# -> State# RealWorld -> State# RealWorld
    setByteArray# :: forall d. MutableByteArray# d -> Int# -> Int# -> Int# -> State# d -> State# d
 +  setOtherThreadAllocationCounter# :: Int64# -> ThreadId# -> State# RealWorld -> State# RealWorld
    setThreadAllocationCounter# :: Int64# -> State# RealWorld -> State# RealWorld
    shiftL# :: Word# -> Int# -> Word#
    shiftRL# :: Word# -> Int# -> Word#
@@ -6693,6 +6694,7 @@ module GHC.Exts where
    seq# :: forall a s. a -> State# s -> (# State# s, a #)
    setAddrRange# :: Addr# -> Int# -> Int# -> State# RealWorld -> State# RealWorld
    setByteArray# :: forall d. MutableByteArray# d -> Int# -> Int# -> Int# -> State# d -> State# d
 +  setOtherThreadAllocationCounter# :: Int64# -> ThreadId# -> State# RealWorld -> State# RealWorld
    setThreadAllocationCounter# :: Int64# -> State# RealWorld -> State# RealWorld
    shiftL# :: Word# -> Int# -> Word#
    shiftRL# :: Word# -> Int# -> Word#

testsuite/tests/interface-stability/base-exports.stdout-javascript-unknown-ghcjs

@@ -4607,6 +4607,7 @@ module GHC.Base where
    sequence :: forall (m :: * -> *) a. Monad m => [m a] -> m [a]
    setAddrRange# :: Addr# -> Int# -> Int# -> State# RealWorld -> State# RealWorld
    setByteArray# :: forall d. MutableByteArray# d -> Int# -> Int# -> Int# -> State# d -> State# d
 +  setOtherThreadAllocationCounter# :: Int64# -> ThreadId# -> State# RealWorld -> State# RealWorld
    setThreadAllocationCounter# :: Int64# -> State# RealWorld -> State# RealWorld
    shiftL# :: Word# -> Int# -> Word#
    shiftRL# :: Word# -> Int# -> Word#
@@ -6665,6 +6666,7 @@ module GHC.Exts where
    seq# :: forall a s. a -> State# s -> (# State# s, a #)
    setAddrRange# :: Addr# -> Int# -> Int# -> State# RealWorld -> State# RealWorld
    setByteArray# :: forall d. MutableByteArray# d -> Int# -> Int# -> Int# -> State# d -> State# d
 +  setOtherThreadAllocationCounter# :: Int64# -> ThreadId# -> State# RealWorld -> State# RealWorld
    setThreadAllocationCounter# :: Int64# -> State# RealWorld -> State# RealWorld
    shiftL# :: Word# -> Int# -> Word#
    shiftRL# :: Word# -> Int# -> Word#

testsuite/tests/interface-stability/base-exports.stdout-mingw32

@@ -4610,6 +4610,7 @@ module GHC.Base where
    sequence :: forall (m :: * -> *) a. Monad m => [m a] -> m [a]
    setAddrRange# :: Addr# -> Int# -> Int# -> State# RealWorld -> State# RealWorld
    setByteArray# :: forall d. MutableByteArray# d -> Int# -> Int# -> Int# -> State# d -> State# d
 +  setOtherThreadAllocationCounter# :: Int64# -> ThreadId# -> State# RealWorld -> State# RealWorld
    setThreadAllocationCounter# :: Int64# -> State# RealWorld -> State# RealWorld
    shiftL# :: Word# -> Int# -> Word#
    shiftRL# :: Word# -> Int# -> Word#
@@ -6836,6 +6837,7 @@ module GHC.Exts where
    seq# :: forall a s. a -> State# s -> (# State# s, a #)
    setAddrRange# :: Addr# -> Int# -> Int# -> State# RealWorld -> State# RealWorld
    setByteArray# :: forall d. MutableByteArray# d -> Int# -> Int# -> Int# -> State# d -> State# d
 +  setOtherThreadAllocationCounter# :: Int64# -> ThreadId# -> State# RealWorld -> State# RealWorld
    setThreadAllocationCounter# :: Int64# -> State# RealWorld -> State# RealWorld
    shiftL# :: Word# -> Int# -> Word#
    shiftRL# :: Word# -> Int# -> Word#

testsuite/tests/interface-stability/base-exports.stdout-ws-32

@@ -4607,6 +4607,7 @@ module GHC.Base where
    sequence :: forall (m :: * -> *) a. Monad m => [m a] -> m [a]
    setAddrRange# :: Addr# -> Int# -> Int# -> State# RealWorld -> State# RealWorld
    setByteArray# :: forall d. MutableByteArray# d -> Int# -> Int# -> Int# -> State# d -> State# d
 +  setOtherThreadAllocationCounter# :: Int64# -> ThreadId# -> State# RealWorld -> State# RealWorld
    setThreadAllocationCounter# :: Int64# -> State# RealWorld -> State# RealWorld
    shiftL# :: Word# -> Int# -> Word#
    shiftRL# :: Word# -> Int# -> Word#
@@ -6693,6 +6694,7 @@ module GHC.Exts where
    seq# :: forall a s. a -> State# s -> (# State# s, a #)
    setAddrRange# :: Addr# -> Int# -> Int# -> State# RealWorld -> State# RealWorld
    setByteArray# :: forall d. MutableByteArray# d -> Int# -> Int# -> Int# -> State# d -> State# d
 +  setOtherThreadAllocationCounter# :: Int64# -> ThreadId# -> State# RealWorld -> State# RealWorld
    setThreadAllocationCounter# :: Int64# -> State# RealWorld -> State# RealWorld
    shiftL# :: Word# -> Int# -> Word#
    shiftRL# :: Word# -> Int# -> Word#

testsuite/tests/interface-stability/ghc-experimental-exports.stdout

@@ -10916,6 +10916,16 @@ module Prelude.Experimental where
    data Unit# = ...
    getSolo :: forall a. Solo a -> a
 +module System.Mem.Experimental where
 +  -- Safety: None
 +  type AllocationLimitKillBehaviour :: *
 +  data AllocationLimitKillBehaviour = KillOnAllocationLimit | DontKillOnAllocationLimit
 +  disableAllocationLimitFor :: GHC.Internal.Conc.Sync.ThreadId -> GHC.Types.IO ()
 +  enableAllocationLimitFor :: GHC.Internal.Conc.Sync.ThreadId -> GHC.Types.IO ()
 +  getAllocationCounterFor :: GHC.Internal.Conc.Sync.ThreadId -> GHC.Types.IO GHC.Internal.Int.Int64
 +  setAllocationCounterFor :: GHC.Internal.Int.Int64 -> GHC.Internal.Conc.Sync.ThreadId -> GHC.Types.IO ()
 +  setGlobalAllocationLimitHandler :: AllocationLimitKillBehaviour -> GHC.Internal.Maybe.Maybe (GHC.Internal.Conc.Sync.ThreadId -> GHC.Types.IO ()) -> GHC.Types.IO ()
++
  -- Instances:
  instance GHC.Internal.Base.Alternative GHC.Internal.Types.IO -- Defined in ‘GHC.Internal.Base’

testsuite/tests/rts/T22859.hs

 +{-# LANGUAGE NumericUnderscores #-}
 +{-# LANGUAGE MagicHash #-}
 +{-# LANGUAGE UnboxedTuples #-}
++
 +import Control.Exception
 +import Control.Exception.Backtrace
 +import Control.Concurrent
 +import Control.Concurrent.MVar
 +import System.Mem
 +import System.Mem.Experimental
 +import GHC.IO (IO (..))
 +import GHC.Exts
 +import System.IO
++
 +-- | Just do some work and hPutStrLn to stderr to indicate that we are making progress
 +worker :: IO ()
 +worker = loop [] 2
 +  where
 +    loop !m !n
 +      | n > 30 = hPutStrLn stderr . show $ length m
 +      | otherwise = do
 +          let x = show n
 +          hPutStrLn stderr x
 +          -- just to bulk out the allocations
 +          IO (\s -> case newByteArray# 900000# s of (# s', arr# #) -> (# s', () #))
 +          yield
 +          loop (x:m) (n + 1)
++
 +main :: IO ()
 +main = do
 +  done <- newMVar () -- we use this lock to wait for the worker to finish
 +  started <- newEmptyMVar
 +  let runWorker = do
 +        forkIO . withMVar done $ \_ -> flip onException (hPutStrLn stderr "worker died") $ do
 +          hPutStrLn stderr "worker starting"
 +          putMVar started ()
 +          setAllocationCounter 1_000_000
 +          enableAllocationLimit
 +          worker
 +          hPutStrLn stderr "worker done"
 +        takeMVar started
 +        readMVar done
 +        hFlush stderr
 +        threadDelay 1000
 +  -- default behaviour:
 +  -- kill it after the limit is exceeded
 +  hPutStrLn stderr "default behaviour"
 +  runWorker
 +  hPutStrLn stderr "just log once on the hook being triggered"
 +  setGlobalAllocationLimitHandler DontKillOnAllocationLimit (Just $ \_ -> hPutStrLn stderr "allocation limit triggered 1")
 +  runWorker
 +  hPutStrLn stderr "just log on the hook being triggered"
 +  setGlobalAllocationLimitHandler DontKillOnAllocationLimit . Just $ \tid -> do
 +    hPutStrLn stderr "allocation limit triggered 2"
 +    -- re-enable the hook
 +    setAllocationCounterFor 1_000_000 tid
 +    enableAllocationLimitFor tid
 +  runWorker
 +  hPutStrLn stderr "kill from the hook"
 +  setGlobalAllocationLimitHandler DontKillOnAllocationLimit . Just $ \tId -> throwTo tId AllocationLimitExceeded
 +  runWorker
 +  -- not super helpful, but let's test it anyway
 +  hPutStrLn stderr "do nothing"
 +  setGlobalAllocationLimitHandler DontKillOnAllocationLimit Nothing
 +  runWorker
 +  -- this is possible to handle using an exception handler instead.
 +  hPutStrLn stderr "kill and log"
 +  setGlobalAllocationLimitHandler KillOnAllocationLimit (Just $ \_ -> hPutStrLn stderr "allocation limit triggered 3")
 +  runWorker
 +  threadDelay 1000
 +  hPutStrLn stderr "done"

testsuite/tests/rts/T22859.stderr

 +default behaviour
 +worker starting
 +2
 +3
 +worker died
 +T22859: allocation limit exceeded
 +HasCallStack backtrace:
 +  collectBacktraces, called at libraries/ghc-internal/src/GHC/Internal/Exception.hs:92:13 in ghc-internal:GHC.Internal.Exception
 +  toExceptionWithBacktrace, called at libraries/ghc-internal/src/GHC/Internal/IO.hs:260:11 in ghc-internal:GHC.Internal.IO
 +  throwIO, called at libraries/ghc-internal/src/GHC/Internal/Control/Exception/Base.hs:195:43 in ghc-internal:GHC.Internal.Control.Exception.Base
++
++
 +just log once on the hook being triggered
 +worker starting
 +2
 +3
 +allocation limit triggered 1
 +4
 +5
 +6
 +7
 +8
 +9
 +10
 +11
 +12
 +13
 +14
 +15
 +16
 +17
 +18
 +19
 +20
 +21
 +22
 +23
 +24
 +25
 +26
 +27
 +28
 +29
 +30
 +29
 +worker done
 +just log on the hook being triggered
 +worker starting
 +2
 +3
 +allocation limit triggered 2
 +4
 +5
 +allocation limit triggered 2
 +6
 +7
 +allocation limit triggered 2
 +8
 +9
 +allocation limit triggered 2
 +10
 +11
 +allocation limit triggered 2
 +12
 +13
 +allocation limit triggered 2
 +14
 +15
 +allocation limit triggered 2
 +16
 +17
 +allocation limit triggered 2
 +18
 +19
 +allocation limit triggered 2
 +20
 +21
 +allocation limit triggered 2
 +22
 +23
 +allocation limit triggered 2
 +24
 +25
 +allocation limit triggered 2
 +26
 +27
 +allocation limit triggered 2
 +28
 +29
 +allocation limit triggered 2
 +30
 +29
 +worker done
 +kill from the hook
 +worker starting
 +2
 +3
 +worker died
 +T22859: allocation limit exceeded
 +HasCallStack backtrace:
 +  collectBacktraces, called at libraries/ghc-internal/src/GHC/Internal/Exception.hs:92:13 in ghc-internal:GHC.Internal.Exception
 +  toExceptionWithBacktrace, called at libraries/ghc-internal/src/GHC/Internal/IO.hs:260:11 in ghc-internal:GHC.Internal.IO
 +  throwIO, called at libraries/ghc-internal/src/GHC/Internal/Control/Exception/Base.hs:195:43 in ghc-internal:GHC.Internal.Control.Exception.Base
++
++
 +do nothing
 +worker starting
 +2
 +3
 +4
 +5
 +6
 +7
 +8
 +9
 +10
 +11
 +12
 +13
 +14
 +15
 +16
 +17
 +18
 +19
 +20
 +21
 +22
 +23
 +24
 +25
 +26
 +27
 +28
 +29
 +30
 +29
 +worker done
 +kill and log
 +worker starting
 +2
 +3
 +allocation limit triggered 3
 +worker died
 +T22859: allocation limit exceeded
 +HasCallStack backtrace:
 +  collectBacktraces, called at libraries/ghc-internal/src/GHC/Internal/Exception.hs:92:13 in ghc-internal:GHC.Internal.Exception
 +  toExceptionWithBacktrace, called at libraries/ghc-internal/src/GHC/Internal/IO.hs:260:11 in ghc-internal:GHC.Internal.IO
 +  throwIO, called at libraries/ghc-internal/src/GHC/Internal/Control/Exception/Base.hs:195:43 in ghc-internal:GHC.Internal.Control.Exception.Base
++
++
 +done

testsuite/tests/rts/all.T

@@ -643,3 +643,4 @@ test('T25280', [unless(opsys('linux'),skip),req_process,js_skip], compile_and_ru
  test('T25560', [req_c_rts, ignore_stderr], compile_and_run, [''])
  test('TestProddableBlockSet', [req_c_rts], multimod_compile_and_run, ['TestProddableBlockSet.c', '-no-hs-main'])
 +test('T22859', [js_skip], compile_and_run, ['-with-rtsopts -A8K'])