[Git][ghc/ghc][wip/T22859] Implement user-defined allocation limit handlers

Teo Camarasu (@teo) gitlab at gitlab.haskell.org
Fri Mar 8 13:26:26 UTC 2024



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


Commits:
d73943ab by Teo Camarasu at 2024-03-08T13:26:20+00:00
Implement user-defined allocation limit handlers

Resolves #22859

- - - - -


8 changed files:

- 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/Schedule.c
- rts/external-symbols.list.in
- rts/include/rts/storage/GC.h


Changes:

=====================================
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,6 @@
+module System.Mem.Experimental
+  ( setGlobalAllocationLimitHandler
+  , AllocationLimitKillBehaviour(..)
+  )
+  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,69 @@
+{-# LANGUAGE MagicHash #-}
+{-# OPTIONS_HADDOCK not-home #-}
+module GHC.Internal.AllocationLimitHandler
+  ( runAllocationLimitHandler
+  , setGlobalAllocationLimitHandler
+  , AllocationLimitKillBehaviour(..)
+  )
+  where
+import GHC.Internal.Foreign.C.Types
+import GHC.Internal.Control.Monad (return, join)
+import GHC.Internal.Data.IORef (IORef, readIORef, writeIORef, newIORef)
+import GHC.Internal.IO (IO, unsafePerformIO)
+import GHC.Internal.System.IO (print)
+import GHC.Internal.Conc.Sync (ThreadId(..))
+import GHC.Internal.Base
+
+
+{-# 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


=====================================
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/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.
@@ -1118,14 +1122,31 @@ schedulePostRunThread (Capability *cap, StgTSO *t)
     // the AllocationLimitExceeded exception.
 
     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) {
+          // 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 +3348,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
    -------------------------------------------------------------------------- */



View it on GitLab: https://gitlab.haskell.org/ghc/ghc/-/commit/d73943ab91d0a2e202cb8096a048b3ae9423bca9

-- 
View it on GitLab: https://gitlab.haskell.org/ghc/ghc/-/commit/d73943ab91d0a2e202cb8096a048b3ae9423bca9
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/20240308/bbb1eac3/attachment-0001.html>


More information about the ghc-commits mailing list