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

{- | The memory sampler: a background loop that reads the collector's statistics and the cgroup's
memory use, and steers the meter's budget through "Ecluse.Core.Server.Admission.Brake".

It runs off the served path. Without @-T@ it still reads the cgroup, and without a cgroup limit it
still reads the collector, so it never has a reason to stop.
-}
module Ecluse.Rts.Sampler (
    SamplerSettings (..),
    samplerPeriodMicros,
    shareWindowPeriods,
    readCollector,
    runMemorySampler,
) where

import Control.Concurrent (threadDelay)
import GHC.Stats (GCDetails (gcdetails_live_bytes), RTSStats (cpu_ns, gc, gc_cpu_ns, major_gcs), getRTSStats, getRTSStatsEnabled)

import Ecluse.Core.Server.Admission.Brake (
    BrakeMarks,
    BrakeState (bsBudget, bsLevel),
    CollectorReading (..),
    brakeStep,
    initialBrakeState,
    newSampleWindow,
    windowSample,
 )
import Ecluse.Core.Server.Admission.Meter (MemoryMeter, meterSnapshot, steerMeter, takeLargestCharge)
import Ecluse.Core.Server.Admission.Types (BrakeBounds, MeterSnapshot (snChargedBytes))

-- | What one sampler loop reads and steers.
data SamplerSettings = SamplerSettings
    { SamplerSettings -> MemoryMeter
ssMeter :: MemoryMeter
    , SamplerSettings -> BrakeMarks
ssMarks :: BrakeMarks
    , SamplerSettings -> BrakeBounds
ssBounds :: BrakeBounds
    , SamplerSettings -> IO (Maybe CollectorReading)
ssCollector :: IO (Maybe CollectorReading)
    -- ^ The collector's counters, absent without @-T@.
    , SamplerSettings -> IO (Maybe Int)
ssKernelPermille :: IO (Maybe Int)
    -- ^ The cgroup's non-reclaimable memory use against its limit, when a limit binds.
    , SamplerSettings -> Int
ssPeriodMicros :: Int
    , SamplerSettings -> Int
ssWindowPeriods :: Int
    -- ^ How many periods the collector's CPU share averages over.
    }

-- | The sampling period: ten readings a second.
samplerPeriodMicros :: Int
samplerPeriodMicros :: Int
samplerPeriodMicros = Int
100_000

-- | The collector's CPU share averages over one second at the default period.
shareWindowPeriods :: Int
shareWindowPeriods :: Int
shareWindowPeriods = Int
10

-- | A reader for the collector's counters, or one that always reads nothing when @-T@ is off.
readCollector :: IO (IO (Maybe CollectorReading))
readCollector :: IO (IO (Maybe CollectorReading))
readCollector = do
    statsOn <- IO Bool
getRTSStatsEnabled
    pure (if statsOn then Just . collectorReading <$> getRTSStats else pure Nothing)
  where
    collectorReading :: RTSStats -> CollectorReading
collectorReading RTSStats
stats =
        CollectorReading
            { crCpuNs :: Int64
crCpuNs = RTSStats -> Int64
cpu_ns RTSStats
stats
            , crGcCpuNs :: Int64
crGcCpuNs = RTSStats -> Int64
gc_cpu_ns RTSStats
stats
            , crMajorCollections :: Word32
crMajorCollections = RTSStats -> Word32
major_gcs RTSStats
stats
            , crLiveBytes :: Int
crLiveBytes = Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (GCDetails -> Word64
gcdetails_live_bytes (RTSStats -> GCDetails
gc RTSStats
stats))
            }

-- | Sample and steer forever. The caller runs it under supervision and cancels it at shutdown.
runMemorySampler :: SamplerSettings -> IO ()
runMemorySampler :: SamplerSettings -> IO ()
runMemorySampler SamplerSettings
settings =
    BrakeState -> SampleWindow -> IO ()
forall {b}. BrakeState -> SampleWindow -> IO b
loop (BrakeBounds -> BrakeState
initialBrakeState (SamplerSettings -> BrakeBounds
ssBounds SamplerSettings
settings)) (Int -> SampleWindow
newSampleWindow (SamplerSettings -> Int
ssWindowPeriods SamplerSettings
settings))
  where
    loop :: BrakeState -> SampleWindow -> IO b
loop BrakeState
current SampleWindow
window = do
        Int -> IO ()
threadDelay (SamplerSettings -> Int
ssPeriodMicros SamplerSettings
settings)
        reading <- SamplerSettings -> IO (Maybe CollectorReading)
ssCollector SamplerSettings
settings
        charged <- snChargedBytes <$> meterSnapshot (ssMeter settings)
        largest <- takeLargestCharge (ssMeter settings)
        kernel <- ssKernelPermille settings
        let (sample, window') = windowSample window reading charged largest kernel
            next = BrakeMarks -> BrakeBounds -> BrakeState -> GcSample -> BrakeState
brakeStep (SamplerSettings -> BrakeMarks
ssMarks SamplerSettings
settings) (SamplerSettings -> BrakeBounds
ssBounds SamplerSettings
settings) BrakeState
current GcSample
sample
        when (bsBudget next /= bsBudget current || bsLevel next /= bsLevel current) $
            steerMeter (ssMeter settings) (bsBudget next) (bsLevel next)
        loop next window'