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.
-- The default behaviour is to throw an @AllocationLimitExceeded@ async exception to the thread.
-- This can be overriden using @AllocationLimitKillBehaviour@.
--
-- We can set a user-specified handler, which can be run 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.
-- You may wish to do this during a user-specified allocation limit handler.
--
-- 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 ()