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

{- | The weighted door\/wait\/shed admission core behind serve admission
("Ecluse.Core.Server.Admission") and byte-weighted publish admission
("Ecluse.Core.Server.Admission.Bytes").

A handle caps the aggregate weight held at once and keeps a bounded room of waiters. Capacity
is taken directly only when the room is empty, so a newcomer never jumps a non-empty room,
though wake order within the room is not FIFO. The wait budget equals the shed path's
@Retry-After: 1@ hint, so nothing is refused faster than the client was told to wait.
-}
module Ecluse.Core.Server.Admission.Weighted (
    WeightedAdmission,
    newWeightedAdmission,
    withWeightedAdmission,
    AdmissionObservers (..),
    admissionWaitMicros,
) where

import Control.Concurrent.STM (retry)
import GHC.Conc (registerDelay)
import UnliftIO (MonadUnliftIO)
import UnliftIO.Exception qualified as UE

{- | The bounded handle's mutable state and its tuning. The constructor stays hidden so only
the checked acquire, wait, and release operations can mutate it.
-}
data WeightedAdmission = WeightedAdmission
    { WeightedAdmission -> TVar Int
waAvailable :: TVar Int
    , WeightedAdmission -> TVar Int
waWaiting :: TVar Int
    , WeightedAdmission -> Int
waWaitingRoom :: Int
    , WeightedAdmission -> Int
waWaitMicros :: Int
    }

{- | The metric hooks the door\/wait\/release bracket calls. The shared machine owns no
telemetry vocabulary, so each instance records under its own signals.
-}
data AdmissionObservers = AdmissionObservers
    { AdmissionObservers -> IO ()
onQueued :: IO ()
    {- ^ A request that had to wait cleared the wait and is now admitted. Serve
    admission records its queued metric here. Byte admission does nothing.
    -}
    , AdmissionObservers -> IO ()
onShed :: IO ()
    {- ^ The request was shed: refused at a full door, or its wait outlived the
    budget. Byte admission records its shed metric here. Serve admission is silent.
    -}
    , AdmissionObservers -> Int -> IO ()
onInFlightDelta :: Int -> IO ()
    {- ^ Move the in-flight gauge by the signed weight. Both calls run under the
    acquire mask, so the gauge is paired on every path.
    -}
    }

{- | The wait budget (microseconds) before a busy acquisition is shed, deliberately equal to
the shed path's @Retry-After: 1@ hint, so nothing is refused faster than the client was told.
-}
admissionWaitMicros :: Int
admissionWaitMicros :: Int
admissionWaitMicros = Int
1_000_000

{- | Allocate a handle over a capacity, a waiter-room bound, and a wait budget (microseconds).
The capacity is verbatim, the wrapper owning that policy, and the other two floor at zero.
-}
newWeightedAdmission :: Int -> Int -> Int -> IO WeightedAdmission
newWeightedAdmission :: Int -> Int -> Int -> IO WeightedAdmission
newWeightedAdmission Int
capacity Int
room Int
waitMicros = do
    available <- Int -> IO (TVar Int)
forall (m :: * -> *) a. MonadIO m => a -> m (TVar a)
newTVarIO Int
capacity
    waiting <- newTVarIO 0
    pure
        WeightedAdmission
            { waAvailable = available
            , waWaiting = waiting
            , waWaitingRoom = max 0 room
            , waWaitMicros = max 0 waitMicros
            }

-- The outcome of the door transaction: the weight taken directly, a place taken in the
-- waiting room, or a refusal (the room was full).
data Gate = Admitted | Queued | Refused

-- The door transaction: decide a 'Gate' in one STM step. The weight is taken directly
-- only when no one is waiting, so a newcomer never jumps a non-empty waiting room.
doorDecision :: WeightedAdmission -> Int -> STM Gate
doorDecision :: WeightedAdmission -> Int -> STM Gate
doorDecision WeightedAdmission
wa Int
weight = do
    available <- TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar (WeightedAdmission -> TVar Int
waAvailable WeightedAdmission
wa)
    waiting <- readTVar (waWaiting wa)
    if available >= weight && waiting == 0
        then writeTVar (waAvailable wa) (available - weight) $> Admitted
        else
            if waiting >= waWaitingRoom wa
                then pure Refused
                else writeTVar (waWaiting wa) (waiting + 1) $> Queued

-- Take the weight the moment it fits, or report expiry. One transaction, so a
-- timeout can never race a committed acquire into leaked weight.
acquireOrExpire :: WeightedAdmission -> Int -> TVar Bool -> STM Bool
acquireOrExpire :: WeightedAdmission -> Int -> TVar Bool -> STM Bool
acquireOrExpire WeightedAdmission
wa Int
weight TVar Bool
deadline = do
    available <- TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar (WeightedAdmission -> TVar Int
waAvailable WeightedAdmission
wa)
    if available >= weight
        then writeTVar (waAvailable wa) (available - weight) $> True
        else do
            expired <- readTVar deadline
            if expired then pure False else retry

{- | Run an action holding the given weight. 'Nothing' is a shed, at a full room or an expired
wait. The weight is used as given, and released on every exit path, cancellation included.
-}

-- Inlined with its arm helpers so each wrapper's literal observers vanish at the call site.
{-# INLINE withWeightedAdmission #-}
withWeightedAdmission ::
    (MonadUnliftIO m) =>
    AdmissionObservers ->
    WeightedAdmission ->
    Int ->
    m a ->
    m (Maybe a)
withWeightedAdmission :: forall (m :: * -> *) a.
MonadUnliftIO m =>
AdmissionObservers
-> WeightedAdmission -> Int -> m a -> m (Maybe a)
withWeightedAdmission AdmissionObservers
obs WeightedAdmission
wa Int
weight m a
action =
    ((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
        gate <- STM Gate -> m Gate
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (WeightedAdmission -> Int -> STM Gate
doorDecision WeightedAdmission
wa Int
weight)
        case gate of
            Gate
Refused -> AdmissionObservers -> m (Maybe a)
forall (m :: * -> *) a.
MonadIO m =>
AdmissionObservers -> m (Maybe a)
shedRecording AdmissionObservers
obs
            Gate
Admitted -> AdmissionObservers
-> WeightedAdmission
-> Int
-> IO ()
-> (m a -> m a)
-> m a
-> m (Maybe a)
forall (m :: * -> *) a.
MonadUnliftIO m =>
AdmissionObservers
-> WeightedAdmission
-> Int
-> IO ()
-> (m a -> m a)
-> m a
-> m (Maybe a)
admittedRun AdmissionObservers
obs WeightedAdmission
wa Int
weight (() -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()) m a -> m a
forall a. m a -> m a
restore m a
action
            Gate
Queued -> AdmissionObservers
-> WeightedAdmission -> Int -> (m a -> m a) -> m a -> m (Maybe a)
forall (m :: * -> *) a.
MonadUnliftIO m =>
AdmissionObservers
-> WeightedAdmission -> Int -> (m a -> m a) -> m a -> m (Maybe a)
queuedWait AdmissionObservers
obs WeightedAdmission
wa Int
weight m a -> m a
forall a. m a -> m a
restore m a
action

-- Record the shed and refuse. A room place taken on the queued path is already
-- surrendered before this runs.
{-# INLINE shedRecording #-}
shedRecording :: (MonadIO m) => AdmissionObservers -> m (Maybe a)
shedRecording :: forall (m :: * -> *) a.
MonadIO m =>
AdmissionObservers -> m (Maybe a)
shedRecording AdmissionObservers
obs = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (AdmissionObservers -> IO ()
onShed AdmissionObservers
obs) m () -> Maybe a -> m (Maybe a)
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> Maybe a
forall a. Maybe a
Nothing

-- A blocked STM retry stays interruptible under the mask, so a cancellation aborts it taking
-- nothing while a committed acquire returns with the weight held and exceptions still masked.
{-# INLINE queuedWait #-}
queuedWait ::
    (MonadUnliftIO m) =>
    AdmissionObservers ->
    WeightedAdmission ->
    Int ->
    (m a -> m a) ->
    m a ->
    m (Maybe a)
queuedWait :: forall (m :: * -> *) a.
MonadUnliftIO m =>
AdmissionObservers
-> WeightedAdmission -> Int -> (m a -> m a) -> m a -> m (Maybe a)
queuedWait AdmissionObservers
obs WeightedAdmission
wa Int
weight m a -> m a
restore m a
action = do
    deadline <- IO (TVar Bool) -> m (TVar Bool)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Int -> IO (TVar Bool)
registerDelay (WeightedAdmission -> Int
waWaitMicros WeightedAdmission
wa))
    acquired <-
        atomically (acquireOrExpire wa weight deadline)
            `UE.finally` atomically (modifyTVar' (waWaiting wa) (subtract 1))
    if acquired
        then admittedRun obs wa weight (onQueued obs) restore action
        else shedRecording obs

-- The gauge increment runs under the enclosing mask, before 'restore', and 'afterArm' with it, so a
-- cancellation or a throwing observer cannot fire the finaliser's decrement without the increment.
{-# INLINE admittedRun #-}
admittedRun ::
    (MonadUnliftIO m) =>
    AdmissionObservers ->
    WeightedAdmission ->
    Int ->
    IO () ->
    (m a -> m a) ->
    m a ->
    m (Maybe a)
admittedRun :: forall (m :: * -> *) a.
MonadUnliftIO m =>
AdmissionObservers
-> WeightedAdmission
-> Int
-> IO ()
-> (m a -> m a)
-> m a
-> m (Maybe a)
admittedRun AdmissionObservers
obs WeightedAdmission
wa Int
weight IO ()
afterArm m a -> m a
restore m a
action =
    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
<$> ( (IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (AdmissionObservers -> Int -> IO ()
onInFlightDelta AdmissionObservers
obs Int
weight IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> IO ()
afterArm) 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
>> m a -> m a
restore m a
action)
                m a -> m () -> m a
forall (m :: * -> *) a b. MonadUnliftIO m => m a -> m b -> m a
`UE.finally` AdmissionObservers -> WeightedAdmission -> Int -> m ()
forall (m :: * -> *).
MonadUnliftIO m =>
AdmissionObservers -> WeightedAdmission -> Int -> m ()
releaseWeight AdmissionObservers
obs WeightedAdmission
wa Int
weight
            )

-- Publish the gauge decrement before returning capacity, or a woken waiter increments while the departing
-- holder is still observable. The STM release is the finaliser, so a throwing observer cannot leak capacity.
{-# INLINE releaseWeight #-}
releaseWeight :: (MonadUnliftIO m) => AdmissionObservers -> WeightedAdmission -> Int -> m ()
releaseWeight :: forall (m :: * -> *).
MonadUnliftIO m =>
AdmissionObservers -> WeightedAdmission -> Int -> m ()
releaseWeight AdmissionObservers
obs WeightedAdmission
wa Int
weight =
    IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (AdmissionObservers -> Int -> IO ()
onInFlightDelta AdmissionObservers
obs (Int -> Int
forall a. Num a => a -> a
negate Int
weight))
        m () -> m () -> m ()
forall (m :: * -> *) a b. MonadUnliftIO m => m a -> m b -> m a
`UE.finally` STM () -> m ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (TVar Int -> (Int -> Int) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' (WeightedAdmission -> TVar Int
waAvailable WeightedAdmission
wa) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
weight))