[Git][ghc/ghc][wip/T22859] Implement user-defined allocation limit handlers
Teo Camarasu (@teo)
gitlab at gitlab.haskell.org
Tue Mar 19 10:26:37 UTC 2024
Teo Camarasu pushed to branch wip/T22859 at Glasgow Haskell Compiler / GHC
Commits:
52b27a8c by Teo Camarasu at 2024-03-19T10:26:23+00:00
Implement user-defined allocation limit handlers
Resolves #22859
- - - - -
16 changed files:
- compiler/GHC/Builtin/primops.txt.pp
- compiler/GHC/StgToCmm/Prim.hs
- compiler/GHC/StgToJS/Prim.hs
- libraries/ghc-experimental/ghc-experimental.cabal
- + libraries/ghc-experimental/src/System/Mem/Experimental.hs
- libraries/ghc-internal/ghc-internal.cabal
- + libraries/ghc-internal/src/GHC/Internal/AllocationLimitHandler.hs
- rts/Prelude.h
- rts/PrimOps.cmm
- 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/ghc-experimental-exports.stdout
Changes:
=====================================
compiler/GHC/Builtin/primops.txt.pp
=====================================
@@ -3952,6 +3952,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
=====================================
@@ -1746,6 +1746,7 @@ emitPrimOp cfg primop =
TraceEventBinaryOp -> alwaysExternal
TraceMarkerOp -> alwaysExternal
SetThreadAllocationCounter -> alwaysExternal
+ SetOtherThreadAllocationCounter -> alwaysExternal
KeepAliveOp -> alwaysExternal
CastWord32ToFloatOp -> alwaysExternal
CastWord64ToDoubleOp -> alwaysExternal
=====================================
compiler/GHC/StgToJS/Prim.hs
=====================================
@@ -1167,6 +1167,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
=====================================
@@ -27,6 +27,7 @@ library
Data.Tuple.Experimental
Data.Sum.Experimental
Prelude.Experimental
+ System.Mem.Experimental
if arch(wasm32)
exposed-modules: GHC.Wasm.Prim
other-extensions:
=====================================
libraries/ghc-experimental/src/System/Mem/Experimental.hs
=====================================
@@ -0,0 +1,10 @@
+module System.Mem.Experimental
+ ( setGlobalAllocationLimitHandler
+ , AllocationLimitKillBehaviour(..)
+ , getAllocationCounterFor
+ , setAllocationCounterFor
+ , enableAllocationLimitFor
+ , disableAllocationLimitFor
+ )
+ where
+import GHC.Internal.AllocationLimitHandler
=====================================
libraries/ghc-internal/ghc-internal.cabal
=====================================
@@ -91,6 +91,7 @@ Library
ghc-bignum >= 1.0 && < 2.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
=====================================
@@ -0,0 +1,106 @@
+{-# 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 at .
+--
+-- 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 at .
+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# #)
+
+getAllocationCounterFor :: ThreadId -> IO Int64
+getAllocationCounterFor (ThreadId t#) = IO $ \s ->
+ case getOtherThreadAllocationCounter# t# s of (# s', i# #) -> (# s', I64# i# #)
+
+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 at .
+enableAllocationLimitFor :: ThreadId -> IO ()
+enableAllocationLimitFor (ThreadId t) = do
+ rts_enableThreadAllocationLimit t
+
+-- | Disable allocation limit processing the thread @t at .
+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(ghczmprim_GHCziCString_unpackCStringzh_info);
PRELUDE_INFO(ghczmprim_GHCziTypes_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_GHCziInternalziIOziException_allocationLimitExceeded_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
=====================================
@@ -2923,6 +2923,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
@@ -2935,6 +2940,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/RtsSymbols.c
=====================================
@@ -914,7 +914,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
=====================================
@@ -94,6 +94,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.
@@ -1113,19 +1117,35 @@ 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);
+ // Schedule the handler to be run immediatelly.
+ pushOnRunQueue(cap, hookThread);
+ }
+
}
/* some statistics gathering in the parallel case */
@@ -3327,3 +3347,9 @@ resurrectThreads (StgTSO *threads)
}
}
}
+
+void setAllocLimitKill(bool shouldKill, bool shouldHook)
+{
+ allocLimitKill = shouldKill;
+ allocLimitRunHook = shouldHook;
+}
=====================================
rts/external-symbols.list.in
=====================================
@@ -38,6 +38,7 @@ ghczminternal_GHCziInternalziConcziIO_ioManagerCapabilitiesChanged_closure
ghczminternal_GHCziInternalziConcziSignal_runHandlersPtr_closure
ghczminternal_GHCziInternalziTopHandler_flushStdHandles_closure
ghczminternal_GHCziInternalziTopHandler_runMainIO_closure
+ghczminternal_GHCziInternalziAllocationLimitHandler_runAllocationLimitHandler_closure
ghczmprim_GHCziTypes_Czh_con_info
ghczmprim_GHCziTypes_Izh_con_info
ghczmprim_GHCziTypes_Fzh_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
=====================================
@@ -588,7 +588,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/ghc-experimental-exports.stdout
=====================================
@@ -8638,6 +8638,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.Classes.Eq GHC.Types.Bool -- Defined in ‘GHC.Classes’
View it on GitLab: https://gitlab.haskell.org/ghc/ghc/-/commit/52b27a8c01e9712ab7c5c36319474580062d93fa
--
View it on GitLab: https://gitlab.haskell.org/ghc/ghc/-/commit/52b27a8c01e9712ab7c5c36319474580062d93fa
You're receiving this email because of your account on gitlab.haskell.org.
-------------- next part --------------
An HTML attachment was scrubbed...
URL: <http://mail.haskell.org/pipermail/ghc-commits/attachments/20240319/6ac8473a/attachment-0001.html>
More information about the ghc-commits
mailing list