{-
  Copyright (c) Meta Platforms, Inc. and affiliates.
  All rights reserved.

  This source code is licensed under the BSD-style license found in the
  LICENSE file in the root directory of this source tree.
-}

{-# OPTIONS_GHC -fno-warn-unused-imports #-}

module Util.AllocLimit
  ( limitAllocs
  , limitAllocsThrow
  , withAllocLimit
  , withSavedAllocLimit
  , withAllocCounter
  , handleAllExceptions
  ) where

import GHC.Conc
import Data.Int
import Control.Exception

-- | Impose an allocation limit on the given 'IO' action.  If the
-- action exceeds the limit, then 'Nothing' will be returned.
limitAllocs :: Int64 -> IO a -> IO (Maybe a)
limitAllocs limit action =
  handle (\AllocationLimitExceeded -> return Nothing) $
    Just <$> limitAllocsThrow limit action

-- | Impose an allocation limit on the given 'IO' action.  If the
-- action exceeds the limit, then 'AllocationLimitExceeded' will be
-- thrown.
limitAllocsThrow :: Int64 -> IO a -> IO a
limitAllocsThrow limit action = do
  e <- handleAllExceptions (return . Left) $
    Right <$> withSavedAllocLimit limit action
  either throwIO return e

-- | Impose an allocation limit on the given 'IO' action.  If the
-- action exceeds the limit, then 'AllocationLimitExceeded' will be
-- thrown. (Note: most uses should use 'limitAllocs' instead).
--
-- Note that if this function throws an exception
-- ('AllocationLimitExceeded' or otherwise), even though there will be
-- no /more/ 'AllocationLimitExceeded' exceptions, there might be
-- multiple 'AllocationLimitExceeded' exceptions already in flight.
-- This happens as follows: the first exception gets thrown, which
-- causes an exception handler to run, which triggers another
-- 'AllocationLimitExceeded' exception, which doesn't get thrown
-- immediately because the exception handler is implicitly masked, and
-- so on.
--
-- To deal with this we provide 'handleAllExceptions', which you
-- should wrap around 'withAllocLimit'.
--
-- Note: this function clobbers the allocation counter, so
-- e.g. `Util.Timing.timeIt` will give bogus results if the
-- computation being measured calls `withAllocLimit`. To avoid this,
-- use `withSavedAllocLimit` instead.
withAllocLimit :: Int64 -> IO a -> IO a
withAllocLimit limit =
  bracket_
    (do setAllocationCounter limit; enableAllocationLimit)
    disableAllocationLimit

-- | Like `withAllocLimit`, but does not clobber the allocation
-- counter.
withSavedAllocLimit :: Int64 -> IO a -> IO a
withSavedAllocLimit limit io = bracket set unset $ \_ -> io
  where
  set = do
    prev <- getAllocationCounter
    setAllocationCounter limit
    enableAllocationLimit
    return prev
  unset prev = do
    cur <- getAllocationCounter
    disableAllocationLimit
    setAllocationCounter (prev - (limit - cur))

-- | Sets the allocation counter for the given 'IO' action, then restores the
-- original counter value after the 'IO' action is done.
--
-- Note that as opposed to 'withAllocLimit' this function provides no guarantee
-- that if the set alloc counter is consumed, an 'AllocationLimitExceeded' will
-- be thrown, that depends on whether 'enableAllocationLimit' has been called
-- or not. If that guarantee is something you need, use 'withAllocLimit', but be
-- mindful of the fact that 'withAllocLimit' calls should not be nested inside
-- each other while 'withAllocCounter' calls can.
--
-- This function can be useful when there's a need to "isolate" certain
-- calls from the current allocation limit by setting a separate allocation
-- limit for them, then restoring the initial limit when done.
--
-- Once this is done: T30781590. We will be able to use the new primitive in
-- 'withAllocLimit' and deprecate 'withAllocCounter'
withAllocCounter :: Int64 -> IO a -> IO a
withAllocCounter counter action = do
  initialAllocCounter <- getAllocationCounter
  bracket_
    (setAllocationCounter counter)
    (setAllocationCounter initialAllocCounter)
    action

-- | Like 'Control.Exception.handle', but if there are pending async
-- exceptions, all are swallowed except for the last one, which is
-- passed to the handler.  This is mainly for wrapping around
-- 'withAllocLimit', to ensure that no 'AllocationLimitExceeded'
-- exceptions can leak out.
--
handleAllExceptions :: (SomeException -> IO a) -> IO a -> IO a
handleAllExceptions handler = handle errorHandler
  where
    errorHandler exn =
      handle errorHandler $ do
        allowInterrupt -- let any pending async exceptions throw
        handler exn