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

{- | The progress watchdog a registry transfer runs under. It counts the request-body bytes
handed to the connection and the response-body bytes read from it, and the time spent blocked on
the upstream in either direction. A transfer that waits a whole 'ProgressFloor' window without
moving the floor's bytes is stopped. The consumer's own work between reads, and the wait for the
status line and headers, never count.
-}
module Ecluse.Core.Registry.Progress (
    -- * The watchdog
    Watch,
    watched,
    watchedRaising,

    -- * Metered transfers
    meteredReader,
    meteredUpload,
) where

import Control.Exception (throwTo)
import Data.ByteString qualified as BS
import Data.ByteString.Builder (toLazyByteString)
import Data.ByteString.Lazy qualified as LBS
import GHC.Clock (getMonotonicTimeNSec)
import Network.HTTP.Client (BodyReader, GivesPopper, Popper, RequestBody (..))
import UnliftIO (try, withAsync)
import UnliftIO.Concurrent (ThreadId, myThreadId, threadDelay)

import Ecluse.Core.Security (ProgressFloor, floorMinBytes, floorWindowMicros)

-- | One transfer's progress in the current window, shared by its meters and its watchdog.
data Watch = Watch ProgressFloor (IORef Window)

-- When the current wait on the upstream began, the nanoseconds waited before it, and the bytes moved.
data Window = Window (Maybe Word64) Word64 Int

-- Raised in a transfer's thread as a synchronous failure of the transfer, as a socket error would be.
data BelowProgressFloor = BelowProgressFloor
    deriving stock (Int -> BelowProgressFloor -> ShowS
[BelowProgressFloor] -> ShowS
BelowProgressFloor -> String
(Int -> BelowProgressFloor -> ShowS)
-> (BelowProgressFloor -> String)
-> ([BelowProgressFloor] -> ShowS)
-> Show BelowProgressFloor
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> BelowProgressFloor -> ShowS
showsPrec :: Int -> BelowProgressFloor -> ShowS
$cshow :: BelowProgressFloor -> String
show :: BelowProgressFloor -> String
$cshowList :: [BelowProgressFloor] -> ShowS
showList :: [BelowProgressFloor] -> ShowS
Show)

instance Exception BelowProgressFloor

{- | Run a transfer under the watchdog. 'Nothing' means the transfer fell below the floor and was
stopped. The catch sits outside the watchdog's lifetime, so a raise can never land after it.
-}
watched :: ProgressFloor -> (Watch -> IO a) -> IO (Maybe a)
watched :: forall a. ProgressFloor -> (Watch -> IO a) -> IO (Maybe a)
watched ProgressFloor
progress Watch -> IO a
transfer =
    IO a -> IO (Either BelowProgressFloor a)
forall (m :: * -> *) e a.
(MonadUnliftIO m, Exception e) =>
m a -> m (Either e a)
try (ProgressFloor -> (Watch -> IO a) -> IO a
forall a. ProgressFloor -> (Watch -> IO a) -> IO a
watchedRaising ProgressFloor
progress Watch -> IO a
transfer) IO (Either BelowProgressFloor a)
-> (Either BelowProgressFloor a -> Maybe a) -> IO (Maybe a)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \case
        Left BelowProgressFloor
BelowProgressFloor -> Maybe a
forall a. Maybe a
Nothing
        Right a
result -> a -> Maybe a
forall a. a -> Maybe a
Just a
result

{- | 'watched' for a transfer already committed to a client, which a floor miss can only abort. The
failure propagates as an exception, so the response tears down instead of ending cleanly.
-}
watchedRaising :: ProgressFloor -> (Watch -> IO a) -> IO a
watchedRaising :: forall a. ProgressFloor -> (Watch -> IO a) -> IO a
watchedRaising ProgressFloor
progress Watch -> IO a
transfer = do
    owner <- IO ThreadId
forall (m :: * -> *). MonadIO m => m ThreadId
myThreadId
    watch <- Watch progress <$> newIORef (Window Nothing 0 0)
    withAsync (watchdog watch owner) (\Async ()
_ -> Watch -> IO a
transfer Watch
watch)

{- Base 'throwTo' delivers the plain exception, which the synchronous catch in 'watched' sees.
UnliftIO's would wrap it as asynchronous. -}
watchdog :: Watch -> ThreadId -> IO ()
watchdog :: Watch -> ThreadId -> IO ()
watchdog (Watch ProgressFloor
progress IORef Window
window) ThreadId
owner = IO ()
go
  where
    budget :: Word64
budget = Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ProgressFloor -> Int
floorWindowMicros ProgressFloor
progress) Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
* Word64
1_000
    go :: IO ()
go = do
        now <- IO Word64
getMonotonicTimeNSec
        current <- readIORef window
        let spent = Word64 -> Window -> Word64
waitedBy Word64
now Window
current
        if spent >= budget
            then throwTo owner BelowProgressFloor
            else threadDelay (fromIntegral (pause budget spent current `div` 1_000) + 1) >> go

{- With a wait open the watchdog wakes as its budget runs out. With none open the waited time cannot
grow, so it sleeps at least a sixteenth of the window instead of spinning on a near-spent budget. -}
pause :: Word64 -> Word64 -> Window -> Word64
pause :: Word64 -> Word64 -> Window -> Word64
pause Word64
budget Word64
spent = \case
    Window (Just Word64
_) Word64
_ Int
_ -> Word64
budget Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
spent
    Window Maybe Word64
Nothing Word64
_ Int
_ -> Word64 -> Word64 -> Word64
forall a. Ord a => a -> a -> a
max (Word64
budget Word64 -> Word64 -> Word64
forall a. Integral a => a -> a -> a
`div` Word64
16) (Word64
budget Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
spent)

-- | A reader whose every read counts as a wait on the upstream, and its chunk as progress.
meteredReader :: Watch -> BodyReader -> BodyReader
meteredReader :: Watch -> BodyReader -> BodyReader
meteredReader Watch
watch BodyReader
readChunk = do
    Watch -> IO ()
beginWait Watch
watch
    chunk <- BodyReader
readChunk
    endWait watch (BS.length chunk)
    pure chunk

{- | The same body, handed over in slices the watch counts once each is written. A body with no
bytes has no upload to watch and stays as it is.
-}
meteredUpload :: Watch -> RequestBody -> RequestBody
meteredUpload :: Watch -> RequestBody -> RequestBody
meteredUpload Watch
watch RequestBody
body = case RequestBody
body of
    RequestBodyLBS ByteString
bytes | Bool -> Bool
not (ByteString -> Bool
LBS.null ByteString
bytes) -> Int64 -> [ByteString] -> RequestBody
fromChunks (ByteString -> Int64
LBS.length ByteString
bytes) (ByteString -> [ByteString]
LBS.toChunks ByteString
bytes)
    RequestBodyBS ByteString
bytes | Bool -> Bool
not (ByteString -> Bool
BS.null ByteString
bytes) -> Int64 -> [ByteString] -> RequestBody
fromChunks (Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Int
BS.length ByteString
bytes)) [ByteString
bytes]
    RequestBodyBuilder Int64
size Builder
builder | Int64
size Int64 -> Int64 -> Bool
forall a. Ord a => a -> a -> Bool
> Int64
0 -> Int64 -> [ByteString] -> RequestBody
fromChunks Int64
size (ByteString -> [ByteString]
LBS.toChunks (Builder -> ByteString
toLazyByteString Builder
builder))
    RequestBodyStream Int64
size GivesPopper ()
gives -> Int64 -> GivesPopper () -> RequestBody
RequestBodyStream Int64
size (Watch -> GivesPopper () -> GivesPopper ()
meteredGives Watch
watch GivesPopper ()
gives)
    RequestBodyStreamChunked GivesPopper ()
gives -> GivesPopper () -> RequestBody
RequestBodyStreamChunked (Watch -> GivesPopper () -> GivesPopper ()
meteredGives Watch
watch GivesPopper ()
gives)
    RequestBodyIO IO RequestBody
io -> IO RequestBody -> RequestBody
RequestBodyIO (Watch -> RequestBody -> RequestBody
meteredUpload Watch
watch (RequestBody -> RequestBody) -> IO RequestBody -> IO RequestBody
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO RequestBody
io)
    RequestBody
_ -> RequestBody
body
  where
    fromChunks :: Int64 -> [ByteString] -> RequestBody
fromChunks Int64
size [ByteString]
chunks = Int64 -> GivesPopper () -> RequestBody
RequestBodyStream Int64
size (Watch -> GivesPopper () -> GivesPopper ()
meteredGives Watch
watch ([ByteString] -> GivesPopper ()
givesChunks [ByteString]
chunks))

givesChunks :: [ByteString] -> GivesPopper ()
givesChunks :: [ByteString] -> GivesPopper ()
givesChunks [ByteString]
chunks NeedsPopper ()
needsPopper = do
    remaining <- [ByteString] -> IO (IORef [ByteString])
forall (m :: * -> *) a. MonadIO m => a -> m (IORef a)
newIORef [ByteString]
chunks
    needsPopper . atomicModifyIORef' remaining $ \case
        [] -> ([], ByteString
BS.empty)
        ByteString
chunk : [ByteString]
rest -> ([ByteString]
rest, ByteString
chunk)

{- The connection asks for the next slice only once it has written the last, so each call closes
the wait that slice opened. Slicing keeps a large chunk from counting only when it all lands. -}
meteredGives :: Watch -> GivesPopper () -> GivesPopper ()
meteredGives :: Watch -> GivesPopper () -> GivesPopper ()
meteredGives Watch
watch GivesPopper ()
gives NeedsPopper ()
needsPopper = GivesPopper ()
gives GivesPopper () -> GivesPopper ()
forall a b. (a -> b) -> a -> b
$ \BodyReader
popper -> do
    handed <- Int -> IO (IORef Int)
forall (m :: * -> *) a. MonadIO m => a -> m (IORef a)
newIORef Int
0
    pending <- newIORef BS.empty
    needsPopper $ do
        readIORef handed >>= \Int
size -> Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
size Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (Watch -> Int -> IO ()
endWait Watch
watch Int
size)
        slice <- nextSlice pending popper
        writeIORef handed (BS.length slice)
        unless (BS.null slice) (beginWait watch)
        pure slice

nextSlice :: IORef ByteString -> Popper -> IO ByteString
nextSlice :: IORef ByteString -> BodyReader -> BodyReader
nextSlice IORef ByteString
pending BodyReader
popper = do
    held <- IORef ByteString -> BodyReader
forall (m :: * -> *) a. MonadIO m => IORef a -> m a
readIORef IORef ByteString
pending
    chunk <- if BS.null held then popper else pure held
    let (slice, rest) = BS.splitAt uploadSliceBytes chunk
    writeIORef pending rest
    pure slice

uploadSliceBytes :: Int
uploadSliceBytes :: Int
uploadSliceBytes = Int
64 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1024

beginWait :: Watch -> IO ()
beginWait :: Watch -> IO ()
beginWait (Watch ProgressFloor
_ IORef Window
window) = do
    now <- IO Word64
getMonotonicTimeNSec
    atomicModifyIORef' window (\(Window Maybe Word64
_ Word64
waited Int
moved) -> (Maybe Word64 -> Word64 -> Int -> Window
Window (Word64 -> Maybe Word64
forall a. a -> Maybe a
Just Word64
now) Word64
waited Int
moved, ()))

endWait :: Watch -> Int -> IO ()
endWait :: Watch -> Int -> IO ()
endWait (Watch ProgressFloor
progress IORef Window
window) Int
bytes = do
    now <- IO Word64
getMonotonicTimeNSec
    atomicModifyIORef' window $ \current :: Window
current@(Window Maybe Word64
_ Word64
_ Int
moved) ->
        if Int
moved Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
bytes Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= ProgressFloor -> Int
floorMinBytes ProgressFloor
progress
            then (Maybe Word64 -> Word64 -> Int -> Window
Window Maybe Word64
forall a. Maybe a
Nothing Word64
0 Int
0, ())
            else (Maybe Word64 -> Word64 -> Int -> Window
Window Maybe Word64
forall a. Maybe a
Nothing (Word64 -> Window -> Word64
waitedBy Word64
now Window
current) (Int
moved Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
bytes), ())

-- The watchdog can read its clock just before the transfer stamps a later wait, so the gap floors at zero.
waitedBy :: Word64 -> Window -> Word64
waitedBy :: Word64 -> Window -> Word64
waitedBy Word64
now (Window Maybe Word64
since Word64
waited Int
_) = Word64
waited Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64 -> (Word64 -> Word64) -> Maybe Word64 -> Word64
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Word64
0 (\Word64
began -> if Word64
now Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
> Word64
began then Word64
now Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
began else Word64
0) Maybe Word64
since