-- SPDX-FileCopyrightText: 2026 Alexandra de Wit
--
-- SPDX-License-Identifier: MIT

{- | The metadata memory gate's meter: one byte budget that requests pay into as they read and build.

A request takes a small entry step before its CPU slot and is shed, as a value, when that step does
not fit within the admission wait. A started request pays for what it reads one step at a time, and
pauses instead of failing when the budget is spent. One ticket at a time holds the overdraw token
until its request ends. Work that request waits on, a shared fetch or render, may overdraw on the
same token, so the overshoot stays near one request with its shared work, and a pause never deadlocks.
-}
module Ecluse.Core.Server.Admission.Meter (
    -- * The meter
    MemoryMeter,
    MeterSettings (..),
    newMemoryMeter,
    steerMeter,
    meterSnapshot,
    meterFigures,
    takeLargestCharge,

    -- * A request's ticket
    MemoryTicket,
    withMemoryEntry,
    charge,

    -- * Shared work
    awaitingFlight,
    servingFlight,
) where

import Control.Concurrent.STM (retry, stateTVar)
import Data.IntSet qualified as IntSet
import Data.Map.Strict qualified as Map
import GHC.Conc (registerDelay)
import UnliftIO (MonadUnliftIO)
import UnliftIO.Exception qualified as UE

import Ecluse.Core.Server.Admission.Budget (
    EntryGate (..),
    GrowthGate (..),
    MeterView (..),
    entryDecision,
    entryReady,
    growthDecision,
    roundUpToStep,
 )
import Ecluse.Core.Server.Admission.Types (BrakeLevel (Calm), FlightKey, MeterSnapshot (..))
import Ecluse.Core.Telemetry.Record (MetricsPort (..))

-- | The boot-time shape of a meter.
data MeterSettings = MeterSettings
    { MeterSettings -> Int
msBudgetBytes :: Int
    -- ^ The starting budget. The sampler may move it later.
    , MeterSettings -> Int
msStepBytes :: Int
    -- ^ The entry step, and the unit growth is paid in.
    , MeterSettings -> Int
msEntryRoom :: Int
    -- ^ How many new requests may wait at the gate at once.
    , MeterSettings -> Int
msEntryWaitMicros :: Int
    -- ^ How long a new request waits for its entry step before it is shed.
    }
    deriving stock (MeterSettings -> MeterSettings -> Bool
(MeterSettings -> MeterSettings -> Bool)
-> (MeterSettings -> MeterSettings -> Bool) -> Eq MeterSettings
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: MeterSettings -> MeterSettings -> Bool
== :: MeterSettings -> MeterSettings -> Bool
$c/= :: MeterSettings -> MeterSettings -> Bool
/= :: MeterSettings -> MeterSettings -> Bool
Eq, Int -> MeterSettings -> ShowS
[MeterSettings] -> ShowS
MeterSettings -> String
(Int -> MeterSettings -> ShowS)
-> (MeterSettings -> String)
-> ([MeterSettings] -> ShowS)
-> Show MeterSettings
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> MeterSettings -> ShowS
showsPrec :: Int -> MeterSettings -> ShowS
$cshow :: MeterSettings -> String
show :: MeterSettings -> String
$cshowList :: [MeterSettings] -> ShowS
showList :: [MeterSettings] -> ShowS
Show)

-- | The process-wide meter. The constructor stays hidden so only the checked operations change it.
data MemoryMeter = MemoryMeter
    { MemoryMeter -> TVar Int
mmBudget :: TVar Int
    , MemoryMeter -> TVar Int
mmCharged :: TVar Int
    , MemoryMeter -> TVar (Map Int (Maybe FlightKey))
mmWaiters :: TVar (Map Int (Maybe FlightKey))
    -- ^ Paused tickets, each with the shared work its paused charge serves.
    , MemoryMeter -> TVar (Map FlightKey IntSet)
mmInterest :: TVar (Map FlightKey IntSet)
    -- ^ The tickets waiting on each unit of shared work.
    , MemoryMeter -> TVar (Maybe Int)
mmToken :: TVar (Maybe Int)
    , MemoryMeter -> TVar Int
mmEntryWaiting :: TVar Int
    , MemoryMeter -> TVar BrakeLevel
mmBrakeLevel :: TVar BrakeLevel
    , MemoryMeter -> TVar Int
mmLargest :: TVar Int
    -- ^ The largest total any one ticket reached since the sampler last read it.
    , MemoryMeter -> IORef Int
mmNextTicket :: IORef Int
    , MemoryMeter -> Int
mmStepBytes :: Int
    , MemoryMeter -> Int
mmEntryRoom :: Int
    , MemoryMeter -> Int
mmEntryWaitMicros :: Int
    }

-- | Build a meter. The step floors at one byte, and the room and the wait at zero.
newMemoryMeter :: MeterSettings -> IO MemoryMeter
newMemoryMeter :: MeterSettings -> IO MemoryMeter
newMemoryMeter MeterSettings
settings =
    TVar Int
-> TVar Int
-> TVar (Map Int (Maybe FlightKey))
-> TVar (Map FlightKey IntSet)
-> TVar (Maybe Int)
-> TVar Int
-> TVar BrakeLevel
-> TVar Int
-> IORef Int
-> Int
-> Int
-> Int
-> MemoryMeter
MemoryMeter
        (TVar Int
 -> TVar Int
 -> TVar (Map Int (Maybe FlightKey))
 -> TVar (Map FlightKey IntSet)
 -> TVar (Maybe Int)
 -> TVar Int
 -> TVar BrakeLevel
 -> TVar Int
 -> IORef Int
 -> Int
 -> Int
 -> Int
 -> MemoryMeter)
-> IO (TVar Int)
-> IO
     (TVar Int
      -> TVar (Map Int (Maybe FlightKey))
      -> TVar (Map FlightKey IntSet)
      -> TVar (Maybe Int)
      -> TVar Int
      -> TVar BrakeLevel
      -> TVar Int
      -> IORef Int
      -> Int
      -> Int
      -> Int
      -> MemoryMeter)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> IO (TVar Int)
forall (m :: * -> *) a. MonadIO m => a -> m (TVar a)
newTVarIO (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (MeterSettings -> Int
msBudgetBytes MeterSettings
settings))
        IO
  (TVar Int
   -> TVar (Map Int (Maybe FlightKey))
   -> TVar (Map FlightKey IntSet)
   -> TVar (Maybe Int)
   -> TVar Int
   -> TVar BrakeLevel
   -> TVar Int
   -> IORef Int
   -> Int
   -> Int
   -> Int
   -> MemoryMeter)
-> IO (TVar Int)
-> IO
     (TVar (Map Int (Maybe FlightKey))
      -> TVar (Map FlightKey IntSet)
      -> TVar (Maybe Int)
      -> TVar Int
      -> TVar BrakeLevel
      -> TVar Int
      -> IORef Int
      -> Int
      -> Int
      -> Int
      -> MemoryMeter)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> IO (TVar Int)
forall (m :: * -> *) a. MonadIO m => a -> m (TVar a)
newTVarIO Int
0
        IO
  (TVar (Map Int (Maybe FlightKey))
   -> TVar (Map FlightKey IntSet)
   -> TVar (Maybe Int)
   -> TVar Int
   -> TVar BrakeLevel
   -> TVar Int
   -> IORef Int
   -> Int
   -> Int
   -> Int
   -> MemoryMeter)
-> IO (TVar (Map Int (Maybe FlightKey)))
-> IO
     (TVar (Map FlightKey IntSet)
      -> TVar (Maybe Int)
      -> TVar Int
      -> TVar BrakeLevel
      -> TVar Int
      -> IORef Int
      -> Int
      -> Int
      -> Int
      -> MemoryMeter)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Map Int (Maybe FlightKey) -> IO (TVar (Map Int (Maybe FlightKey)))
forall (m :: * -> *) a. MonadIO m => a -> m (TVar a)
newTVarIO Map Int (Maybe FlightKey)
forall k a. Map k a
Map.empty
        IO
  (TVar (Map FlightKey IntSet)
   -> TVar (Maybe Int)
   -> TVar Int
   -> TVar BrakeLevel
   -> TVar Int
   -> IORef Int
   -> Int
   -> Int
   -> Int
   -> MemoryMeter)
-> IO (TVar (Map FlightKey IntSet))
-> IO
     (TVar (Maybe Int)
      -> TVar Int
      -> TVar BrakeLevel
      -> TVar Int
      -> IORef Int
      -> Int
      -> Int
      -> Int
      -> MemoryMeter)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Map FlightKey IntSet -> IO (TVar (Map FlightKey IntSet))
forall (m :: * -> *) a. MonadIO m => a -> m (TVar a)
newTVarIO Map FlightKey IntSet
forall k a. Map k a
Map.empty
        IO
  (TVar (Maybe Int)
   -> TVar Int
   -> TVar BrakeLevel
   -> TVar Int
   -> IORef Int
   -> Int
   -> Int
   -> Int
   -> MemoryMeter)
-> IO (TVar (Maybe Int))
-> IO
     (TVar Int
      -> TVar BrakeLevel
      -> TVar Int
      -> IORef Int
      -> Int
      -> Int
      -> Int
      -> MemoryMeter)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe Int -> IO (TVar (Maybe Int))
forall (m :: * -> *) a. MonadIO m => a -> m (TVar a)
newTVarIO Maybe Int
forall a. Maybe a
Nothing
        IO
  (TVar Int
   -> TVar BrakeLevel
   -> TVar Int
   -> IORef Int
   -> Int
   -> Int
   -> Int
   -> MemoryMeter)
-> IO (TVar Int)
-> IO
     (TVar BrakeLevel
      -> TVar Int -> IORef Int -> Int -> Int -> Int -> MemoryMeter)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> IO (TVar Int)
forall (m :: * -> *) a. MonadIO m => a -> m (TVar a)
newTVarIO Int
0
        IO
  (TVar BrakeLevel
   -> TVar Int -> IORef Int -> Int -> Int -> Int -> MemoryMeter)
-> IO (TVar BrakeLevel)
-> IO (TVar Int -> IORef Int -> Int -> Int -> Int -> MemoryMeter)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> BrakeLevel -> IO (TVar BrakeLevel)
forall (m :: * -> *) a. MonadIO m => a -> m (TVar a)
newTVarIO BrakeLevel
Calm
        IO (TVar Int -> IORef Int -> Int -> Int -> Int -> MemoryMeter)
-> IO (TVar Int)
-> IO (IORef Int -> Int -> Int -> Int -> MemoryMeter)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> IO (TVar Int)
forall (m :: * -> *) a. MonadIO m => a -> m (TVar a)
newTVarIO Int
0
        IO (IORef Int -> Int -> Int -> Int -> MemoryMeter)
-> IO (IORef Int) -> IO (Int -> Int -> Int -> MemoryMeter)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> IO (IORef Int)
forall (m :: * -> *) a. MonadIO m => a -> m (IORef a)
newIORef Int
0
        IO (Int -> Int -> Int -> MemoryMeter)
-> IO Int -> IO (Int -> Int -> MemoryMeter)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> IO Int
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 (MeterSettings -> Int
msStepBytes MeterSettings
settings))
        IO (Int -> Int -> MemoryMeter) -> IO Int -> IO (Int -> MemoryMeter)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> IO Int
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (MeterSettings -> Int
msEntryRoom MeterSettings
settings))
        IO (Int -> MemoryMeter) -> IO Int -> IO MemoryMeter
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> IO Int
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (MeterSettings -> Int
msEntryWaitMicros MeterSettings
settings))

-- | Move the budget and record the brake level that moved it. Paused reads and queued entries retry at once.
steerMeter :: MemoryMeter -> Int -> BrakeLevel -> IO ()
steerMeter :: MemoryMeter -> Int -> BrakeLevel -> IO ()
steerMeter MemoryMeter
meter Int
budget BrakeLevel
level = STM () -> IO ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
    TVar Int -> Int -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (MemoryMeter -> TVar Int
mmBudget MemoryMeter
meter) (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 Int
budget)
    TVar BrakeLevel -> BrakeLevel -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (MemoryMeter -> TVar BrakeLevel
mmBrakeLevel MemoryMeter
meter) BrakeLevel
level

-- | Read the meter's figures without blocking a request.
meterSnapshot :: MemoryMeter -> IO MeterSnapshot
meterSnapshot :: MemoryMeter -> IO MeterSnapshot
meterSnapshot = STM MeterSnapshot -> IO MeterSnapshot
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM MeterSnapshot -> IO MeterSnapshot)
-> (MemoryMeter -> STM MeterSnapshot)
-> MemoryMeter
-> IO MeterSnapshot
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MemoryMeter -> STM MeterSnapshot
meterFigures

-- | The meter's figures inside a transaction, so a caller can wait for them to change.
meterFigures :: MemoryMeter -> STM MeterSnapshot
meterFigures :: MemoryMeter -> STM MeterSnapshot
meterFigures MemoryMeter
meter =
    Int -> Int -> Int -> Int -> BrakeLevel -> MeterSnapshot
MeterSnapshot
        (Int -> Int -> Int -> Int -> BrakeLevel -> MeterSnapshot)
-> STM Int
-> STM (Int -> Int -> Int -> BrakeLevel -> MeterSnapshot)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar (MemoryMeter -> TVar Int
mmBudget MemoryMeter
meter)
        STM (Int -> Int -> Int -> BrakeLevel -> MeterSnapshot)
-> STM Int -> STM (Int -> Int -> BrakeLevel -> MeterSnapshot)
forall a b. STM (a -> b) -> STM a -> STM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar (MemoryMeter -> TVar Int
mmCharged MemoryMeter
meter)
        STM (Int -> Int -> BrakeLevel -> MeterSnapshot)
-> STM Int -> STM (Int -> BrakeLevel -> MeterSnapshot)
forall a b. STM (a -> b) -> STM a -> STM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar (MemoryMeter -> TVar Int
mmEntryWaiting MemoryMeter
meter)
        STM (Int -> BrakeLevel -> MeterSnapshot)
-> STM Int -> STM (BrakeLevel -> MeterSnapshot)
forall a b. STM (a -> b) -> STM a -> STM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Map Int (Maybe FlightKey) -> Int
forall k a. Map k a -> Int
Map.size (Map Int (Maybe FlightKey) -> Int)
-> STM (Map Int (Maybe FlightKey)) -> STM Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar (Map Int (Maybe FlightKey)) -> STM (Map Int (Maybe FlightKey))
forall a. TVar a -> STM a
readTVar (MemoryMeter -> TVar (Map Int (Maybe FlightKey))
mmWaiters MemoryMeter
meter))
        STM (BrakeLevel -> MeterSnapshot)
-> STM BrakeLevel -> STM MeterSnapshot
forall a b. STM (a -> b) -> STM a -> STM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TVar BrakeLevel -> STM BrakeLevel
forall a. TVar a -> STM a
readTVar (MemoryMeter -> TVar BrakeLevel
mmBrakeLevel MemoryMeter
meter)

-- | The largest total one ticket reached since the last call, for the sampler alone.
takeLargestCharge :: MemoryMeter -> IO Int
takeLargestCharge :: MemoryMeter -> IO Int
takeLargestCharge MemoryMeter
meter = STM Int -> IO Int
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (TVar Int -> (Int -> (Int, Int)) -> STM Int
forall s a. TVar s -> (s -> (a, s)) -> STM a
stateTVar (MemoryMeter -> TVar Int
mmLargest MemoryMeter
meter) (,Int
0))

-- | One request's account with the meter, or a view of it that pays for shared work.
data MemoryTicket = MemoryTicket
    { MemoryTicket -> MemoryMeter
tMeter :: MemoryMeter
    , MemoryTicket -> Int
tId :: Int
    , MemoryTicket -> TVar Int
tCharged :: TVar Int
    -- ^ Everything this ticket took from the budget, returned in one transaction at the end.
    , MemoryTicket -> IORef Int
tHeadroom :: IORef Int
    -- ^ Paid but not yet used, so most chunks touch no shared state.
    , MemoryTicket -> Maybe FlightKey
tServes :: Maybe FlightKey
    -- ^ The shared work this view pays for, whose waiters lend it their priority.
    , MemoryTicket -> MetricsPort
tMetrics :: MetricsPort
    }

{- | Run a request under the meter. 'Nothing' is a shed: the waiting room was full, or the entry
step did not fit within the wait. Everything the ticket paid returns on every exit path.
-}
withMemoryEntry :: (MonadUnliftIO m) => MetricsPort -> MemoryMeter -> (MemoryTicket -> m a) -> m (Maybe a)
withMemoryEntry :: forall (m :: * -> *) a.
MonadUnliftIO m =>
MetricsPort -> MemoryMeter -> (MemoryTicket -> m a) -> m (Maybe a)
withMemoryEntry MetricsPort
metrics MemoryMeter
meter MemoryTicket -> m a
body =
    ((forall a. m a -> m a) -> m (Maybe a)) -> m (Maybe a)
forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> m a) -> m b) -> m b
UE.mask (((forall a. m a -> m a) -> m (Maybe a)) -> m (Maybe a))
-> ((forall a. m a -> m a) -> m (Maybe a)) -> m (Maybe a)
forall a b. (a -> b) -> a -> b
$ \forall a. m a -> m a
restore -> do
        ticket <- IO MemoryTicket -> m MemoryTicket
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (MetricsPort -> MemoryMeter -> IO MemoryTicket
newTicket MetricsPort
metrics MemoryMeter
meter)
        entered <- liftIO (enter ticket)
        case entered of
            Maybe Bool
Nothing -> IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (MetricsPort -> IO ()
mpMemoryAdmissionShed MetricsPort
metrics) m () -> Maybe a -> m (Maybe a)
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> Maybe a
forall a. Maybe a
Nothing
            Just Bool
waited ->
                a -> Maybe a
forall a. a -> Maybe a
Just
                    (a -> Maybe a) -> m a -> m (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ( m a -> m a
forall a. m a -> m a
restore (IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
waited (MetricsPort -> IO ()
mpMemoryAdmissionQueued MetricsPort
metrics)) m () -> m a -> m a
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> MemoryTicket -> m a
body MemoryTicket
ticket)
                            m a -> m () -> m a
forall (m :: * -> *) a b. MonadUnliftIO m => m a -> m b -> m a
`UE.finally` IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (MemoryTicket -> IO ()
releaseTicket MemoryTicket
ticket)
                        )

newTicket :: MetricsPort -> MemoryMeter -> IO MemoryTicket
newTicket :: MetricsPort -> MemoryMeter -> IO MemoryTicket
newTicket MetricsPort
metrics MemoryMeter
meter = do
    ticketId <- IORef Int -> (Int -> (Int, Int)) -> IO Int
forall (m :: * -> *) a b.
MonadIO m =>
IORef a -> (a -> (a, b)) -> m b
atomicModifyIORef' (MemoryMeter -> IORef Int
mmNextTicket MemoryMeter
meter) (\Int
n -> (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, Int
n))
    charged <- newTVarIO 0
    headroom <- newIORef 0
    pure MemoryTicket{tMeter = meter, tId = ticketId, tCharged = charged, tHeadroom = headroom, tServes = Nothing, tMetrics = metrics}

-- The gate. 'Just' carries whether the request had to wait. It runs under the caller's mask, and
-- the queued count taken in the deciding transaction is returned before any interruptible step.
enter :: MemoryTicket -> IO (Maybe Bool)
enter :: MemoryTicket -> IO (Maybe Bool)
enter MemoryTicket
ticket = do
    gate <- STM EntryGate -> IO EntryGate
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM EntryGate -> IO EntryGate) -> STM EntryGate -> IO EntryGate
forall a b. (a -> b) -> a -> b
$ do
        view <- MemoryMeter -> STM MeterView
readView MemoryMeter
meter
        waiting <- readTVar (mmEntryWaiting meter)
        case entryDecision view waiting (mmEntryRoom meter) step of
            EntryGate
EntryAdmit -> MemoryTicket -> Int -> STM ()
takeBytes MemoryTicket
ticket Int
step STM () -> EntryGate -> STM EntryGate
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> EntryGate
EntryAdmit
            EntryGate
EntryQueue -> TVar Int -> Int -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (MemoryMeter -> TVar Int
mmEntryWaiting MemoryMeter
meter) (Int
waiting Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) STM () -> EntryGate -> STM EntryGate
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> EntryGate
EntryQueue
            EntryGate
EntryRefuse -> EntryGate -> STM EntryGate
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure EntryGate
EntryRefuse
    case gate of
        EntryGate
EntryAdmit -> IO ()
openHeadroom IO () -> Maybe Bool -> IO (Maybe Bool)
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> Bool -> Maybe Bool
forall a. a -> Maybe a
Just Bool
False
        EntryGate
EntryRefuse -> Maybe Bool -> IO (Maybe Bool)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Bool
forall a. Maybe a
Nothing
        EntryGate
EntryQueue -> do
            admitted <-
                (Int -> IO (TVar Bool)
registerDelay (MemoryMeter -> Int
mmEntryWaitMicros MemoryMeter
meter) IO (TVar Bool) -> (TVar Bool -> IO Bool) -> IO Bool
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= STM Bool -> IO Bool
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM Bool -> IO Bool)
-> (TVar Bool -> STM Bool) -> TVar Bool -> IO Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TVar Bool -> STM Bool
entryOrExpire)
                    IO Bool -> IO () -> IO Bool
forall (m :: * -> *) a b. MonadUnliftIO m => m a -> m b -> m a
`UE.finally` STM () -> IO ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (TVar Int -> (Int -> Int) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' (MemoryMeter -> TVar Int
mmEntryWaiting MemoryMeter
meter) (Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
1))
            if admitted then openHeadroom $> Just True else pure Nothing
  where
    meter :: MemoryMeter
meter = MemoryTicket -> MemoryMeter
tMeter MemoryTicket
ticket
    step :: Int
step = MemoryMeter -> Int
mmStepBytes MemoryMeter
meter
    openHeadroom :: IO ()
openHeadroom = IORef Int -> Int -> IO ()
forall (m :: * -> *) a. MonadIO m => IORef a -> a -> m ()
writeIORef (MemoryTicket -> IORef Int
tHeadroom MemoryTicket
ticket) Int
step
    entryOrExpire :: TVar Bool -> STM Bool
entryOrExpire TVar Bool
deadline = do
        view <- MemoryMeter -> STM MeterView
readView MemoryMeter
meter
        if entryReady view step
            then takeBytes ticket step $> True
            else do
                expired <- readTVar deadline
                if expired then pure False else retry

-- The oldest waiter counts the tickets waiting on its shared work, so it may carry their priority.
readView :: MemoryMeter -> STM MeterView
readView :: MemoryMeter -> STM MeterView
readView MemoryMeter
meter = do
    interest <- TVar (Map FlightKey IntSet) -> STM (Map FlightKey IntSet)
forall a. TVar a -> STM a
readTVar (MemoryMeter -> TVar (Map FlightKey IntSet)
mmInterest MemoryMeter
meter)
    waiters <- readTVar (mmWaiters meter)
    MeterView
        <$> readTVar (mmBudget meter)
        <*> readTVar (mmCharged meter)
        <*> readTVar (mmToken meter)
        <*> pure (smallest (IntSet.fromList [oldestOf interest waiter serves | (waiter, serves) <- Map.toList waiters]))
  where
    oldestOf :: Map FlightKey IntSet -> Int -> Maybe FlightKey -> Int
oldestOf Map FlightKey IntSet
interest Int
waiter Maybe FlightKey
serves = Int -> (Int -> Int) -> Maybe Int -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int
waiter (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
waiter) (IntSet -> Maybe Int
smallest (Map FlightKey IntSet -> Maybe FlightKey -> IntSet
waitingOn Map FlightKey IntSet
interest Maybe FlightKey
serves))

smallest :: IntSet -> Maybe Int
smallest :: IntSet -> Maybe Int
smallest = ((Int, IntSet) -> Int) -> Maybe (Int, IntSet) -> Maybe Int
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Int, IntSet) -> Int
forall a b. (a, b) -> a
fst (Maybe (Int, IntSet) -> Maybe Int)
-> (IntSet -> Maybe (Int, IntSet)) -> IntSet -> Maybe Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntSet -> Maybe (Int, IntSet)
IntSet.minView

waitingOn :: Map FlightKey IntSet -> Maybe FlightKey -> IntSet
waitingOn :: Map FlightKey IntSet -> Maybe FlightKey -> IntSet
waitingOn Map FlightKey IntSet
interest = IntSet -> (FlightKey -> IntSet) -> Maybe FlightKey -> IntSet
forall b a. b -> (a -> b) -> Maybe a -> b
maybe IntSet
IntSet.empty (\FlightKey
key -> IntSet -> FlightKey -> Map FlightKey IntSet -> IntSet
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault IntSet
IntSet.empty FlightKey
key Map FlightKey IntSet
interest)

-- Take bytes onto both the meter and the ticket in one step, so a release can never miss them.
takeBytes :: MemoryTicket -> Int -> STM ()
takeBytes :: MemoryTicket -> Int -> STM ()
takeBytes MemoryTicket
ticket Int
bytes = do
    let meter :: MemoryMeter
meter = MemoryTicket -> MemoryMeter
tMeter MemoryTicket
ticket
    TVar Int -> (Int -> Int) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' (MemoryMeter -> TVar Int
mmCharged MemoryMeter
meter) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
bytes)
    total <- (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
bytes) (Int -> Int) -> STM Int -> STM Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar (MemoryTicket -> TVar Int
tCharged MemoryTicket
ticket)
    writeTVar (tCharged ticket) total
    modifyTVar' (mmLargest meter) (max total)

releaseTicket :: MemoryTicket -> IO ()
releaseTicket :: MemoryTicket -> IO ()
releaseTicket MemoryTicket
ticket = STM () -> IO ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
    let meter :: MemoryMeter
meter = MemoryTicket -> MemoryMeter
tMeter MemoryTicket
ticket
    held <- TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar (MemoryTicket -> TVar Int
tCharged MemoryTicket
ticket)
    writeTVar (tCharged ticket) 0
    modifyTVar' (mmCharged meter) (subtract held)
    holder <- readTVar (mmToken meter)
    when (holder == Just (tId ticket)) (writeTVar (mmToken meter) Nothing)

{- | Mark this request as waiting on a unit of shared work for the length of the action, so the
request that does that work may pause and overdraw with this request's priority.
-}
awaitingFlight :: (MonadUnliftIO m) => MemoryTicket -> FlightKey -> m a -> m a
awaitingFlight :: forall (m :: * -> *) a.
MonadUnliftIO m =>
MemoryTicket -> FlightKey -> m a -> m a
awaitingFlight MemoryTicket
ticket FlightKey
key =
    m () -> m () -> m a -> m a
forall (m :: * -> *) a b c.
MonadUnliftIO m =>
m a -> m b -> m c -> m c
UE.bracket_ (IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (STM () -> IO ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically ((IntSet -> IntSet) -> STM ()
update (Int -> IntSet -> IntSet
IntSet.insert (MemoryTicket -> Int
tId MemoryTicket
ticket))))) (IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (STM () -> IO ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically ((IntSet -> IntSet) -> STM ()
update (Int -> IntSet -> IntSet
IntSet.delete (MemoryTicket -> Int
tId MemoryTicket
ticket)))))
  where
    interest :: TVar (Map FlightKey IntSet)
interest = MemoryMeter -> TVar (Map FlightKey IntSet)
mmInterest (MemoryTicket -> MemoryMeter
tMeter MemoryTicket
ticket)
    update :: (IntSet -> IntSet) -> STM ()
update IntSet -> IntSet
change = TVar (Map FlightKey IntSet)
-> (Map FlightKey IntSet -> Map FlightKey IntSet) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar (Map FlightKey IntSet)
interest ((Maybe IntSet -> Maybe IntSet)
-> FlightKey -> Map FlightKey IntSet -> Map FlightKey IntSet
forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
Map.alter (IntSet -> Maybe IntSet
nonEmptySet (IntSet -> Maybe IntSet)
-> (Maybe IntSet -> IntSet) -> Maybe IntSet -> Maybe IntSet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntSet -> IntSet
change (IntSet -> IntSet)
-> (Maybe IntSet -> IntSet) -> Maybe IntSet -> IntSet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntSet -> Maybe IntSet -> IntSet
forall a. a -> Maybe a -> a
fromMaybe IntSet
IntSet.empty) FlightKey
key)
    nonEmptySet :: IntSet -> Maybe IntSet
nonEmptySet IntSet
set = if IntSet -> Bool
IntSet.null IntSet
set then Maybe IntSet
forall a. Maybe a
Nothing else IntSet -> Maybe IntSet
forall a. a -> Maybe a
Just IntSet
set

-- | The same ticket, paying for shared work: its charges carry the priority of every request waiting on it.
servingFlight :: FlightKey -> MemoryTicket -> MemoryTicket
servingFlight :: FlightKey -> MemoryTicket -> MemoryTicket
servingFlight FlightKey
key MemoryTicket
ticket = MemoryTicket
ticket{tServes = Just key}

-- | Pay for bytes the request is about to hold, already scaled to their charge.
charge :: MemoryTicket -> Int -> IO ()
charge :: MemoryTicket -> Int -> IO ()
charge MemoryTicket
ticket Int
cost = Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
cost Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
    shortfall <- IORef Int -> (Int -> (Int, Int)) -> IO Int
forall (m :: * -> *) a b.
MonadIO m =>
IORef a -> (a -> (a, b)) -> m b
atomicModifyIORef' (MemoryTicket -> IORef Int
tHeadroom MemoryTicket
ticket) ((Int -> (Int, Int)) -> IO Int) -> (Int -> (Int, Int)) -> IO Int
forall a b. (a -> b) -> a -> b
$ \Int
headroom ->
        if Int
cost Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
headroom then (Int
headroom Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
cost, Int
0) else (Int
0, Int
cost Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
headroom)
    when (shortfall > 0) $ do
        let want = Int -> Int -> Int
roundUpToStep (MemoryMeter -> Int
mmStepBytes (MemoryTicket -> MemoryMeter
tMeter MemoryTicket
ticket)) Int
shortfall
        grow ticket want
        atomicModifyIORef' (tHeadroom ticket) (\Int
headroom -> (Int
headroom Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
want Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
shortfall, ()))

-- Take a growth step, pausing while it neither fits nor may overdraw. A pause registers the
-- ticket as a waiter, which holds new entries back and names the oldest claimant.
grow :: MemoryTicket -> Int -> IO ()
grow :: MemoryTicket -> Int -> IO ()
grow MemoryTicket
ticket Int
want = do
    immediate <- STM (Maybe Bool) -> IO (Maybe Bool)
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (MemoryTicket -> Int -> STM (Maybe Bool)
attempt MemoryTicket
ticket Int
want)
    overdrew <- case immediate of
        Just Bool
overdrew -> Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
overdrew
        Maybe Bool
Nothing -> do
            MetricsPort -> IO ()
mpMemoryAdmissionPause (MemoryTicket -> MetricsPort
tMetrics MemoryTicket
ticket)
            (STM () -> IO ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (TVar (Map Int (Maybe FlightKey))
-> (Map Int (Maybe FlightKey) -> Map Int (Maybe FlightKey))
-> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar (Map Int (Maybe FlightKey))
waiters (Int
-> Maybe FlightKey
-> Map Int (Maybe FlightKey)
-> Map Int (Maybe FlightKey)
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert (MemoryTicket -> Int
tId MemoryTicket
ticket) (MemoryTicket -> Maybe FlightKey
tServes MemoryTicket
ticket))) IO () -> IO Bool -> IO Bool
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> STM Bool -> IO Bool
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (MemoryTicket -> Int -> STM (Maybe Bool)
attempt MemoryTicket
ticket Int
want STM (Maybe Bool) -> (Maybe Bool -> STM Bool) -> STM Bool
forall a b. STM a -> (a -> STM b) -> STM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= STM Bool -> (Bool -> STM Bool) -> Maybe Bool -> STM Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe STM Bool
forall a. STM a
retry Bool -> STM Bool
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure))
                IO Bool -> IO () -> IO Bool
forall (m :: * -> *) a b. MonadUnliftIO m => m a -> m b -> m a
`UE.finally` STM () -> IO ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (TVar (Map Int (Maybe FlightKey))
-> (Map Int (Maybe FlightKey) -> Map Int (Maybe FlightKey))
-> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar (Map Int (Maybe FlightKey))
waiters (Int -> Map Int (Maybe FlightKey) -> Map Int (Maybe FlightKey)
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete (MemoryTicket -> Int
tId MemoryTicket
ticket)))
    when overdrew (mpMemoryAdmissionOverdraw (tMetrics ticket))
  where
    waiters :: TVar (Map Int (Maybe FlightKey))
waiters = MemoryMeter -> TVar (Map Int (Maybe FlightKey))
mmWaiters (MemoryTicket -> MemoryMeter
tMeter MemoryTicket
ticket)

-- One growth attempt: 'Just' whether it overdrew, or 'Nothing' to pause.
attempt :: MemoryTicket -> Int -> STM (Maybe Bool)
attempt :: MemoryTicket -> Int -> STM (Maybe Bool)
attempt MemoryTicket
ticket Int
want = do
    let meter :: MemoryMeter
meter = MemoryTicket -> MemoryMeter
tMeter MemoryTicket
ticket
    view <- MemoryMeter -> STM MeterView
readView MemoryMeter
meter
    served <- IntSet.insert (tId ticket) . (`waitingOn` tServes ticket) <$> readTVar (mmInterest meter)
    case growthDecision view served want of
        GrowthGate
GrowWithin -> MemoryTicket -> Int -> STM ()
takeBytes MemoryTicket
ticket Int
want STM () -> Maybe Bool -> STM (Maybe Bool)
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> Bool -> Maybe Bool
forall a. a -> Maybe a
Just Bool
False
        GrowthGate
GrowOnToken -> MemoryTicket -> Int -> STM ()
takeBytes MemoryTicket
ticket Int
want STM () -> Maybe Bool -> STM (Maybe Bool)
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> Bool -> Maybe Bool
forall a. a -> Maybe a
Just Bool
True
        GrowthGate
GrowTakeToken -> do
            TVar (Maybe Int) -> Maybe Int -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (MemoryMeter -> TVar (Maybe Int)
mmToken MemoryMeter
meter) (Int -> Maybe Int
forall a. a -> Maybe a
Just (MemoryTicket -> Int
tId MemoryTicket
ticket))
            MemoryTicket -> Int -> STM ()
takeBytes MemoryTicket
ticket Int
want STM () -> Maybe Bool -> STM (Maybe Bool)
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> Bool -> Maybe Bool
forall a. a -> Maybe a
Just Bool
True
        GrowthGate
GrowWait -> Maybe Bool -> STM (Maybe Bool)
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Bool
forall a. Maybe a
Nothing