diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,48 @@ # Changelog +## 0.11.0.0 — 2026-10-25 + +### Breaking Changes + +- `MetricsServerConfig` gains `enableControl`, `corsAllowedOrigins`, + `wsMaxQueuedFrames`, `wsMaxQueuedBytes`, and `wsSendTimeoutUs`; direct construction + must supply them. `MetricsServer` gains the required `serverShutdown` action. +- `ServerMessage` gains `ServerError`; exhaustive matches must handle it. + `ProcessorHealth` gains `paused`; direct construction must supply the count. +- Require `shibuya-core ^>=0.11.0.0`, including its extended processor metrics and + paused-state API. + +### New Features + +- Add opt-in control discovery and pause/resume POST routes with structured errors, + draining refusal, and actual-request Origin enforcement. Control is disabled by + default. Export `Shibuya.Metrics.Control`. +- Expose optional latency, oldest in-flight, and bounded acknowledged progress in + JSON and WebSocket metrics. Add the processing-duration Prometheus histogram and + optional oldest-in-flight gauge; report paused health counts and state value 5. +- Bound queued WebSocket frames and encoded bytes with independent sampling and a + single application writer. Recover dropped deltas with an adjacent + `queue_overflow` error and replacement snapshot; bound sends and shutdown flushes. +- Export the pure `Shibuya.Metrics.Outbound` queue model and optional + `WebSocketHooks` instrumentation. Subscription changes establish ordered snapshots. +- Send in-band WebSocket errors for malformed messages and before subscription-limit + policy closes; expose stable `code` members in HTTP error responses. +- Add configurable HTTP CORS and preflight handling, and validate WebSocket upgrade + Origins before acquiring connection slots. Export `Shibuya.Metrics.Cors` and + `Shibuya.Metrics.Error` with reusable middleware and error helpers. + +### Bug Fixes + +- Suspend deltas after terminal-frame admission evicts updates so the next update + establishes a replacement snapshot. + +### Other Changes + +- Bound the observer-load executable's `streamly-core` dependency to `^>=0.3`, + matching the range already required by `shibuya-core`. +- Avoid waking WebSocket writers for empty terminal batches and make observer + captures reproducible with bounded delivery acceptance evidence. + ## 0.10.0.1 — 2026-10-05 ### Other Changes diff --git a/bench/ObserverLoad.hs b/bench/ObserverLoad.hs new file mode 100644 --- /dev/null +++ b/bench/ObserverLoad.hs @@ -0,0 +1,154 @@ +-- | Sustained real WebSocket observers around the production serial runner. +-- Copy this harness and its Cabal stanza unchanged to the baseline checkout. +module Main (main) where + +import Control.Concurrent (threadDelay) +import Control.Concurrent.Async (async, cancel) +import Control.Concurrent.NQE.Supervisor (Strategy (IgnoreAll)) +import Control.Concurrent.STM (atomically, check, modifyTVar', newTVarIO, readTVar, writeTVar) +import Control.Exception (bracket, catch, evaluate, throwIO) +import Control.Monad (forM, forever, unless, when) +import Data.Aeson (encode, object, (.=)) +import Data.ByteString.Lazy qualified as LBS +import Data.IORef (atomicModifyIORef', newIORef, readIORef) +import Data.List (sort) +import Data.Text qualified as Text +import Data.Word (Word64) +import Effectful (liftIO, runEff) +import GHC.Clock (getMonotonicTimeNSec) +import GHC.Stats (GCDetails (..), RTSStats (..), getRTSStats) +import Network.Wai.Handler.Warp qualified as Warp +import Network.WebSockets qualified as WS +import Shibuya.Adapter (Adapter (..)) +import Shibuya.Core.Ack (AckDecision (AckOk)) +import Shibuya.Core.AckHandle (AckHandle (..)) +import Shibuya.Core.Ingested (mkIngested) +import Shibuya.Core.Metrics (ProcessorId (..)) +import Shibuya.Core.Types (MessageId (..), mkEnvelope) +import Shibuya.Internal.Runner.Master (markMasterRunning, startMaster, stopMaster) +import Shibuya.Internal.Runner.Supervised (isDone, runSupervised) +import Shibuya.Metrics.Config (MetricsServerConfig (..), defaultConfig) +import Shibuya.Metrics.Server (combinedApp) +import Shibuya.Metrics.WebSocket (WebSocketState (..), newWebSocketState, shutdownWebSockets) +import Shibuya.Policy (Concurrency (Serial), OrderingPolicy (Unordered)) +import Shibuya.Telemetry.Effect (runTracingNoop) +import Streamly.Data.Stream qualified as Stream +import System.Environment (lookupEnv) +import System.Mem (performMajorGC) +import System.Timeout (timeout) +import Text.Read (readMaybe) + +main :: IO () +main = do + clients <- envInt "CLIENTS" 16 + count <- envInt "MESSAGES" 2_000_000 + output <- maybe "observer-load.json" id <$> lookupEnv "OUTPUT_JSON" + report <- bracket acquire (runEff . stopMaster) $ \master -> do + ready <- newTVarIO False + acknowledged <- newIORef ([] :: [Word64]) + frames <- newIORef (0 :: Int) + let source = + Stream.mapM + ( \n -> do + liftIO $ atomically $ readTVar ready >>= check + started <- liftIO getMonotonicTimeNSec + let envelope = mkEnvelope (MessageId $ Text.pack $ show n) n + ack = AckHandle $ \_ -> liftIO $ do + end <- getMonotonicTimeNSec + atomicModifyIORef' acknowledged (\xs -> ((end - started) : xs, ())) + pure $ mkIngested envelope ack + ) + $ Stream.fromList [1 .. count] + adapter = Adapter "observer-load" source (pure ()) + processor <- + runEff $ + runTracingNoop $ + runSupervised master 16 (ProcessorId "observer-load") Unordered Serial adapter (const $ pure AckOk) + wsState <- newWebSocketState defaultConfig.wsMaxConnections + let config = defaultConfig {wsPushIntervalUs = 10_000} + Warp.testWithApplication (pure $ combinedApp config master wsState []) $ \port -> do + connected <- newTVarIO (0 :: Int) + let observer index = + WS.runClient + "127.0.0.1" + port + "/ws" + ( \conn -> do + _ <- WS.receiveData conn :: IO LBS.ByteString + atomically $ modifyTVar' connected (+ 1) + -- A quarter of observers read ten times slower than the producer. + forever $ do + when ((index - 1) `mod` 4 == 0) $ threadDelay 100_000 + bytes <- WS.receiveData conn :: IO LBS.ByteString + _ <- evaluate $ LBS.length bytes + atomicModifyIORef' frames (\n -> (n + 1, ())) + ) + `catch` normalClose + bracket (forM [1 .. clients] $ async . observer) (mapM_ cancel) $ \observers -> do + connectedAll <- timeout 5_000_000 $ atomically $ readTVar connected >>= check . (== clients) + unless (connectedAll == Just ()) $ fail "observers did not connect" + performMajorGC + before <- getRTSStats + start <- getMonotonicTimeNSec + atomically $ writeTVar ready True + let awaitDone = runEff (isDone processor) >>= \done -> unless done (threadDelay 1_000 >> awaitDone) + completed <- timeout 30_000_000 awaitDone + unless (completed == Just ()) $ fail "production runner did not drain" + end <- getMonotonicTimeNSec + -- Keep the latency sample list live through the GC, identically on both builds. + samples <- readIORef acknowledged + _ <- evaluate $ length samples + performMajorGC + after <- getRTSStats + shutdownStart <- getMonotonicTimeNSec + shutdownWebSockets wsState + released <- timeout 6_000_000 $ atomically $ readTVar wsState.connectionCount >>= check . (== 0) + shutdownEnd <- getMonotonicTimeNSec + unless (released == Just ()) $ fail "connection slots were not restored" + -- Cancel and join slow readers after server slots are released; buffered + -- client data is not part of the server shutdown latency measurement. + mapM_ cancel observers + actual <- evaluate $ length samples + unless (actual == count) $ fail "missing acknowledgements" + delivered <- readIORef frames + pure $ + object + [ "schemaVersion" .= (1 :: Int), + "clients" .= clients, + "slowReaders" .= ((clients + 3) `div` 4), + "elapsedSeconds" .= (fromIntegral (end - start) / 1e9 :: Double), + "messages" .= count, + "completed" .= actual, + "observerFrames" .= delivered, + "throughputPerSecond" .= (fromIntegral count * 1e9 / fromIntegral (end - start) :: Double), + "latencyP95Ms" .= percentile 0.95 samples, + "latencyP99Ms" .= percentile 0.99 samples, + "allocatedBytesPerMessage" .= (fromIntegral (after.allocated_bytes - before.allocated_bytes) / fromIntegral count :: Double), + "maxLiveBytes" .= after.max_live_bytes, + "liveBytes" .= after.gc.gcdetails_live_bytes, + "shutdownLatencyMs" .= (fromIntegral (shutdownEnd - shutdownStart) / 1e6 :: Double), + "finalWebSocketConnections" .= (0 :: Int) + ] + LBS.writeFile output (encode report <> "\n") + LBS.putStr $ encode report <> "\n" + where + acquire = runEff $ do + master <- startMaster IgnoreAll + markMasterRunning master + pure master + +normalClose :: WS.ConnectionException -> IO () +normalClose = \case + WS.ConnectionClosed -> pure () + WS.CloseRequest _ _ -> pure () + other -> throwIO other + +envInt :: String -> Int -> IO Int +envInt key fallback = maybe fallback id . (>>= readMaybe) <$> lookupEnv key + +percentile :: Double -> [Word64] -> Double +percentile _ [] = 0 +percentile fraction samples = + let ordered = sort samples + index = min (length ordered - 1) (ceiling (fraction * fromIntegral (length ordered)) - 1) + in fromIntegral (ordered !! max 0 index) / 1e6 diff --git a/bench/SaturationLoad.hs b/bench/SaturationLoad.hs new file mode 100644 --- /dev/null +++ b/bench/SaturationLoad.hs @@ -0,0 +1,119 @@ +-- | Deterministic saturation with a real connection and a blocked writer hook. +module Main (main) where + +import Control.Concurrent (threadDelay) +import Control.Concurrent.NQE.Supervisor (Strategy (IgnoreAll)) +import Control.Concurrent.STM + ( atomically, + check, + newTQueueIO, + newTVarIO, + orElse, + readTQueue, + readTVar, + writeTQueue, + writeTVar, + ) +import Control.Exception (bracket) +import Control.Monad (replicateM, when) +import Data.Aeson (encode, object, (.=)) +import Data.ByteString.Lazy qualified as LBS +import Data.Text qualified as Text +import Data.Time.Clock (getCurrentTime) +import Effectful (runEff) +import GHC.Clock (getMonotonicTimeNSec) +import GHC.Stats (GCDetails (..), RTSStats (..), getRTSStats) +import Network.Wai.Handler.Warp qualified as Warp +import Network.Wai.Handler.WebSockets qualified as WaiWS +import Network.WebSockets qualified as WS +import Shibuya.Core.Metrics (ProcessorId (..), incrementReceived, newMetricsHandle) +import Shibuya.Internal.Runner.Master (markMasterRunning, registerProcessor, startMaster, stopMaster) +import Shibuya.Metrics.Config (MetricsServerConfig (..), defaultConfig) +import Shibuya.Metrics.Outbound qualified as Out +import Shibuya.Metrics.Server (combinedApp) +import Shibuya.Metrics.WebSocket + ( WebSocketHooks (..), + WebSocketState (..), + defaultWebSocketHooks, + newWebSocketState, + websocketAppWithHooks, + ) +import System.Environment (lookupEnv) +import System.Mem (performMajorGC) +import System.Timeout (timeout) + +main :: IO () +main = bracket acquire (runEff . stopMaster) $ \master -> do + handle <- getCurrentTime >>= newMetricsHandle + runEff $ registerProcessor master (ProcessorId $ Text.replicate 8192 "p") handle + gate <- newTVarIO False + entered <- newTQueueIO + ticks <- newTQueueIO + observed <- newTQueueIO + state <- newWebSocketState 1 + let config = defaultConfig {wsMaxQueuedBytes = 65_536, wsSendTimeoutUs = 1_000_000} + hooks = + defaultWebSocketHooks + { beforeWrite = \_ -> do + atomically $ readTVar gate >>= \blocked -> when blocked (writeTQueue entered ()) + atomically $ readTVar gate >>= check . not, + afterProduce = atomically . writeTQueue observed, + waitForPush = \_ wsState -> + atomically $ + (readTVar wsState.shutdownRequested >>= \stopped -> check stopped >> pure True) + `orElse` (readTQueue ticks >> pure False) + } + app = + WaiWS.websocketsOr + WS.defaultConnectionOptions + (websocketAppWithHooks hooks config master state) + (combinedApp config {enableWebSocket = False} master state []) + tick = + incrementReceived handle + >> atomically (writeTQueue ticks ()) + >> timeout 50_000 (atomically $ readTQueue observed) + fill n previous = do + result <- tick + case result of + Nothing -> pure (n, previous) + Just q -> do + when (Out.queuedFrames q > config.wsMaxQueuedFrames || Out.queuedBytes q > config.wsMaxQueuedBytes) $ + fail "queue exceeded a budget" + if n >= (500 :: Int) then fail "required saturation never stopped production" else fill (n + 1) q + report <- Warp.testWithApplication (pure app) $ \port -> + WS.runClient "127.0.0.1" port "/ws" $ \conn -> do + _ <- WS.receiveData conn :: IO LBS.ByteString + atomically $ writeTVar gate True + _ <- tick + blocked <- timeout 500_000 $ atomically $ readTQueue entered + when (blocked /= Just ()) $ fail "writer was not blocked" + (count, q) <- fill 1 Out.empty + heap <- replicateM 4 $ do + performMajorGC + stats <- getRTSStats + threadDelay 20_000 + pure stats.gc.gcdetails_live_bytes + started <- getMonotonicTimeNSec + released <- timeout 1_500_000 $ atomically $ readTVar state.connectionCount >>= check . (== 0) + ended <- getMonotonicTimeNSec + when (released /= Just ()) $ fail "stalled connection slot was not released" + pure $ + object + [ "schemaVersion" .= (1 :: Int), + "producerTicksBeforeSaturation" .= count, + "frameBudget" .= config.wsMaxQueuedFrames, + "byteBudget" .= config.wsMaxQueuedBytes, + "retainedFramesAtSaturation" .= Out.queuedFrames q, + "retainedEncodedBytesAtSaturation" .= Out.queuedBytes q, + "liveHeapSamplesAfterSaturation" .= heap, + "slotRecoveryMicros" .= ((ended - started) `div` 1000), + "finalWebSocketConnections" .= (0 :: Int) + ] + path <- maybe "saturation-load.json" id <$> lookupEnv "OUTPUT_JSON" + LBS.writeFile path $ encode report <> "\n" + LBS.putStr $ encode report <> "\n" + where + acquire = runEff $ do + master <- startMaster IgnoreAll + markMasterRunning master + pure master diff --git a/shibuya-metrics.cabal b/shibuya-metrics.cabal --- a/shibuya-metrics.cabal +++ b/shibuya-metrics.cabal @@ -1,6 +1,6 @@ cabal-version: 3.12 name: shibuya-metrics -version: 0.10.0.1 +version: 0.11.0.0 synopsis: Metrics web server for Shibuya queue processing framework description: Provides HTTP/JSON, Prometheus, and WebSocket endpoints for @@ -22,8 +22,12 @@ exposed-modules: Shibuya.Metrics Shibuya.Metrics.Config + Shibuya.Metrics.Control + Shibuya.Metrics.Cors + Shibuya.Metrics.Error Shibuya.Metrics.Health Shibuya.Metrics.JSON + Shibuya.Metrics.Outbound Shibuya.Metrics.Prometheus Shibuya.Metrics.Server Shibuya.Metrics.Types @@ -49,7 +53,7 @@ containers ^>=0.7, http-types ^>=0.12, prometheus-client ^>=1.1, - shibuya-core ^>=0.10.0.1, + shibuya-core ^>=0.11.0.0, stm ^>=2.5, text ^>=2.1, time ^>=1.14, @@ -83,8 +87,11 @@ other-modules: Paths_shibuya_metrics + Shibuya.Metrics.ControlSpec + Shibuya.Metrics.CorsSpec Shibuya.Metrics.HealthSpec Shibuya.Metrics.JSONSpec + Shibuya.Metrics.OutboundQueueSpec Shibuya.Metrics.PrometheusSpec Shibuya.Metrics.ServerSpec Shibuya.Metrics.TestSupport @@ -103,15 +110,18 @@ containers, effectful >=2.6.1 && <2.8, hspec ^>=2.11.17, + http-client ^>=0.7.19, http-types, nqe ^>=0.6, - shibuya-core, + shibuya-core ^>=0.11.0.0, shibuya-metrics, stm, + streamly-core, text, time, wai, wai-extra ^>=3.1.18, + wai-websockets, warp, websockets, @@ -148,5 +158,74 @@ shibuya-metrics, stm, time, + warp, + websockets, + +-- Sustained real observers while the production runner processes messages. +executable metrics-observer-load + import: warnings + default-language: GHC2024 + hs-source-dirs: bench + main-is: ObserverLoad.hs + default-extensions: + DerivingStrategies + LambdaCase + NoFieldSelectors + OverloadedRecordDot + OverloadedStrings + + ghc-options: + -threaded + -rtsopts + "-with-rtsopts=-N4 -T -A32m" + -O2 + + build-depends: + aeson, + async, + base ^>=4.21.0.0, + bytestring, + effectful >=2.6.1 && <2.8, + nqe ^>=0.6, + shibuya-core, + shibuya-metrics, + stm, + streamly-core ^>=0.3, + text, + warp, + websockets, + +-- Writer-stall saturation, retained heap plateau, and slot recovery fixture. +executable metrics-saturation-load + import: warnings + default-language: GHC2024 + hs-source-dirs: bench + main-is: SaturationLoad.hs + default-extensions: + DerivingStrategies + LambdaCase + NoFieldSelectors + OverloadedRecordDot + OverloadedStrings + + ghc-options: + -threaded + -rtsopts + "-with-rtsopts=-N4 -T -A32m" + -O2 + + build-depends: + aeson, + async, + base ^>=4.21.0.0, + bytestring, + effectful >=2.6.1 && <2.8, + nqe ^>=0.6, + shibuya-core, + shibuya-metrics, + stm, + text, + time, + wai-websockets, warp, websockets, diff --git a/src/Shibuya/Metrics.hs b/src/Shibuya/Metrics.hs --- a/src/Shibuya/Metrics.hs +++ b/src/Shibuya/Metrics.hs @@ -6,7 +6,7 @@ -- == Quick Start -- -- @ --- import Shibuya.App (runApp, AppHandle) +-- import Shibuya.App (runApp, AppHandle, getAppMaster) -- import Shibuya.Metrics -- -- main :: IO () @@ -15,7 +15,7 @@ -- Right appHandle <- runEff $ runApp ... -- -- -- Start metrics server on port 9090 --- withMetricsServer defaultConfig appHandle.master $ \\server -> do +-- withMetricsServer defaultConfig (getAppMaster appHandle) $ \\server -> do -- putStrLn $ "Metrics server running on port " ++ show server.serverPort -- -- Your app logic here... -- @ @@ -30,6 +30,49 @@ -- * @GET /metrics/prometheus@ - Prometheus-format metrics -- * @WS /ws@ - WebSocket for real-time updates -- +-- == Latency and in-flight detail +-- +-- Opt in through @AppConfig.metricsDetail = MetricsDetailed defaultDetailOptions@. +-- JSON and WebSocket frames gain optional @latency@ (count, sumSeconds, finite +-- cumulative buckets) and @oldestInFlight@ (messageId, startedAt, ageSeconds). +-- Timing includes finalization retries once per attempt or emitted batch. +-- Prometheus exposes @shibuya_message_processing_seconds@ and the optional +-- @shibuya_processor_oldest_in_flight_seconds@ gauge. Age-only changes do not +-- emit WebSocket deltas: clients extrapolate from their last received frame. +-- HTTP reads and snapshots return freshly sampled age. Basic accounting remains +-- the default and omits these members. +-- +-- == Cross-origin access +-- +-- Set 'corsAllowedOrigins' to explicit serialized browser origins, for example +-- @defaultConfig {corsAllowedOrigins = ["https://ops.example.com"]}@. The empty +-- default preserves read-route headers and WebSocket upgrade behavior. Wildcards, +-- whitespace, paths, queries and fragments are refused by the built-in server. +-- Matching is exact, without case folding or default-port normalization. +-- Allowed HTTP responses (including errors) gain @Access-Control-Allow-Origin@; +-- enabled responses merge @Vary: Origin@ even for absent or refused origins. +-- Preflight @OPTIONS@ requests allow @GET, POST, OPTIONS@ and valid requested +-- header tokens (or @Content-Type@), returning empty 204 responses with a 600-second +-- max age and variation by requested method and headers. Credentials are never granted. +-- Disallowed WebSocket Origins are refused before acquiring a connection slot; +-- headerless clients remain supported. This policy is not authentication. Use a +-- trusted network or an authenticating reverse proxy. Hosts mounting 'combinedApp' +-- directly must enforce the documented configuration constraints themselves. +-- Same-origin reads need no origin list; browser control POSTs must explicitly +-- allow their page origin, even in same-origin deployments. +-- +-- == HTTP errors +-- +-- Error bodies contain a human-readable @error@, a stable @code@, and optional +-- context. Unknown paths return 404 @not_found@; unknown @/metrics/:id@ processors +-- return 404 @processor_not_found@ with @processor@; plain HTTP @/ws@ returns 404 +-- @websocket_upgrade_required@ when WebSockets are enabled. Preflight and upgrade +-- origin refusals return 403 @origin_not_allowed@. Unsupported preflight methods +-- return 403 @cors_method_not_allowed@; invalid requested header lists return 403 +-- @cors_headers_not_allowed@. Existing error members, statuses and content types +-- are preserved, and @code@ is an additive member. 'errorResponse' protects its +-- reserved @error@ and @code@ members from overrides by context. +-- -- == WebSocket Protocol -- -- Clients can send: @@ -46,17 +89,51 @@ -- * @{"type": "terminal", "processor": "id", "status": "stopped"}@ - Processor left the live registry -- * @{"type": "pong"}@ - Response to ping -- * @{"type": "goodbye"}@ - Server shutting down +-- * @{"type": "error", "code": "invalid_message", "message": "..."}@ - Invalid client JSON/tag +-- +-- Error codes are @invalid_message@ (receive continues), @subscription_limit_exceeded@ +-- (error before close 1008), and @queue_overflow@ (immediately followed by a current +-- snapshot). Subscribe and unsubscribe changes establish ordered snapshot boundaries. +-- One writer drains a queue bounded by 'wsMaxQueuedFrames' (64, minimum 2) and +-- 'wsMaxQueuedBytes' (1 MiB). All queued frames count, including snapshots, errors, +-- terminal events and goodbye. Oldest expendable updates/pongs may be dropped; +-- required-frame saturation or oversized snapshots close with 1013. Reconnect and +-- read a fresh snapshot on close; metrics have no replay cursor. +-- 'wsSendTimeoutUs' bounds sends and the whole close flush (5 seconds). Responsive +-- shutdown peers receive goodbye; stalled peers may time out without it. At most +-- one popped frame is outside the queue; transient encoding and the entire process +-- heap are not covered by the encoded-payload budget. Closed decoders must add the +-- new @error@ tag; raw JSON consumers may ignore unknown frame types. +-- See @docs/architecture/METRICS.md@ for the complete shared-convention mapping. +-- Detailed processor metrics also carry optional acknowledged cursor progress: latest, +-- bounded partitions and sticky truncation, with lossless tagged cursor values. +-- Paused intake is reported with live in-flight counts and the existing last-progress +-- clock; drained pauses remain ready and stuck in-flight pauses do not. +-- Control routes (enableJSON required): GET /control discovers enabled/operations; +-- POST /control/processors//pause or /resume returns atomic intake intent. +-- enableControl defaults False: actions return 403 control_disabled without mutation. +-- Success is {processor,state:"running"} or {processor,state:"paused",pausedAt}. +-- Actual POST Origin must be explicitly allowed, even when CORS is disabled. +-- Refusals: 403 origin_not_allowed; 404 processor_not_found/not_found; +-- 409 processor_terminal/processor_not_controllable/processor_draining; +-- 405 method_not_allowed with Allow GET or POST. Lifecycle refusals include processor. +-- Control adds no authentication or WebSocket mutation frames. module Shibuya.Metrics ( -- * Server Lifecycle startMetricsServer, startMetricsServerWithDeps, stopMetricsServer, withMetricsServer, + combinedApp, -- * Configuration MetricsServerConfig (..), defaultConfig, + -- * Host Application Helpers + corsMiddleware, + errorResponse, + -- * Server Handle MetricsServer (..), @@ -78,6 +155,8 @@ where import Shibuya.Metrics.Config (MetricsServerConfig (..), defaultConfig) +import Shibuya.Metrics.Cors (corsMiddleware) +import Shibuya.Metrics.Error (errorResponse) import Shibuya.Metrics.Health ( ApplicationStatus (..), DependencyCheck, @@ -88,5 +167,5 @@ ReadinessStatus (..), defaultHealthConfig, ) -import Shibuya.Metrics.Server (startMetricsServer, startMetricsServerWithDeps, stopMetricsServer, withMetricsServer) +import Shibuya.Metrics.Server (combinedApp, startMetricsServer, startMetricsServerWithDeps, stopMetricsServer, withMetricsServer) import Shibuya.Metrics.Types (ClientMessage (..), MetricsServer (..), ProcessorTerminalStatus (..), ServerMessage (..)) diff --git a/src/Shibuya/Metrics/Config.hs b/src/Shibuya/Metrics/Config.hs --- a/src/Shibuya/Metrics/Config.hs +++ b/src/Shibuya/Metrics/Config.hs @@ -5,6 +5,7 @@ ) where +import Data.Text (Text) import Data.Time.Clock (NominalDiffTime) import GHC.Generics (Generic) @@ -23,6 +24,14 @@ enablePrometheus :: !Bool, -- | Enable WebSocket endpoint (default: True) enableWebSocket :: !Bool, + -- | Enable reversible control POSTs (default False). The server performs + -- no authentication; use a trusted network or authenticating reverse proxy. + enableControl :: !Bool, + -- | Exact browser Origins allowed for HTTP CORS and WebSocket upgrades. + -- An empty list disables the policy (default). Entries must be serialized + -- @scheme://host[:port]@ origins without whitespace, paths or trailing slashes. + -- Origins are compared byte-for-byte; the wildcard @*@ is refused. + corsAllowedOrigins :: ![Text], -- | WebSocket push interval in microseconds (default: 100_000 = 100ms) wsPushIntervalUs :: !Int, -- | Maximum WebSocket connections (default: 100) @@ -30,6 +39,13 @@ -- | Maximum retained processor identifiers per WebSocket connection -- (default: 1,000). A client exceeding this policy is disconnected. wsMaxSubscriptions :: !Int, + -- | Maximum retained outbound frames, including required frames (default: 64; minimum: 2). + wsMaxQueuedFrames :: !Int, + -- | Maximum retained encoded outbound bytes (default: 1 MiB). + -- At most one popped frame and a transient encoding are outside this budget. + wsMaxQueuedBytes :: !Int, + -- | Deadline for a send and the entire orderly close flush (default: 5 seconds). + wsSendTimeoutUs :: !Int, -- | Timeout for liveness check in microseconds (default: 1_000_000 = 1s) livenessTimeoutMicros :: !Int, -- | Timeout for each dependency readiness check in microseconds (default: 1s) @@ -48,9 +64,14 @@ enableJSON = True, enablePrometheus = True, enableWebSocket = True, + enableControl = False, + corsAllowedOrigins = [], wsPushIntervalUs = 100_000, -- 100ms wsMaxConnections = 100, wsMaxSubscriptions = 1_000, + wsMaxQueuedFrames = 64, + wsMaxQueuedBytes = 1_048_576, + wsSendTimeoutUs = 5_000_000, livenessTimeoutMicros = 1_000_000, -- 1 second dependencyTimeoutMicros = 1_000_000, -- 1 second per dependency stuckThreshold = 60 -- 60 seconds diff --git a/src/Shibuya/Metrics/Control.hs b/src/Shibuya/Metrics/Control.hs new file mode 100644 --- /dev/null +++ b/src/Shibuya/Metrics/Control.hs @@ -0,0 +1,69 @@ +-- | Opt-in reversible intake control. This surface performs no authentication. +module Shibuya.Metrics.Control (controlApp) where + +import Data.Aeson (encode, object, (.=)) +import Data.Text (Text) +import Network.HTTP.Types (methodGet, methodPost, status200, status403, status404, status405, status409) +import Network.Wai (Application, mapResponseHeaders, pathInfo, requestHeaders, requestMethod, responseLBS) +import Shibuya.App (Master) +import Shibuya.Core.Metrics (ProcessorId (..)) +import Shibuya.Internal.Runner.Master + ( ControlOutcome (..), + PauseState (..), + pauseProcessorIO, + resumeProcessorIO, + ) +import Shibuya.Metrics.Config (MetricsServerConfig (..)) +import Shibuya.Metrics.Cors (isAllowedOrigin) +import Shibuya.Metrics.Error (errorResponse) + +-- | Discovery and POST actions. Successful bodies describe the intake state at +-- the atomic decision; another operation may supersede it before delivery. +controlApp :: MetricsServerConfig -> Master -> Application +controlApp config master = \req respond -> case pathInfo req of + ["control"] + | requestMethod req == methodGet -> + respond $ + json $ + object + ["enabled" .= config.enableControl, "operations" .= (["pause", "resume"] :: [Text])] + | otherwise -> respond $ methodRefusal "GET" + ["control", "processors", pid, operation] + | operation == "pause" || operation == "resume" -> + if requestMethod req /= methodPost + then respond $ methodRefusal "POST" + else + if not config.enableControl + then respond $ errorResponse status403 "Control operations are disabled" "control_disabled" [] + else + if not (allowedRequest req) + then respond $ errorResponse status403 "Origin not allowed" "origin_not_allowed" [] + else do + outcome <- (if operation == "pause" then pauseProcessorIO else resumeProcessorIO) master (ProcessorId pid) + respond $ case outcome of + ControlApplied state -> success pid state + ControlAlreadyInState state -> success pid state + ControlNotFound -> refusal pid status404 "Processor not found" "processor_not_found" + ControlNotControllable -> refusal pid status409 "Processor is not controllable" "processor_not_controllable" + ControlTerminal -> refusal pid status409 "Processor has terminated" "processor_terminal" + ControlDraining -> refusal pid status409 "Processor is draining" "processor_draining" + _ -> respond $ errorResponse status404 "Not found" "not_found" [] + where + -- Compile once outside the request loop, including the empty policy. A + -- bodyless/simple POST still requires its supplied Origin to be allowed. + permits = isAllowedOrigin config.corsAllowedOrigins + allowedRequest req = case [value | (name, value) <- requestHeaders req, name == "Origin"] of + [] -> True + [origin] -> permits origin + _ -> False + json = responseLBS status200 [("Content-Type", "application/json")] . encode + methodRefusal allowed = + mapResponseHeaders (("Allow", allowed) :) $ + errorResponse status405 "Method not allowed" "method_not_allowed" [] + refusal pid status message code = errorResponse status message code ["processor" .= pid] + success pid state = + json $ + object $ + ["processor" .= pid] ++ case state of + IntakeRunning -> ["state" .= ("running" :: Text)] + IntakePaused since -> ["state" .= ("paused" :: Text), "pausedAt" .= since] diff --git a/src/Shibuya/Metrics/Cors.hs b/src/Shibuya/Metrics/Cors.hs new file mode 100644 --- /dev/null +++ b/src/Shibuya/Metrics/Cors.hs @@ -0,0 +1,176 @@ +-- | Opt-in, exact-origin browser response policy. This is not authentication. +module Shibuya.Metrics.Cors + ( corsMiddleware, + isAllowedOrigin, + isValidOrigin, + originGuard, + ) +where + +import Data.Aeson (encode, object, (.=)) +import Data.ByteString (ByteString) +import Data.ByteString.Char8 qualified as BS +import Data.ByteString.Lazy qualified as LBS +import Data.Char (isAsciiLower, isAsciiUpper, isSpace, toLower) +import Data.List (foldl') +import Data.Set qualified as Set +import Data.Text (Text) +import Data.Text qualified as Text +import Data.Text.Encoding (encodeUtf8) +import Network.HTTP.Types (HeaderName, RequestHeaders, ResponseHeaders, status204, status403) +import Network.Wai (Middleware, mapResponseHeaders, requestHeaders, requestMethod, responseLBS) +import Network.WebSockets qualified as WS +import Shibuya.Metrics.Error (errorResponse) + +-- | Check configured origin syntax. Entries must still be written exactly as the +-- browser serializes them; matching does not normalize scheme, host or port. +isValidOrigin :: Text -> Bool +isValidOrigin origin = + not (Text.any isSpace origin) + && case Text.breakOn "://" origin of + (scheme, suffix) -> + case Text.uncons scheme of + Nothing -> False + Just (first, rest) -> + asciiLetter first + && Text.all (\c -> asciiLetter c || asciiDigit c || c `elem` ("+-." :: String)) rest + && not (Text.null suffix) + && let authority = Text.drop 3 suffix + in not (Text.null authority) + && Text.all (\c -> c > '\x20' && c < '\x7f' && c `notElem` ("/?#,@\\" :: String)) authority + where + asciiLetter c = isAsciiLower c || isAsciiUpper c + asciiDigit c = c >= '0' && c <= '9' + +-- | Compile the allow list once when constructing an application. Matching is +-- byte-for-byte: no case folding, default-port removal or prefix matching. +isAllowedOrigin :: [Text] -> ByteString -> Bool +isAllowedOrigin origins = (`Set.member` compiled) + where + compiled = Set.fromList $ map encodeUtf8 origins + +-- | Apply the same exact-origin policy before the inner WebSocket application +-- acquires a slot. Empty policy and headerless non-browser clients pass through. +-- A repeated Origin is refused even if each value individually matches. +originGuard :: [Text] -> WS.ServerApp -> WS.ServerApp +originGuard [] = id +originGuard origins = \app pending -> + let headers = WS.requestHeaders (WS.pendingRequest pending) + in if not (hasHeader "Origin" headers) || maybe False permits (singleHeader "Origin" headers) + then app pending + else + WS.rejectRequestWith + pending + WS.defaultRejectRequest + { WS.rejectCode = 403, + WS.rejectMessage = "Forbidden", + WS.rejectHeaders = [("Content-Type", "application/json"), ("Vary", "Origin")], + WS.rejectBody = + LBS.toStrict $ + encode $ + object + ["error" .= ("Origin not allowed" :: Text), "code" .= ("origin_not_allowed" :: Text)] + } + where + permits = isAllowedOrigin origins + +-- | Disabled by an empty list. When enabled, ordinary reads keep their status and +-- body, and every response varies by Origin. Only explicit allowed origins get +-- response permission. Preflights validate methods and header names before 204. +-- This middleware does not authorize mutation; control routes must validate the +-- actual request's single Origin before acting, even when CORS is disabled. +corsMiddleware :: [Text] -> Middleware +corsMiddleware [] = id +corsMiddleware origins = \app req respond -> + let headers = requestHeaders req + suppliedOrigin = singleHeader "Origin" headers + allowed = maybe False permits suppliedOrigin + preflight = + requestMethod req == "OPTIONS" + && hasHeader "Origin" headers + && hasHeader "Access-Control-Request-Method" headers + decorate = + mapResponseHeaders $ + if preflight + then mergeVary preflightVary + else decorateHeaders (if allowed then suppliedOrigin else Nothing) + in if preflight + then respond $ decorate $ case () of + _ + | not allowed -> errorResponse status403 "Origin not allowed" "origin_not_allowed" [] + | maybe True (`notElem` ["GET", "POST", "OPTIONS"]) (singleHeader "Access-Control-Request-Method" headers) -> + errorResponse status403 "CORS method not allowed" "cors_method_not_allowed" [] + | otherwise -> case requestedHeaders headers of + Nothing -> errorResponse status403 "CORS headers not allowed" "cors_headers_not_allowed" [] + Just names -> + responseLBS + status204 + [ ("Access-Control-Allow-Origin", maybe "" id suppliedOrigin), + ("Access-Control-Allow-Methods", "GET, POST, OPTIONS"), + ("Access-Control-Allow-Headers", names), + ("Access-Control-Max-Age", "600") + ] + "" + else app req (respond . decorate) + where + permits = isAllowedOrigin origins + +singleHeader :: HeaderName -> RequestHeaders -> Maybe ByteString +singleHeader name headers = case [value | (key, value) <- headers, key == name] of + [value] -> Just value + _ -> Nothing + +hasHeader :: HeaderName -> RequestHeaders -> Bool +hasHeader name = any ((== name) . fst) + +requestedHeaders :: RequestHeaders -> Maybe ByteString +requestedHeaders headers = case [value | (key, value) <- headers, key == "Access-Control-Request-Headers"] of + [] -> Just "Content-Type" + values -> + let names = concatMap (map trim . BS.split ',') values + in if all (not . BS.null) values && not (null names) && all validToken names + then Just (BS.intercalate ", " names) + else Nothing + where + validToken token = not (BS.null token) && token /= "*" && BS.all tokenChar token + tokenChar c = + isAsciiLower c + || isAsciiUpper c + || (c >= '0' && c <= '9') + || c `elem` ("!#$%&'*+-.^_`|~" :: String) + +trim :: ByteString -> ByteString +trim = BS.dropWhileEnd horizontalSpace . BS.dropWhile horizontalSpace + where + horizontalSpace c = c == ' ' || c == '\t' + +preflightVary :: [ByteString] +preflightVary = ["Origin", "Access-Control-Request-Method", "Access-Control-Request-Headers"] + +-- Preserve unrelated Vary tokens and collapse all Vary fields, including a +-- pre-existing wildcard. Strip stale permission/credential headers from hosts. +decorateHeaders :: Maybe ByteString -> ResponseHeaders -> ResponseHeaders +decorateHeaders origin headers = + mergeVary ["Origin"] (filter (\(name, _) -> name `notElem` permissionHeaders) headers) + <> maybe [] (\value -> [("Access-Control-Allow-Origin", value)]) origin + where + permissionHeaders = + [ "Access-Control-Allow-Origin", + "Access-Control-Allow-Credentials", + "Access-Control-Allow-Methods", + "Access-Control-Allow-Headers", + "Access-Control-Max-Age", + "Access-Control-Expose-Headers" + ] + +mergeVary :: [ByteString] -> ResponseHeaders -> ResponseHeaders +mergeVary required headers = + let existing = concatMap (map trim . BS.split ',') [value | (name, value) <- headers, name == "Vary"] + tokens = foldl' addToken [] (filter (not . BS.null) existing <> required) + vary = if "*" `elem` tokens then "*" else BS.intercalate ", " tokens + kept = filter ((/= "Vary") . fst) headers + in kept <> [("Vary", vary)] + where + addToken tokens token + | any (\existing -> BS.map toLower existing == BS.map toLower token) tokens = tokens + | otherwise = tokens <> [token] diff --git a/src/Shibuya/Metrics/Error.hs b/src/Shibuya/Metrics/Error.hs new file mode 100644 --- /dev/null +++ b/src/Shibuya/Metrics/Error.hs @@ -0,0 +1,17 @@ +-- | The shared HTTP error shape: human-readable text, a stable code, and context. +module Shibuya.Metrics.Error (errorResponse) where + +import Data.Aeson (encode, object, (.=)) +import Data.Aeson.Types (Pair) +import Data.Text (Text) +import Network.HTTP.Types (Status, hContentType) +import Network.Wai (Response, responseLBS) + +-- | Render an error as JSON. Extra context cannot override @error@ or @code@. +errorResponse :: Status -> Text -> Text -> [Pair] -> Response +errorResponse status message code extras = + responseLBS status [(hContentType, "application/json")] $ + encode $ + object $ + ["error" .= message, "code" .= code] + <> filter (\(key, _) -> key /= "error" && key /= "code") extras diff --git a/src/Shibuya/Metrics/Health.hs b/src/Shibuya/Metrics/Health.hs --- a/src/Shibuya/Metrics/Health.hs +++ b/src/Shibuya/Metrics/Health.hs @@ -40,7 +40,8 @@ import Data.Time.Clock (NominalDiffTime, UTCTime, diffUTCTime, getCurrentTime) import Shibuya.App (Master, getAllMetricsIO) import Shibuya.Core.Metrics - ( MetricsMap, + ( InFlightInfo (..), + MetricsMap, ProcessorId, ProcessorMetrics (..), ProcessorState (..), @@ -141,7 +142,8 @@ { total :: !Int, healthy :: !Int, failed :: !Int, - stuck :: !Int + stuck :: !Int, + paused :: !Int } deriving stock (Eq, Show) @@ -151,7 +153,8 @@ [ "total" .= ph.total, "healthy" .= ph.healthy, "failed" .= ph.failed, - "stuck" .= ph.stuck + "stuck" .= ph.stuck, + "paused" .= ph.paused ] -- | Status of an external dependency. @@ -249,16 +252,17 @@ analyzeProcessorHealth config now metrics lifecycles = let processorIds = Map.keysSet metrics <> Map.keysSet lifecycles total = length processorIds - (healthy, failed, stuck) = + (healthy, failed, stuck, paused) = foldr (categorize config now metrics lifecycles) - (0, 0, 0) + (0, 0, 0, 0) processorIds in ProcessorHealth { total = total, healthy = healthy, failed = failed, - stuck = stuck + stuck = stuck, + paused = paused } -- | Categorize a processor as healthy, failed, or stuck. @@ -268,22 +272,25 @@ MetricsMap -> LifecycleSnapshot -> ProcessorId -> - (Int, Int, Int) -> - (Int, Int, Int) -categorize config now metrics lifecycles pid counts@(h, f, s) = + (Int, Int, Int, Int) -> + (Int, Int, Int, Int) +categorize config now metrics lifecycles pid counts@(h, f, s, p) = case Map.lookup pid lifecycles of - Just LifecycleFailed {} -> (h, f + 1, s) + Just LifecycleFailed {} -> (h, f + 1, s, p) _ -> case Map.lookup pid metrics of Nothing -> counts Just pm -> case pm.state of - Idle -> (h + 1, f, s) + Idle -> (h + 1, f, s, p) Stopped -> counts - Failed _ _ -> (h, f + 1, s) + Failed _ _ -> (h, f + 1, s, p) + Paused info _ lastProgress + | info.inFlight > 0 && diffUTCTime now lastProgress > config.stuckThreshold -> (h, f, s + 1, p) + | otherwise -> (h, f, s, p + 1) Processing _ _ lastProgress -> let timeSinceProgress = diffUTCTime now lastProgress in if timeSinceProgress > config.stuckThreshold - then (h, f, s + 1) - else (h + 1, f, s) + then (h, f, s + 1, p) + else (h + 1, f, s, p) classifyApplication :: MasterPhase -> LifecycleSnapshot -> ApplicationStatus classifyApplication masterPhase lifecycles diff --git a/src/Shibuya/Metrics/JSON.hs b/src/Shibuya/Metrics/JSON.hs --- a/src/Shibuya/Metrics/JSON.hs +++ b/src/Shibuya/Metrics/JSON.hs @@ -18,6 +18,7 @@ import Network.Wai (Application, Response, pathInfo, responseLBS) import Shibuya.App (Master, getAllMetricsIO, getProcessorMetricsIO) import Shibuya.Core.Metrics (ProcessorId (..)) +import Shibuya.Metrics.Error (errorResponse) import Shibuya.Metrics.Health ( DependencyCheck, HealthConfig, @@ -60,12 +61,7 @@ Just m -> pure $ jsonResponse status200 $ encode m Nothing -> pure $ - jsonResponse status404 $ - encode $ - object - [ "error" .= ("Processor not found" :: Text), - "processor" .= procIdText - ] + errorResponse status404 "Processor not found" "processor_not_found" ["processor" .= procIdText] -- GET /health/live - liveness probe (fast, simple) ["health", "live"] -> do liveness <- checkLiveness config master @@ -92,10 +88,7 @@ ] -- Not found _ -> - pure $ - jsonResponse status404 $ - encode $ - object ["error" .= ("Not found" :: Text)] + pure $ errorResponse status404 "Not found" "not_found" [] jsonResponse :: Status -> LBS.ByteString -> Response jsonResponse status = responseLBS status [(hContentType, "application/json")] diff --git a/src/Shibuya/Metrics/Outbound.hs b/src/Shibuya/Metrics/Outbound.hs new file mode 100644 --- /dev/null +++ b/src/Shibuya/Metrics/Outbound.hs @@ -0,0 +1,101 @@ +-- | Pure, strictly bounded outbound admission. Constructors for frames and queues +-- are private so cached sizes cannot diverge from the retained strict payloads. +module Shibuya.Metrics.Outbound + ( Retention (..), + Frame, + frame, + payload, + generation, + retention, + Limits (..), + OutboundQueue, + empty, + queuedFrames, + queuedBytes, + entries, + Outcome (..), + enqueue, + pop, + discardUpdates, + ) +where + +import Data.ByteString (ByteString) +import Data.ByteString qualified as BS +import Data.Foldable (toList) +import Data.Sequence (Seq, ViewL (..), (><)) +import Data.Sequence qualified as Seq + +-- | Pongs are expendable but their loss does not invalidate the metrics ledger. +data Retention = Required | Update | Ephemeral + deriving stock (Eq, Show) + +data Frame = Frame !Int !Retention !ByteString + deriving stock (Eq, Show) + +frame :: Int -> Retention -> ByteString -> Frame +frame = Frame + +payload :: Frame -> ByteString +payload (Frame _ _ bytes) = bytes + +generation :: Frame -> Int +generation (Frame value _ _) = value + +retention :: Frame -> Retention +retention (Frame _ value _) = value + +data Limits = Limits {maxFrames :: !Int, maxBytes :: !Int} + deriving stock (Eq, Show) + +data OutboundQueue = Queue !(Seq Frame) !Int + deriving stock (Eq, Show) + +empty :: OutboundQueue +empty = Queue Seq.empty 0 + +queuedFrames :: OutboundQueue -> Int +queuedFrames (Queue values _) = Seq.length values + +queuedBytes :: OutboundQueue -> Int +queuedBytes (Queue _ bytes) = bytes + +entries :: OutboundQueue -> [Frame] +entries (Queue values _) = toList values + +data Outcome = Enqueued | DroppedUpdates | MustClose + deriving stock (Eq, Show) + +-- | Admit the entire batch atomically or none of it. Evict expendable frames +-- oldest-first; never evict a required frame. Saturation by required frames +-- closes the connection instead of exceeding either limit. A rejected update +-- marks the comparison ledger as needing a full resynchronization. +enqueue :: Limits -> [Frame] -> OutboundQueue -> (Outcome, OutboundQueue) +enqueue limits batch original = go False original + where + incoming = Seq.fromList batch + bytes = sum $ map (toInteger . BS.length . payload) batch + fits q = + toInteger (queuedFrames q) + toInteger (length batch) <= toInteger limits.maxFrames + && toInteger (queuedBytes q) + bytes <= toInteger limits.maxBytes + go dropped q@(Queue values size) + | fits q = (if dropped then DroppedUpdates else Enqueued, Queue (values >< incoming) (size + fromInteger bytes)) + | otherwise = case Seq.findIndexL ((/= Required) . retention) values of + Just index -> + let victim = Seq.index values index + next = Queue (Seq.deleteAt index values) (size - BS.length (payload victim)) + in go (dropped || retention victim == Update) next + Nothing + | any ((== Required) . retention) batch -> (MustClose, original) + | otherwise -> (if dropped || any ((== Update) . retention) batch then DroppedUpdates else Enqueued, q) + +pop :: OutboundQueue -> Maybe (Frame, OutboundQueue) +pop (Queue values size) = case Seq.viewl values of + EmptyL -> Nothing + value :< rest -> Just (value, Queue rest (size - BS.length (payload value))) + +-- | Remove only unsent metrics updates. Required frames keep their order. +discardUpdates :: OutboundQueue -> OutboundQueue +discardUpdates (Queue values _) = + let kept = Seq.filter ((/= Update) . retention) values + in Queue kept (sum $ fmap (BS.length . payload) kept) diff --git a/src/Shibuya/Metrics/Prometheus.hs b/src/Shibuya/Metrics/Prometheus.hs --- a/src/Shibuya/Metrics/Prometheus.hs +++ b/src/Shibuya/Metrics/Prometheus.hs @@ -20,6 +20,7 @@ ProcessorState (..), StreamStats (..), ) +import Shibuya.Core.Metrics qualified as Metrics -- | WAI application for Prometheus metrics endpoint. -- Handles: GET /metrics/prometheus @@ -46,12 +47,18 @@ renderHelp "shibuya_messages_failed_total" "Total messages that failed processing", renderType "shibuya_messages_failed_total" "counter", renderGaugeMetrics "shibuya_messages_failed_total" (fromIntegral . (.stats.failed)) metrics, - renderHelp "shibuya_processor_state" "Current processor state (1=idle, 2=processing, 3=failed, 4=stopped)", + renderHelp "shibuya_processor_state" "Current processor state (1=idle, 2=processing, 3=failed, 4=stopped, 5=paused)", renderType "shibuya_processor_state" "gauge", renderGaugeMetrics "shibuya_processor_state" (fromIntegral . stateToInt . (.state)) metrics, renderHelp "shibuya_processor_in_flight" "Number of messages currently being processed", renderType "shibuya_processor_in_flight" "gauge", - renderGaugeMetrics "shibuya_processor_in_flight" (fromIntegral . inFlightCount . (.state)) metrics + renderGaugeMetrics "shibuya_processor_in_flight" (fromIntegral . inFlightCount . (.state)) metrics, + renderHelp "shibuya_message_processing_seconds" "Processing attempt duration including finalization retries (one observation per message or emitted batch)", + renderType "shibuya_message_processing_seconds" "histogram", + Map.foldMapWithKey renderLatency metrics, + renderHelp "shibuya_processor_oldest_in_flight_seconds" "Age of oldest in-flight processing attempt", + renderType "shibuya_processor_oldest_in_flight_seconds" "gauge", + Map.foldMapWithKey renderOldest metrics ] renderHelp :: Text -> Text -> Builder.Builder @@ -97,8 +104,33 @@ stateToInt (Processing _ _ _) = 2 stateToInt (Failed _ _) = 3 stateToInt Stopped = 4 +stateToInt (Paused _ _ _) = 5 -- | Get in-flight count from processor state. inFlightCount :: ProcessorState -> Int inFlightCount (Processing info _ _) = info.inFlight +inFlightCount (Paused info _ _) = info.inFlight inFlightCount _ = 0 + +renderLatency :: ProcessorId -> ProcessorMetrics -> Builder.Builder +renderLatency (ProcessorId pid) pm = case pm.latency of + Nothing -> mempty + Just summary -> + let labels = Builder.byteString "{processor=\"" <> Builder.byteString (Text.encodeUtf8 (escapeLabel pid)) + prefix suffix = Builder.byteString "shibuya_message_processing_seconds" <> Builder.byteString suffix <> labels + bucket label count = + prefix "_bucket" + <> Builder.byteString "\",le=\"" + <> Builder.byteString (Text.encodeUtf8 label) + <> Builder.byteString "\"} " + <> Builder.integerDec count + <> Builder.char8 '\n' + scalar suffix value = prefix suffix <> Builder.byteString "\"} " <> value <> Builder.char8 '\n' + in mconcat (zipWith bucket Metrics.latencyBucketLabels (map (.count) summary.buckets ++ [summary.count])) + <> scalar "_sum" (Builder.doubleDec summary.sumSeconds) + <> scalar "_count" (Builder.integerDec summary.count) + +renderOldest :: ProcessorId -> ProcessorMetrics -> Builder.Builder +renderOldest pid pm = case pm.oldestInFlight of + Nothing -> mempty + Just oldest -> renderMetric "shibuya_processor_oldest_in_flight_seconds" (const oldest.ageSeconds) pid pm diff --git a/src/Shibuya/Metrics/Server.hs b/src/Shibuya/Metrics/Server.hs --- a/src/Shibuya/Metrics/Server.hs +++ b/src/Shibuya/Metrics/Server.hs @@ -16,22 +16,28 @@ where import Control.Concurrent.Async (async, cancel) +import Control.Concurrent.STM (atomically, check, readTVar) import Control.Exception (bracket, finally) -import Data.Aeson (encode, object, (.=)) +import Control.Monad (void) +import Data.List (find) import Data.String (fromString) -import Data.Text (Text) -import Network.HTTP.Types (hContentType, status404) -import Network.Wai (Application, Response, pathInfo, responseLBS) +import Data.Text qualified as Text +import Network.HTTP.Types (status404) +import Network.Wai (Application, pathInfo) import Network.Wai.Handler.Warp qualified as Warp import Network.Wai.Handler.WebSockets qualified as WaiWS import Network.WebSockets qualified as WS import Shibuya.App (Master) import Shibuya.Metrics.Config (MetricsServerConfig (..), defaultConfig) +import Shibuya.Metrics.Control (controlApp) +import Shibuya.Metrics.Cors (corsMiddleware, isValidOrigin, originGuard) +import Shibuya.Metrics.Error (errorResponse) import Shibuya.Metrics.Health (DependencyCheck, HealthConfig (..)) import Shibuya.Metrics.JSON (jsonAppWithHealth) import Shibuya.Metrics.Prometheus (prometheusApp) import Shibuya.Metrics.Types (MetricsServer (..)) -import Shibuya.Metrics.WebSocket (WebSocketState, newWebSocketState, shutdownWebSockets, websocketApp) +import Shibuya.Metrics.WebSocket (WebSocketState (..), newWebSocketState, shutdownWebSockets, websocketApp) +import System.Timeout (timeout) -- | Start the metrics server without dependency checks. -- Returns a handle that can be used to stop the server. @@ -50,15 +56,22 @@ wsState <- newWebSocketState config.wsMaxConnections let app = combinedApp config master wsState depChecks settings = - Warp.setPort config.port $ - Warp.setHost - (fromString config.host) - Warp.defaultSettings + Warp.setGracefulShutdownTimeout (Just 0) $ + Warp.setPort config.port $ + Warp.setHost + (fromString config.host) + Warp.defaultSettings serverAsync <- async $ Warp.runSettings settings app `finally` shutdownWebSockets wsState pure MetricsServer { serverThread = serverAsync, - serverPort = config.port + serverPort = config.port, + serverShutdown = do + shutdownWebSockets wsState + void $ + timeout (config.wsSendTimeoutUs + 100_000) $ + atomically $ + readTVar wsState.connectionCount >>= check . (== 0) } validateConfig :: MetricsServerConfig -> IO () @@ -68,14 +81,19 @@ | config.wsPushIntervalUs <= 0 = fail "MetricsServerConfig.wsPushIntervalUs must be positive" | config.wsMaxConnections <= 0 = fail "MetricsServerConfig.wsMaxConnections must be positive" | config.wsMaxSubscriptions <= 0 = fail "MetricsServerConfig.wsMaxSubscriptions must be positive" + | config.wsMaxQueuedFrames < 2 = fail "MetricsServerConfig.wsMaxQueuedFrames must be at least 2" + | config.wsMaxQueuedBytes <= 0 = fail "MetricsServerConfig.wsMaxQueuedBytes must be positive" + | config.wsSendTimeoutUs <= 0 = fail "MetricsServerConfig.wsSendTimeoutUs must be positive" | config.livenessTimeoutMicros <= 0 = fail "MetricsServerConfig.livenessTimeoutMicros must be positive" | config.dependencyTimeoutMicros <= 0 = fail "MetricsServerConfig.dependencyTimeoutMicros must be positive" | config.stuckThreshold <= 0 = fail "MetricsServerConfig.stuckThreshold must be positive" + | Just invalid <- find (not . isValidOrigin) config.corsAllowedOrigins = + fail $ "MetricsServerConfig.corsAllowedOrigins contains an invalid origin: " <> show (Text.unpack invalid) | otherwise = pure () -- | Stop the metrics server. stopMetricsServer :: MetricsServer -> IO () -stopMetricsServer server = cancel server.serverThread +stopMetricsServer server = server.serverShutdown `finally` cancel server.serverThread -- | Run an action with a metrics server, ensuring cleanup. withMetricsServer :: @@ -89,6 +107,9 @@ stopMetricsServer -- | Combined WAI application routing to all endpoints. +-- Hosts mounting this directly must enforce the configuration constraints themselves: +-- positive timeouts/resource limits and valid explicit serialized origins. The built-in +-- 'startMetricsServer' validates before starting; this pure constructor does not. combinedApp :: MetricsServerConfig -> Master -> @@ -100,15 +121,15 @@ then WaiWS.websocketsOr WS.defaultConnectionOptions - (websocketApp config master wsState) + (originGuard config.corsAllowedOrigins (websocketApp config master wsState)) fallback else fallback where - fallback = httpApp config master depChecks + fallback = corsMiddleware config.corsAllowedOrigins (httpApp config master depChecks) -- | HTTP application routing based on path. httpApp :: MetricsServerConfig -> Master -> [DependencyCheck] -> Application -httpApp config master depChecks req respond = do +httpApp config master depChecks = \req respond -> do let path = pathInfo req healthConfig = HealthConfig @@ -138,16 +159,13 @@ ["health", "ready"] | config.enableJSON -> jsonHandler req respond + ("control" : _) + | config.enableJSON -> controlHandler req respond -- WebSocket path info (for documentation, actual upgrade handled above) ["ws"] | config.enableWebSocket -> - respond $ notFoundResponse "WebSocket endpoint - use ws:// protocol" + respond $ errorResponse status404 "WebSocket endpoint - use ws:// protocol" "websocket_upgrade_required" [] -- Not found - _ -> respond $ notFoundResponse "Not found" - -notFoundResponse :: Text -> Response -notFoundResponse msg = - responseLBS - status404 - [(hContentType, "application/json")] - (encode $ object ["error" .= msg]) + _ -> respond $ errorResponse status404 "Not found" "not_found" [] + where + controlHandler = controlApp config master diff --git a/src/Shibuya/Metrics/Types.hs b/src/Shibuya/Metrics/Types.hs --- a/src/Shibuya/Metrics/Types.hs +++ b/src/Shibuya/Metrics/Types.hs @@ -68,6 +68,8 @@ Pong | -- | A formerly visible processor reached a retained terminal state ProcessorTerminal !ProcessorId !ProcessorTerminalStatus + | -- | In-band failure: stable code followed by human-readable message + ServerError !Text !Text | -- | Server is shutting down Goodbye deriving stock (Eq, Show, Generic) @@ -98,6 +100,8 @@ "error" .= failure, "messageId" .= messageId ] + toJSON (ServerError code message) = + object ["type" .= ("error" :: Text), "code" .= code, "message" .= message] toJSON Goodbye = object ["type" .= ("goodbye" :: Text)] instance FromJSON ServerMessage where @@ -115,6 +119,7 @@ "failed" -> TerminalFailed <$> v .: "error" <*> v .:? "messageId" other -> fail $ "Unknown terminal status: " <> Text.unpack other pure $ ProcessorTerminal pid terminal + "error" -> ServerError <$> v .: "code" <*> v .: "message" "goodbye" -> pure Goodbye other -> fail $ "Unknown message type: " <> Text.unpack other @@ -127,6 +132,8 @@ { -- | Async handle for the server thread serverThread :: !(Async ()), -- | Port the server is running on - serverPort :: !Port + serverPort :: !Port, + -- | Request and await bounded WebSocket shutdown before stopping Warp. + serverShutdown :: !(IO ()) } deriving stock (Generic) diff --git a/src/Shibuya/Metrics/WebSocket.hs b/src/Shibuya/Metrics/WebSocket.hs --- a/src/Shibuya/Metrics/WebSocket.hs +++ b/src/Shibuya/Metrics/WebSocket.hs @@ -1,13 +1,17 @@ --- | WebSocket endpoint for real-time metrics updates. +-- | Real-time metrics with independent production and one bounded application writer. module Shibuya.Metrics.WebSocket ( websocketApp, WebSocketState (..), newWebSocketState, shutdownWebSockets, + WebSocketHooks (..), + defaultWebSocketHooks, + websocketAppWithHooks, ) where import Control.Concurrent.Async (race_) +import Control.Concurrent.MVar (MVar, newMVar, withMVar) import Control.Concurrent.STM ( STM, TVar, @@ -19,294 +23,361 @@ readTVar, readTVarIO, registerDelay, + retry, writeTVar, ) -import Control.Exception (catch, finally, mask, throwIO) -import Control.Monad (forever, when) -import Data.Aeson (decode, encode) +import Control.Exception (catch, evaluate, finally, mask, throwIO) +import Control.Monad (forever, unless, when) +import Data.Aeson (eitherDecode, encode) +import Data.ByteString qualified as BS +import Data.ByteString.Lazy qualified as LBS import Data.Map.Strict qualified as Map import Data.Set (Set) import Data.Set qualified as Set import Data.Text (Text) import Data.Text qualified as Text +import Data.Word (Word16, Word64) +import GHC.Clock (getMonotonicTimeNSec) import Network.WebSockets qualified as WS import Shibuya.App (Master, getAllMetricsIO) -import Shibuya.Core.Metrics (MetricsMap, ProcessorId (..), ProcessorMetrics) +import Shibuya.Core.Metrics (MetricsMap, ProcessorId (..)) +import Shibuya.Core.Metrics qualified as Metrics import Shibuya.Core.Types (MessageId (..)) -import Shibuya.Internal.Runner.Master - ( ProcessorLifecycle (..), - getLifecycleSnapshotIO, - ) +import Shibuya.Internal.Runner.Master (ProcessorLifecycle (..), getLifecycleSnapshotIO) import Shibuya.Metrics.Config (MetricsServerConfig (..)) -import Shibuya.Metrics.Types - ( ClientMessage (..), - ProcessorTerminalStatus (..), - ServerMessage (..), - ) - --------------------------------------------------------------------------------- --- WebSocket State --------------------------------------------------------------------------------- +import Shibuya.Metrics.Outbound qualified as Out +import Shibuya.Metrics.Types (ClientMessage (..), ProcessorTerminalStatus (..), ServerMessage (..)) +import System.Timeout (timeout) --- | Shared state for WebSocket connections. +-- | Shared connection accounting and idempotent shutdown signal. data WebSocketState = WebSocketState - { -- | Current number of connections - connectionCount :: !(TVar Int), - -- | Maximum allowed connections - maxConnections :: !Int, - -- | Whether server shutdown has begun - shutdownRequested :: !(TVar Bool) - } + {connectionCount :: !(TVar Int), maxConnections :: !Int, shutdownRequested :: !(TVar Bool)} --- | Create new WebSocket state. newWebSocketState :: Int -> IO WebSocketState -newWebSocketState maxConns = do - countVar <- newTVarIO 0 - shutdownVar <- newTVarIO False - pure - WebSocketState - { connectionCount = countVar, - maxConnections = maxConns, - shutdownRequested = shutdownVar - } +newWebSocketState maxConns = WebSocketState <$> newTVarIO 0 <*> pure maxConns <*> newTVarIO False --- | Ask every active connection to send 'Goodbye' and finish. +-- | Stop admission and wake every producer to request an orderly, bounded flush. shutdownWebSockets :: WebSocketState -> IO () -shutdownWebSockets wsState = - atomically $ writeTVar wsState.shutdownRequested True +shutdownWebSockets state = atomically $ writeTVar state.shutdownRequested True -data AcquireResult - = Acquired - | AtCapacity - | ServerShuttingDown +data AcquireResult = Acquired | AtCapacity | ServerShuttingDown --- | Try to acquire a connection slot. acquireConnection :: WebSocketState -> STM AcquireResult -acquireConnection wsState = do - shuttingDown <- readTVar wsState.shutdownRequested - count <- readTVar wsState.connectionCount - if shuttingDown +acquireConnection state = do + stopping <- readTVar state.shutdownRequested + count <- readTVar state.connectionCount + if stopping then pure ServerShuttingDown else - if count >= wsState.maxConnections + if count >= state.maxConnections then pure AtCapacity - else do - writeTVar wsState.connectionCount (count + 1) - pure Acquired + else writeTVar state.connectionCount (count + 1) >> pure Acquired --- | Release a connection slot. releaseConnection :: WebSocketState -> STM () -releaseConnection wsState = - modifyTVar' wsState.connectionCount (\c -> max 0 (c - 1)) +releaseConnection state = modifyTVar' state.connectionCount (max 0 . subtract 1) --------------------------------------------------------------------------------- --- Connection State --------------------------------------------------------------------------------- +data Subscription = AllProcessors !(Set ProcessorId) | SelectedProcessors !(Set ProcessorId) --- | State for a single WebSocket connection. +-- The lock covers subscription, sampling, ledger and admission. It never covers +-- socket IO. Only the writer pops the queue; all producers obey the close cell. data ConnectionState = ConnectionState - { -- | Processor selection, including exclusions from subscribe-all + { producerLock :: !(MVar ()), subscriptions :: !(TVar Subscription), - -- | Last sent metrics for delta detection - lastMetrics :: !(TVar MetricsMap) + generation :: !(TVar Int), + lastMetrics :: !(TVar MetricsMap), + outbound :: !(TVar Out.OutboundQueue), + resyncPending :: !(TVar Bool), + closing :: !(TVar (Maybe CloseControl)) } -data Subscription - = AllProcessors !(Set ProcessorId) - | SelectedProcessors !(Set ProcessorId) +data CloseControl = CloseControl !Word16 !Text !Word64 --- | Create new connection state. newConnectionState :: IO ConnectionState -newConnectionState = do - subsVar <- newTVarIO $ AllProcessors Set.empty - lastVar <- newTVarIO Map.empty - pure - ConnectionState - { subscriptions = subsVar, - lastMetrics = lastVar - } +newConnectionState = + ConnectionState + <$> newMVar () + <*> newTVarIO (AllProcessors Set.empty) + <*> newTVarIO 0 + <*> newTVarIO Map.empty + <*> newTVarIO Out.empty + <*> newTVarIO False + <*> newTVarIO Nothing --------------------------------------------------------------------------------- --- WebSocket Application --------------------------------------------------------------------------------- +-- | Instrumentation for deterministic transport tests and externally mounted +-- servers. Hooks run inside the send deadline. Production uses no-op hooks and +-- an independent interval timer; a custom wait must also observe shutdown. +data WebSocketHooks = WebSocketHooks + { beforeWrite :: !(BS.ByteString -> IO ()), + afterReceive :: !(Out.OutboundQueue -> IO ()), + afterProduce :: !(Out.OutboundQueue -> IO ()), + waitForPush :: !(Int -> WebSocketState -> IO Bool) + } --- | WebSocket server application. -websocketApp :: - MetricsServerConfig -> - Master -> - WebSocketState -> - WS.ServerApp -websocketApp config master wsState pending = - mask $ \restore -> do - outcome <- atomically $ acquireConnection wsState - case outcome of - AtCapacity -> restore $ WS.rejectRequest pending "Too many connections" - ServerShuttingDown -> restore $ WS.rejectRequest pending "Server shutting down" - Acquired -> - restore (serveConnection config master wsState pending `catch` normalPeerClosure) - `finally` atomically (releaseConnection wsState) +defaultWebSocketHooks :: WebSocketHooks +defaultWebSocketHooks = WebSocketHooks (const $ pure ()) (const $ pure ()) (const $ pure ()) waitForPushOrShutdown +websocketApp :: MetricsServerConfig -> Master -> WebSocketState -> WS.ServerApp +websocketApp = websocketAppWithHooks defaultWebSocketHooks + +websocketAppWithHooks :: WebSocketHooks -> MetricsServerConfig -> Master -> WebSocketState -> WS.ServerApp +websocketAppWithHooks hooks config master state pending = mask $ \restore -> do + outcome <- atomically $ acquireConnection state + case outcome of + AtCapacity -> restore $ WS.rejectRequest pending "Too many connections" + ServerShuttingDown -> restore $ WS.rejectRequest pending "Server shutting down" + Acquired -> + restore (serveConnection hooks config master state pending `catch` normalPeerClosure) + `finally` atomically (releaseConnection state) + normalPeerClosure :: WS.ConnectionException -> IO () normalPeerClosure = \case WS.ConnectionClosed -> pure () WS.CloseRequest _ _ -> pure () unexpected -> throwIO unexpected -serveConnection :: MetricsServerConfig -> Master -> WebSocketState -> WS.PendingConnection -> IO () -serveConnection config master wsState pending = do +serveConnection :: WebSocketHooks -> MetricsServerConfig -> Master -> WebSocketState -> WS.PendingConnection -> IO () +serveConnection hooks config master state pending = do conn <- WS.acceptRequest pending WS.withPingThread conn 30 (pure ()) $ do - connState <- newConnectionState - metrics <- getAllMetricsIO master - WS.sendTextData conn $ encode $ MetricsSnapshot metrics - atomically $ writeTVar connState.lastMetrics metrics + cs <- newConnectionState + withMVar cs.producerLock $ \() -> subscriptionSnapshot config master cs + -- Orderly producers wait rather than returning and cancelling queued data. + -- A peer exception, writer timeout or writer completion cancels and joins both. race_ - (receiveLoop config master connState conn) - (pushLoop config master wsState connState conn) + (writerLoop hooks config cs conn) + (race_ (receiveLoop hooks config master state cs conn) (pushLoop hooks config master state cs)) --------------------------------------------------------------------------------- --- Receive Loop --------------------------------------------------------------------------------- +limits :: MetricsServerConfig -> Out.Limits +limits config = Out.Limits config.wsMaxQueuedFrames config.wsMaxQueuedBytes --- | Handle incoming messages from client. -receiveLoop :: MetricsServerConfig -> Master -> ConnectionState -> WS.Connection -> IO () -receiveLoop config master connState conn = forever $ do +encoded :: Int -> Out.Retention -> ServerMessage -> IO Out.Frame +encoded gen keep message = do + let bytes = LBS.toStrict $ encode message + _ <- evaluate $ BS.length bytes + pure $ Out.frame gen keep bytes + +-- Called under the producer lock. Encoding is forced before entering STM. +admit :: MetricsServerConfig -> ConnectionState -> [Out.Frame] -> IO Out.Outcome +admit config cs frames = do + outcome <- atomically $ do + stopped <- readTVar cs.closing + case stopped of + Just _ -> pure Out.MustClose + -- An empty terminal batch must not write an unchanged queue and wake + -- every waiting writer before the actual delta is ready. + Nothing | null frames -> pure Out.Enqueued + Nothing -> do + q <- readTVar cs.outbound + let (result, next) = Out.enqueue (limits config) frames q + writeTVar cs.outbound next + when (result == Out.DroppedUpdates) $ writeTVar cs.resyncPending True + pure result + when (outcome == Out.MustClose) $ requestClose config cs 1013 "Outbound queue capacity exceeded" + pure outcome + +requestClose :: MetricsServerConfig -> ConnectionState -> Word16 -> Text -> IO () +requestClose config cs code reason = do + now <- getMonotonicTimeNSec + atomically $ do + current <- readTVar cs.closing + case current of + Just _ -> pure () + Nothing -> + writeTVar cs.closing $ + Just $ + CloseControl + code + reason + (now + fromIntegral config.wsSendTimeoutUs * 1_000) + +awaitWriter :: IO a +awaitWriter = atomically retry + +receiveLoop :: WebSocketHooks -> MetricsServerConfig -> Master -> WebSocketState -> ConnectionState -> WS.Connection -> IO () +receiveLoop hooks config master state cs conn = forever $ do msg <- WS.receiveData conn - case decode msg of - Nothing -> pure () -- Ignore invalid messages - Just clientMsg -> handleClientMessage config.wsMaxSubscriptions master connState conn clientMsg + withMVar cs.producerLock $ \() -> do + stopped <- readTVarIO cs.closing + case stopped of + Just _ -> pure () + Nothing -> do + shuttingDown <- readTVarIO state.shutdownRequested + unless shuttingDown $ case eitherDecode msg of + Left reason -> do + gen <- readTVarIO cs.generation + err <- encoded gen Out.Required $ ServerError "invalid_message" (Text.pack reason) + _ <- admit config cs [err] + pure () + Right clientMsg -> handleClientMessage config master cs clientMsg + readTVarIO cs.outbound >>= hooks.afterReceive + stopped <- readTVarIO cs.closing + case stopped of Just _ -> awaitWriter; Nothing -> pure () --- | Handle a client message. -handleClientMessage :: - Int -> - Master -> - ConnectionState -> - WS.Connection -> - ClientMessage -> - IO () -handleClientMessage maxSubscriptions master connState conn = \case - SubscribeAll -> do - atomically $ writeTVar connState.subscriptions $ AllProcessors Set.empty - -- Send snapshot of all metrics - metrics <- getAllMetricsIO master - WS.sendTextData conn $ encode $ MetricsSnapshot metrics - atomically $ writeTVar connState.lastMetrics metrics - Subscribe pids -> do - subscription <- atomically $ do - current <- readTVar connState.subscriptions - let newSubs = case current of +handleClientMessage :: MetricsServerConfig -> Master -> ConnectionState -> ClientMessage -> IO () +handleClientMessage config master cs = \case + Ping -> do + gen <- readTVarIO cs.generation + pong <- encoded gen Out.Ephemeral Pong + _ <- admit config cs [pong] + pure () + message -> do + current <- readTVarIO cs.subscriptions + let next = case message of + SubscribeAll -> AllProcessors Set.empty + Subscribe pids -> case current of AllProcessors _ -> SelectedProcessors $ Set.fromList pids SelectedProcessors existing -> SelectedProcessors $ existing <> Set.fromList pids - if subscriptionSize newSubs > maxSubscriptions - then pure Nothing - else writeTVar connState.subscriptions newSubs >> pure (Just newSubs) - case subscription of - Nothing -> rejectOversizedSubscription conn maxSubscriptions - Just accepted -> do - allMetrics <- getAllMetricsIO master - let filtered = filterMetrics accepted allMetrics - WS.sendTextData conn $ encode $ MetricsSnapshot filtered - atomically $ writeTVar connState.lastMetrics filtered - Unsubscribe pids -> do - accepted <- atomically $ do - current <- readTVar connState.subscriptions - let removed = Set.fromList pids - newSubs = case current of - AllProcessors excluded -> AllProcessors $ excluded <> removed - SelectedProcessors existing -> SelectedProcessors $ Set.difference existing removed - if subscriptionSize newSubs > maxSubscriptions - then pure False - else writeTVar connState.subscriptions newSubs >> pure True - when (not accepted) $ rejectOversizedSubscription conn maxSubscriptions - Ping -> - WS.sendTextData conn $ encode Pong + Unsubscribe pids -> case current of + AllProcessors excluded -> AllProcessors $ excluded <> Set.fromList pids + SelectedProcessors existing -> SelectedProcessors $ existing `Set.difference` Set.fromList pids + if subscriptionSize next > config.wsMaxSubscriptions + then rejectOversizedSubscription config cs + else do + atomically $ do + writeTVar cs.subscriptions next + modifyTVar' cs.generation (+ 1) + modifyTVar' cs.outbound Out.discardUpdates + subscriptionSnapshot config master cs subscriptionSize :: Subscription -> Int subscriptionSize = \case AllProcessors excluded -> Set.size excluded SelectedProcessors selected -> Set.size selected -rejectOversizedSubscription :: WS.Connection -> Int -> IO () -rejectOversizedSubscription conn limit = - WS.sendCloseCode conn 1008 $ - "WebSocket processor subscription limit exceeded (maximum " - <> Text.pack (show limit) - <> ")" +subscriptionSnapshot :: MetricsServerConfig -> Master -> ConnectionState -> IO () +subscriptionSnapshot config master cs = do + subscription <- readTVarIO cs.subscriptions + metrics <- filterMetrics subscription <$> getAllMetricsIO master + gen <- readTVarIO cs.generation + snapshot <- encoded gen Out.Required $ MetricsSnapshot metrics + result <- admit config cs [snapshot] + unless (result == Out.MustClose) $ atomically $ do + writeTVar cs.lastMetrics metrics + writeTVar cs.resyncPending False --------------------------------------------------------------------------------- --- Push Loop --------------------------------------------------------------------------------- +rejectOversizedSubscription :: MetricsServerConfig -> ConnectionState -> IO () +rejectOversizedSubscription config cs = do + let reason = "WebSocket processor subscription limit exceeded (maximum " <> Text.pack (show config.wsMaxSubscriptions) <> ")" + gen <- readTVarIO cs.generation + err <- encoded gen Out.Required $ ServerError "subscription_limit_exceeded" reason + -- Set policy close first in the same producer critical section: capacity failure + -- cannot override 1008. Admission below is atomic with the close signal. + now <- getMonotonicTimeNSec + atomically $ do + q <- readTVar cs.outbound + let (_, next) = Out.enqueue (limits config) [err] q + writeTVar cs.outbound next + writeTVar cs.closing $ Just $ CloseControl 1008 reason (now + fromIntegral config.wsSendTimeoutUs * 1_000) --- | Push metrics updates to client at configured interval. -pushLoop :: - MetricsServerConfig -> - Master -> - WebSocketState -> - ConnectionState -> - WS.Connection -> - IO () -pushLoop config master wsState connState conn = loop +pushLoop :: WebSocketHooks -> MetricsServerConfig -> Master -> WebSocketState -> ConnectionState -> IO () +pushLoop hooks config master state cs = loop where loop = do - shuttingDown <- waitForPushOrShutdown config.wsPushIntervalUs wsState - if shuttingDown - then WS.sendTextData conn $ encode Goodbye - else pushUpdates master connState conn >> loop + shutdown <- hooks.waitForPush config.wsPushIntervalUs state + withMVar cs.producerLock $ \() -> do + stopped <- readTVarIO cs.closing + case stopped of + Just _ -> pure () + Nothing -> do + requested <- readTVarIO state.shutdownRequested + if shutdown || requested then orderlyShutdown else pushUpdates config master cs + readTVarIO cs.outbound >>= hooks.afterProduce + stopped <- readTVarIO cs.closing + case stopped of Just _ -> awaitWriter; Nothing -> loop + orderlyShutdown = do + gen <- readTVarIO cs.generation + goodbye <- encoded gen Out.Required Goodbye + atomically $ modifyTVar' cs.outbound Out.discardUpdates + _ <- admit config cs [goodbye] + requestClose config cs 1000 "Server shutting down" waitForPushOrShutdown :: Int -> WebSocketState -> IO Bool -waitForPushOrShutdown intervalUs wsState = do - intervalElapsed <- registerDelay intervalUs +waitForPushOrShutdown intervalUs state = do + elapsed <- registerDelay intervalUs atomically $ - (readTVar wsState.shutdownRequested >>= \requested -> check requested >> pure True) - `orElse` (readTVar intervalElapsed >>= \elapsed -> check elapsed >> pure False) + (readTVar state.shutdownRequested >>= \requested -> check requested >> pure True) + `orElse` (readTVar elapsed >>= \ready -> check ready >> pure False) -pushUpdates :: Master -> ConnectionState -> WS.Connection -> IO () -pushUpdates master connState conn = do - currentMetrics <- getAllMetricsIO master +pushUpdates :: MetricsServerConfig -> Master -> ConnectionState -> IO () +pushUpdates config master cs = do + current <- getAllMetricsIO master lifecycle <- getLifecycleSnapshotIO master - subscription <- readTVarIO connState.subscriptions - lastSent <- readTVarIO connState.lastMetrics - let filteredMetrics = filterMetrics subscription currentMetrics - _ <- Map.traverseWithKey (sendIfChanged lastSent conn) filteredMetrics - let removed = Map.keysSet lastSent `Set.difference` Map.keysSet currentMetrics - mapM_ (sendTerminal lifecycle conn) $ Set.toList removed - atomically $ writeTVar connState.lastMetrics filteredMetrics + subscription <- readTVarIO cs.subscriptions + previous <- readTVarIO cs.lastMetrics + resync <- readTVarIO cs.resyncPending + gen <- readTVarIO cs.generation + let filtered = filterMetrics subscription current + removed = Map.keysSet previous `Set.difference` Map.keysSet current + terminals = + [ ProcessorTerminal pid status + | pid <- Set.toList removed, + Just status <- [Map.lookup pid lifecycle >>= terminalStatus] + ] + terminalFrames <- mapM (encoded gen Out.Required) terminals + if resync + then do + atomically $ modifyTVar' cs.outbound Out.discardUpdates + err <- encoded gen Out.Required $ ServerError "queue_overflow" "Updates dropped; replacing metrics snapshot" + snapshot <- encoded gen Out.Required $ MetricsSnapshot filtered + result <- admit config cs (terminalFrames <> [err, snapshot]) + unless (result == Out.MustClose) $ atomically $ do + writeTVar cs.lastMetrics filtered + writeTVar cs.resyncPending False + else do + -- Required terminal events precede updates and survive eviction. + result <- admit config cs terminalFrames + case result of + Out.MustClose -> pure () + -- Terminal admission can itself evict a delta. Stop ordinary updates + -- immediately, but retire admitted terminals before the recovery tick. + Out.DroppedUpdates -> atomically $ writeTVar cs.lastMetrics $ Map.withoutKeys previous removed + Out.Enqueued -> do + let updates = [ProcessorUpdate pid pm | (pid, pm) <- Map.toList filtered, fmap stableView (Map.lookup pid previous) /= Just (stableView pm)] + frames <- mapM (encoded gen Out.Update) updates + updateResult <- admit config cs frames + unless (updateResult == Out.MustClose) $ atomically $ do + -- Retain the old live comparison ledger after delta loss, but retire + -- admitted terminals so the recovery tick never repeats them. + pending <- readTVar cs.resyncPending + writeTVar cs.lastMetrics $ if pending then Map.withoutKeys previous removed else filtered -filterMetrics :: Subscription -> MetricsMap -> MetricsMap -filterMetrics subscription = - Map.filterWithKey $ \pid _ -> case subscription of - AllProcessors excluded -> Set.notMember pid excluded - SelectedProcessors selected -> Set.member pid selected +writerLoop :: WebSocketHooks -> MetricsServerConfig -> ConnectionState -> WS.Connection -> IO () +writerLoop hooks config cs conn = loop + where + loop = do + next <- atomically $ do + q <- readTVar cs.outbound + case Out.pop q of + Just (value, rest) -> writeTVar cs.outbound rest >> pure (Left value) + Nothing -> readTVar cs.closing >>= maybe retry (pure . Right) + case next of + Left value -> bounded (hooks.beforeWrite (Out.payload value) >> WS.sendTextData conn (Out.payload value)) >> loop + Right (CloseControl code reason _) -> bounded $ WS.sendCloseCode conn code reason + bounded action = do + now <- getMonotonicTimeNSec + closing <- readTVarIO cs.closing + let remaining = case closing of + Nothing -> config.wsSendTimeoutUs + Just (CloseControl _ _ deadline) -> if now >= deadline then 0 else fromIntegral ((deadline - now) `div` 1_000) + result <- timeout (min config.wsSendTimeoutUs remaining) action + case result of Nothing -> throwIO WS.ConnectionClosed; Just () -> pure () -sendTerminal :: Map.Map ProcessorId ProcessorLifecycle -> WS.Connection -> ProcessorId -> IO () -sendTerminal lifecycle conn pid = - case Map.lookup pid lifecycle >>= terminalStatus of - Nothing -> pure () - Just status -> WS.sendTextData conn $ encode $ ProcessorTerminal pid status +filterMetrics :: Subscription -> MetricsMap -> MetricsMap +filterMetrics subscription = Map.filterWithKey $ \pid _ -> case subscription of + AllProcessors excluded -> Set.notMember pid excluded + SelectedProcessors selected -> Set.member pid selected terminalStatus :: ProcessorLifecycle -> Maybe ProcessorTerminalStatus terminalStatus = \case LifecycleStopped -> Just TerminalStopped - LifecycleFailed failure messageId -> - Just $ TerminalFailed failure (messageIdText <$> messageId) + LifecycleFailed failure messageId -> Just $ TerminalFailed failure (messageIdText <$> messageId) LifecycleRunning -> Nothing LifecycleDraining -> Nothing messageIdText :: MessageId -> Text messageIdText (MessageId value) = value --- | Send update if metrics have changed. -sendIfChanged :: - MetricsMap -> - WS.Connection -> - ProcessorId -> - ProcessorMetrics -> - IO () -sendIfChanged lastSent conn pid metrics = do - let changed = case Map.lookup pid lastSent of - Nothing -> True - Just old -> old /= metrics - when changed $ - WS.sendTextData conn $ - encode $ - ProcessorUpdate pid metrics +-- Ages grow without a state change. Clients extrapolate from their last frame; +-- snapshots and HTTP scrapes always include a freshly sampled age. +stableView :: Metrics.ProcessorMetrics -> Metrics.ProcessorMetrics +stableView pm = pm {Metrics.oldestInFlight = fmap (\oldest -> oldest {Metrics.ageSeconds = 0}) pm.oldestInFlight} diff --git a/test/Main.hs b/test/Main.hs --- a/test/Main.hs +++ b/test/Main.hs @@ -1,7 +1,10 @@ module Main (main) where +import Shibuya.Metrics.ControlSpec qualified as Control +import Shibuya.Metrics.CorsSpec qualified import Shibuya.Metrics.HealthSpec qualified import Shibuya.Metrics.JSONSpec qualified +import Shibuya.Metrics.OutboundQueueSpec qualified import Shibuya.Metrics.PrometheusSpec qualified import Shibuya.Metrics.ServerSpec qualified import Shibuya.Metrics.TypesSpec qualified @@ -10,8 +13,11 @@ main :: IO () main = hspec $ do + Control.spec + Shibuya.Metrics.CorsSpec.spec Shibuya.Metrics.ServerSpec.spec Shibuya.Metrics.JSONSpec.spec + Shibuya.Metrics.OutboundQueueSpec.spec Shibuya.Metrics.PrometheusSpec.spec Shibuya.Metrics.TypesSpec.spec Shibuya.Metrics.WebSocketSpec.spec diff --git a/test/Shibuya/Metrics/ControlSpec.hs b/test/Shibuya/Metrics/ControlSpec.hs new file mode 100644 --- /dev/null +++ b/test/Shibuya/Metrics/ControlSpec.hs @@ -0,0 +1,211 @@ +module Shibuya.Metrics.ControlSpec (spec) where + +import Control.Concurrent.STM +import Control.Exception (bracket) +import Control.Monad (forM_) +import Data.Aeson (Value (..), decode, eitherDecode, object, (.=)) +import Data.Aeson.Key qualified as Key +import Data.Aeson.KeyMap qualified as KM +import Data.ByteString (ByteString) +import Data.ByteString.Lazy qualified as LBS +import Data.Text qualified as Text +import Effectful (liftIO, runEff) +import Network.HTTP.Client qualified as HTTP +import Network.HTTP.Types +import Network.Wai (Application, requestHeaders, requestMethod) +import Network.Wai.Handler.Warp qualified as Warp +import Network.Wai.Test +import Network.WebSockets qualified as WS +import Shibuya qualified as Public +import Shibuya.Adapter (Adapter (..)) +import Shibuya.Core.Metrics +import Shibuya.Internal.Runner.Master +import Shibuya.Internal.Runner.Pause +import Shibuya.Metrics.Config +import Shibuya.Metrics.Server (combinedApp) +import Shibuya.Metrics.TestSupport (fixedTime, registerIdleProcessor, withMaster) +import Shibuya.Metrics.Types (ServerMessage (..)) +import Shibuya.Metrics.WebSocket (newWebSocketState) +import Streamly.Data.Stream qualified as Stream +import System.Timeout (timeout) +import Test.Hspec + +spec :: Spec +spec = do + around withMaster $ describe "gated control routes" $ do + it "discovers both gate positions and refuses disabled mutations" $ \master -> do + _ <- controllableFixture master + forM_ [False, True] $ \enabled -> do + app <- application defaultConfig {enableControl = enabled} master + response <- send app methodGet "/control" [] + response.simpleStatus `shouldBe` status200 + decode response.simpleBody `shouldBe` Just (object ["enabled" .= enabled, "operations" .= (["pause", "resume"] :: [String])]) + app <- application defaultConfig master + forM_ ["pause", "resume"] $ \op -> do + response <- send app methodPost ("/control/processors/orders/" <> op) [("Origin", "https://forbidden.example")] + assertCode response status403 "control_disabled" + isProcessorPausedIO master (ProcessorId "orders") `shouldReturn` Just False + + it "makes repeated pause and resume idempotent with an unchanged pause timestamp" $ \master -> do + _ <- controllableFixture master + app <- application defaultConfig {enableControl = True} master + first <- send app methodPost "/control/processors/orders/pause" [] + second <- send app methodPost "/control/processors/orders/pause" [] + first.simpleStatus `shouldBe` status200 + first.simpleBody `shouldBe` second.simpleBody + field "state" first `shouldBe` Just (String "paused") + field "pausedAt" first `shouldSatisfy` maybe False (const True) + isProcessorPausedIO master (ProcessorId "orders") `shouldReturn` Just True + forM_ [1, 2 :: Int] $ \_ -> do + response <- send app methodPost "/control/processors/orders/resume" [] + decode response.simpleBody `shouldBe` Just (object ["processor" .= ("orders" :: String), "state" .= ("running" :: String)]) + isProcessorPausedIO master (ProcessorId "orders") `shouldReturn` Just False + + it "maps unknown, uncontrollable, terminal and draining outcomes" $ \master -> do + (_, pause) <- controllableFixture master + _ <- registerIdleProcessor master (ProcessorId "fixture") + app <- application defaultConfig {enableControl = True} master + send app methodPost "/control/processors/missing/pause" [] >>= \r -> assertCode r status404 "processor_not_found" + send app methodPost "/control/processors/fixture/pause" [] >>= \r -> assertCode r status409 "processor_not_controllable" + runEff $ drainAndResume master + send app methodPost "/control/processors/orders/pause" [] >>= \r -> assertCode r status409 "processor_draining" + isPaused pause `shouldReturn` False + markProcessorFailedIO master (ProcessorId "orders") "failure" Nothing + runEff $ unregisterProcessor master (ProcessorId "orders") + send app methodPost "/control/processors/orders/resume" [] >>= \r -> assertCode r status409 "processor_terminal" + + it "does not confuse cold Failed observation with terminal lifecycle" $ \master -> do + (h, _) <- controllableFixture master + finishProcessing h (Left "transient") + app <- application defaultConfig {enableControl = True} master + response <- send app methodPost "/control/processors/orders/pause" [] + response.simpleStatus `shouldBe` status200 + isProcessorPausedIO master (ProcessorId "orders") `shouldReturn` Just True + + it "checks methods, unknown paths and the JSON enable flag" $ \master -> do + app <- application defaultConfig {enableControl = True} master + forM_ [(methodDelete, "/control", "GET"), (methodGet, "/control/processors/orders/pause", "POST")] $ \(method, path, allowed) -> do + response <- send app method path [] + assertCode response status405 "method_not_allowed" + lookup "Allow" response.simpleHeaders `shouldBe` Just allowed + send app methodGet "/control/other" [] >>= \r -> assertCode r status404 "not_found" + hidden <- application defaultConfig {enableJSON = False, enableControl = True} master + send hidden methodPost "/control/processors/orders/pause" [] >>= \r -> assertCode r status404 "not_found" + send hidden methodGet "/control" [] >>= \r -> assertCode r status404 "not_found" + + it "refuses forbidden and repeated Origins before a bodyless or simple POST can mutate" $ \master -> do + _ <- controllableFixture master + forM_ [[], ["https://ops.example.com"]] $ \origins -> do + app <- application defaultConfig {enableControl = True, corsAllowedOrigins = origins} master + forM_ [[("Origin", "https://evil.example")], [("Origin", "null")], [("Origin", "")], [("Origin", "https://ops.example.com"), ("Origin", "https://ops.example.com")]] $ \headers -> + forM_ [[], [("Content-Type", "text/plain")]] $ \content -> do + response <- send app methodPost "/control/processors/orders/pause" (headers ++ content) + assertCode response status403 "origin_not_allowed" + isProcessorPausedIO master (ProcessorId "orders") `shouldReturn` Just False + app <- application defaultConfig {enableControl = True, corsAllowedOrigins = ["https://ops.example.com"]} master + response <- send app methodPost "/control/processors/orders/pause" [("Origin", "https://ops.example.com")] + response.simpleStatus `shouldBe` status200 + isProcessorPausedIO master (ProcessorId "orders") `shouldReturn` Just True + + it "serves the allowed-origin action preflight" $ \master -> do + app <- application defaultConfig {enableControl = True, corsAllowedOrigins = ["https://ops.example.com"]} master + response <- send app methodOptions "/control/processors/orders/pause" [("Origin", "https://ops.example.com"), ("Access-Control-Request-Method", "POST")] + response.simpleStatus `shouldBe` status204 + lookup "Access-Control-Allow-Origin" response.simpleHeaders `shouldBe` Just "https://ops.example.com" + lookup "Access-Control-Allow-Methods" response.simpleHeaders `shouldBe` Just "GET, POST, OPTIONS" + + describe "real HTTP and WebSocket control" $ + it "pauses admission, drains an admitted message, observes paused state, resumes and shuts down" $ do + polls <- newTQueueIO + releases <- newTQueueIO + started <- newTQueueIO + handlerGate <- newEmptyTMVarIO + finalized <- newTVarIO (0 :: Int) + stopped <- newTVarIO False + let acquire = runEff $ Public.runTracingNoop $ do + let source = + Stream.unfoldrM + ( \n -> do + liftIO $ atomically $ writeTQueue polls n + next <- + liftIO $ + atomically $ + (readTVar stopped >>= check >> pure Nothing) `orElse` (Just <$> readTQueue releases) + pure $ fmap (\() -> (Public.mkIngested (Public.mkEnvelope (Public.MessageId "live") ()) (Public.AckHandle (\_ -> liftIO $ atomically $ modifyTVar' finalized (+ 1))), n + 1)) next + ) + (1 :: Int) + adapter = Adapter "live" source (liftIO $ atomically $ writeTVar stopped True) + handler _ = do + liftIO $ atomically $ writeTQueue started () + liftIO $ atomically $ takeTMVar handlerGate + pure Public.AckOk + result <- Public.runApp Public.defaultAppConfig [(ProcessorId "orders", Public.mkProcessor adapter handler)] + either (liftIO . ioError . userError . show) pure result + release app = runEff $ Public.runTracingNoop $ Public.stopAppGracefully Public.defaultShutdownConfig app + bracket acquire release $ \app -> do + bounded (atomically $ readTQueue polls) `shouldReturn` 1 + let master = Public.getAppMaster app + server <- application defaultConfig {enableControl = True, wsPushIntervalUs = 10_000} master + Warp.testWithApplication (pure server) $ \port -> do + manager <- HTTP.newManager HTTP.defaultManagerSettings + WS.runClient "127.0.0.1" port "/ws" $ \connection -> do + _ <- bounded $ WS.receiveData @LBS.ByteString connection + let post operation = HTTP.httpLbs (HTTP.parseRequest_ ("http://127.0.0.1:" <> show port <> "/control/processors/orders/" <> operation)) {HTTP.method = methodPost} manager + pause <- post "pause" + HTTP.responseStatus pause `shouldBe` status200 + atomically $ writeTQueue releases () + bounded $ atomically $ readTQueue started + observePaused connection + atomically $ putTMVar handlerGate () + bounded $ atomically $ readTVar finalized >>= check . (== 1) + timeout 100_000 (atomically $ readTQueue polls) `shouldReturn` Nothing + metric <- HTTP.httpLbs (HTTP.parseRequest_ ("http://127.0.0.1:" <> show port <> "/metrics/orders")) manager + case eitherDecode (HTTP.responseBody metric) of + Right m -> (m :: ProcessorMetrics).state `shouldSatisfy` (\case Paused {} -> True; _ -> False) + Left err -> expectationFailure err + resumed <- post "resume" + HTTP.responseStatus resumed `shouldBe` status200 + bounded (atomically $ readTQueue polls) `shouldReturn` 2 + atomically $ writeTQueue releases () + bounded $ atomically $ readTQueue started + atomically $ putTMVar handlerGate () + bounded $ atomically $ readTVar finalized >>= check . (== 2) + pure () + +application :: MetricsServerConfig -> Master -> IO Application +application config master = do + state <- newWebSocketState config.wsMaxConnections + pure $ combinedApp config master state [] + +controllableFixture :: Master -> IO (MetricsHandle, PauseHandle) +controllableFixture master = do + h <- newMetricsHandle fixedTime + pause <- newPauseHandle + runEff $ registerProcessorWithPause master (ProcessorId "orders") h (Just pause) + pure (h, pause) + +send :: Application -> Method -> ByteString -> RequestHeaders -> IO SResponse +send app method path headers = runSession (request $ setPath defaultRequest {requestMethod = method, requestHeaders = headers} path) app + +field :: Key.Key -> SResponse -> Maybe Value +field key response = case decode response.simpleBody of + Just (Object values) -> KM.lookup key values + _ -> Nothing + +assertCode :: SResponse -> Status -> String -> Expectation +assertCode response status code = do + response.simpleStatus `shouldBe` status + field "code" response `shouldBe` Just (String (Text.pack code)) + +observePaused :: WS.Connection -> IO () +observePaused connection = bounded loop + where + loop = do + bytes <- WS.receiveData connection + case eitherDecode bytes of + Right (ProcessorUpdate _ metrics) | Paused {} <- metrics.state -> pure () + Right (_ :: ServerMessage) -> loop + Left err -> expectationFailure err + +bounded :: IO a -> IO a +bounded action = timeout 2_000_000 action >>= maybe (expectationFailure "timed out" >> error "timeout") pure diff --git a/test/Shibuya/Metrics/CorsSpec.hs b/test/Shibuya/Metrics/CorsSpec.hs new file mode 100644 --- /dev/null +++ b/test/Shibuya/Metrics/CorsSpec.hs @@ -0,0 +1,174 @@ +module Shibuya.Metrics.CorsSpec (spec) where + +import Control.Exception (IOException) +import Control.Monad (forM_) +import Data.Aeson (decode, object, (.=)) +import Data.ByteString (ByteString) +import Data.Char (toLower) +import Data.List (isInfixOf) +import Data.Text (Text) +import Network.HTTP.Types (RequestHeaders, hContentType, status200, status204, status403, status404, status503) +import Network.Wai (Application, requestHeaders, requestMethod, responseLBS) +import Network.Wai.Test (SResponse (..), defaultRequest, request, runSession, setPath) +import Shibuya.App (Master) +import Shibuya.Core.Metrics (ProcessorId (..)) +import Shibuya.Metrics.Config (MetricsServerConfig (..), defaultConfig) +import Shibuya.Metrics.Cors (corsMiddleware, isAllowedOrigin) +import Shibuya.Metrics.Server (combinedApp, startMetricsServer) +import Shibuya.Metrics.TestSupport (registerFailedProcessor, registerIdleProcessor, withMaster) +import Shibuya.Metrics.WebSocket (newWebSocketState) +import Test.Hspec (Spec, around, describe, it, shouldBe, shouldSatisfy, shouldThrow) + +spec :: Spec +spec = around withMaster $ do + describe "cross-origin HTTP policy" $ do + it "answers allowed preflights with an empty 204 and explicit policy" $ \master -> do + app <- appFor master + forM_ ["GET", "POST", "OPTIONS"] $ \method -> do + response <- send app "OPTIONS" "/metrics" (preflight method) + response.simpleStatus `shouldBe` status204 + response.simpleBody `shouldBe` "" + lookup "Access-Control-Allow-Origin" response.simpleHeaders `shouldBe` Just origin + lookup "Access-Control-Allow-Methods" response.simpleHeaders `shouldBe` Just "GET, POST, OPTIONS" + lookup "Access-Control-Allow-Headers" response.simpleHeaders `shouldBe` Just "Content-Type" + lookup "Access-Control-Max-Age" response.simpleHeaders `shouldBe` Just "600" + lookup "Vary" response.simpleHeaders `shouldBe` Just "Origin, Access-Control-Request-Method, Access-Control-Request-Headers" + lookup "Access-Control-Allow-Credentials" response.simpleHeaders `shouldBe` Nothing + + it "adds headers to every allowed read route, including 404s" $ \master -> do + _ <- registerIdleProcessor master (ProcessorId "known") + app <- appFor master + forM_ ["/metrics", "/metrics/known", "/metrics/prometheus", "/health", "/health/live", "/health/ready", "/metrics/missing", "/ws", "/unknown"] $ \path -> do + plain <- send app "GET" path [] + response <- send app "GET" path [("Origin", origin)] + response.simpleStatus `shouldBe` plain.simpleStatus + response.simpleBody `shouldBe` plain.simpleBody + lookup "Access-Control-Allow-Origin" response.simpleHeaders `shouldBe` Just origin + lookup "Vary" response.simpleHeaders `shouldBe` Just "Origin" + lookup "Access-Control-Allow-Credentials" response.simpleHeaders `shouldBe` Nothing + missing <- send app "GET" "/metrics/missing" [("Origin", origin)] + missing.simpleStatus `shouldBe` status404 + decode missing.simpleBody + `shouldBe` Just + ( object + ["error" .= ("Processor not found" :: Text), "code" .= ("processor_not_found" :: Text), "processor" .= ("missing" :: Text)] + ) + + it "adds allowed-origin headers to unhealthy 503 responses" $ \master -> do + _ <- registerFailedProcessor master (ProcessorId "failed") + app <- appFor master + response <- send app "GET" "/health/ready" [("Origin", origin)] + response.simpleStatus `shouldBe` status503 + lookup "Access-Control-Allow-Origin" response.simpleHeaders `shouldBe` Just origin + + it "refuses disallowed preflights but preserves ordinary read responses" $ \master -> do + app <- appFor master + response <- send app "OPTIONS" "/metrics" [("Origin", "https://evil.example.com"), ("Access-Control-Request-Method", "GET")] + response.simpleStatus `shouldBe` status403 + decode response.simpleBody `shouldBe` Just (object ["error" .= ("Origin not allowed" :: Text), "code" .= ("origin_not_allowed" :: Text)]) + lookup hContentType response.simpleHeaders `shouldBe` Just "application/json" + noAccessHeaders response + lookup "Vary" response.simpleHeaders `shouldBe` Just "Origin, Access-Control-Request-Method, Access-Control-Request-Headers" + plain <- send app "GET" "/metrics" [] + refused <- send app "GET" "/metrics" [("Origin", "https://evil.example.com")] + refused.simpleStatus `shouldBe` status200 + refused.simpleBody `shouldBe` plain.simpleBody + noAccessHeaders refused + lookup "Vary" refused.simpleHeaders `shouldBe` Just "Origin" + lookup "Vary" plain.simpleHeaders `shouldBe` Just "Origin" + + it "validates and echoes requested header tokens" $ \master -> do + app <- appFor master + response <- send app "OPTIONS" "/metrics" (preflight "GET" <> [("Access-Control-Request-Headers", "x-trace-id, Content-Type")]) + response.simpleStatus `shouldBe` status204 + lookup "Access-Control-Allow-Headers" response.simpleHeaders `shouldBe` Just "x-trace-id, Content-Type" + forM_ ["", "bad header", "x-trace-id,", "x-trace-id\r\nInjected: yes", "*"] $ \headers -> do + invalid <- send app "OPTIONS" "/metrics" (preflight "GET" <> [("Access-Control-Request-Headers", headers)]) + invalid.simpleStatus `shouldBe` status403 + noAccessHeaders invalid + repeatedEmpty <- + send + app + "OPTIONS" + "/metrics" + (preflight "GET" <> [("Access-Control-Request-Headers", ""), ("Access-Control-Request-Headers", "x-trace-id")]) + repeatedEmpty.simpleStatus `shouldBe` status403 + noAccessHeaders repeatedEmpty + + it "rejects forbidden, malformed, and repeated requested methods" $ \master -> do + app <- appFor master + forM_ ["DELETE", "get", "GET, POST", ""] $ \method -> do + response <- send app "OPTIONS" "/metrics" (preflight method) + response.simpleStatus `shouldBe` status403 + decode response.simpleBody `shouldBe` Just (object ["error" .= ("CORS method not allowed" :: Text), "code" .= ("cors_method_not_allowed" :: Text)]) + noAccessHeaders response + repeated <- send app "OPTIONS" "/metrics" (preflight "GET" <> [("Access-Control-Request-Method", "POST")]) + repeated.simpleStatus `shouldBe` status403 + + it "fails closed for malformed or multiple Origin values" $ \master -> do + app <- appFor master + forM_ [[("Origin", "null")], [("Origin", "")], [("Origin", "https://ops.example.com https://evil.example.com")], [("Origin", "https://ops.example.com,https://evil.example.com")], [("Origin", origin <> "/")], [("Origin", origin), ("origin", origin)]] $ \origins -> do + response <- send app "OPTIONS" "/metrics" (origins <> [("Access-Control-Request-Method", "GET")]) + response.simpleStatus `shouldBe` status403 + noAccessHeaders response + ordinary <- send app "GET" "/metrics" origins + noAccessHeaders ordinary + + it "routes plain OPTIONS normally when preflight headers are absent" $ \master -> do + app <- appFor master + response <- send app "OPTIONS" "/unknown" [("Origin", origin)] + response.simpleStatus `shouldBe` status404 + methodOnly <- send app "OPTIONS" "/unknown" [("Access-Control-Request-Method", "GET")] + methodOnly.simpleStatus `shouldBe` status404 + + it "merges Vary tokens and replaces any stale CORS policy headers" $ \_ -> do + let inner _ respond = + respond $ + responseLBS + status200 + [("Vary", "Accept-Encoding, origin"), ("vary", "Accept-Language"), ("Access-Control-Allow-Origin", "*"), ("Access-Control-Allow-Credentials", "true")] + "body" + app = corsMiddleware ["https://ops.example.com"] inner + forM_ [[], [("Origin", origin)], [("Origin", "https://evil.example.com")]] $ \headers -> do + response <- send app "GET" "/" headers + lookup "Vary" response.simpleHeaders `shouldBe` Just "Accept-Encoding, origin, Accept-Language" + length (filter ((== "Vary") . fst) response.simpleHeaders) `shouldBe` 1 + lookup "Access-Control-Allow-Credentials" response.simpleHeaders `shouldBe` Nothing + allowed <- send app "GET" "/" [("Origin", origin)] + filter ((== "Access-Control-Allow-Origin") . fst) allowed.simpleHeaders `shouldBe` [("Access-Control-Allow-Origin", origin)] + refused <- send app "GET" "/" [("Origin", "https://evil.example.com")] + noAccessHeaders refused + wildcard <- send (corsMiddleware ["https://ops.example.com"] (\_ respond -> respond $ responseLBS status200 [("Vary", "*")] "body")) "GET" "/" [] + lookup "Vary" wildcard.simpleHeaders `shouldBe` Just "*" + + it "compares origin bytes exactly without normalizing or prefix matching" $ \_ -> do + let allowed = isAllowedOrigin ["https://ops.example.com"] + allowed origin `shouldBe` True + forM_ ["https://OPS.example.com", "https://ops.example.com:443", "https://ops.example.com.evil", "https://ops.example.com/"] $ \other -> allowed other `shouldBe` False + + it "rejects invalid configured origins before opening a listener" $ \master -> + forM_ ["*", "", "https://ops.example.com/", "ops.example.com", "1https://ops.example.com", "://ops.example.com", "https://", "https://ops.example.com?x=1", "https://ops.example.com#x", "https://ops. example.com", "https://ops.example.com\t"] $ \invalid -> + startMetricsServer defaultConfig {corsAllowedOrigins = [invalid]} master + `shouldThrow` (\err -> "corsAllowedOrigins" `isInfixOf` show (err :: IOException)) + +origin :: ByteString +origin = "https://ops.example.com" + +preflight :: ByteString -> RequestHeaders +preflight method = [("Origin", origin), ("Access-Control-Request-Method", method)] + +send :: Application -> ByteString -> ByteString -> RequestHeaders -> IO SResponse +send app method path headers = + runSession + (request $ setPath defaultRequest {requestMethod = method, requestHeaders = headers} path) + app + +appFor :: Master -> IO Application +appFor master = do + wsState <- newWebSocketState defaultConfig.wsMaxConnections + pure $ combinedApp defaultConfig {corsAllowedOrigins = ["https://ops.example.com"]} master wsState [] + +noAccessHeaders :: SResponse -> IO () +noAccessHeaders response = + response.simpleHeaders + `shouldSatisfy` (not . any (\(name, _) -> "access-control-" `isInfixOf` map toLower (show name))) diff --git a/test/Shibuya/Metrics/HealthSpec.hs b/test/Shibuya/Metrics/HealthSpec.hs --- a/test/Shibuya/Metrics/HealthSpec.hs +++ b/test/Shibuya/Metrics/HealthSpec.hs @@ -22,6 +22,7 @@ ProcessorState (..), beginProcessing, finishProcessing, + markPaused, newMetricsHandle, newMetricsHandleWithClock, sampleMetrics, @@ -59,7 +60,7 @@ `shouldReturn` ReadinessStatus { ready = True, application = ConfiguredEmpty, - processors = ProcessorHealth {total = 0, healthy = 0, failed = 0, stuck = 0}, + processors = ProcessorHealth {total = 0, healthy = 0, failed = 0, stuck = 0, paused = 0}, dependencies = [] } @@ -68,7 +69,7 @@ readiness <- checkReadiness defaultHealthConfig master [] readiness.ready `shouldBe` False readiness.application `shouldBe` Running - readiness.processors `shouldBe` ProcessorHealth {total = 1, healthy = 0, failed = 1, stuck = 0} + readiness.processors `shouldBe` ProcessorHealth {total = 1, healthy = 0, failed = 1, stuck = 0, paused = 0} (detailed, metrics) <- checkDetailedHealth defaultHealthConfig master [] detailed `shouldBe` readiness @@ -127,6 +128,23 @@ readiness.ready `shouldBe` True readiness.processors.stuck `shouldBe` 0 + it "counts drained pauses separately while preserving stuck readiness failures" $ \master -> do + now <- getCurrentTime + clock <- newIORef 0 + handle <- newMetricsHandleWithClock (readIORef clock) (addUTCTime (-120) now) + runEff $ registerProcessor master (ProcessorId "paused") handle + markPaused handle now + drained <- checkReadiness defaultHealthConfig master [] + drained.ready `shouldBe` True + drained.processors.paused `shouldBe` 1 + _ <- beginProcessing handle 1 + _ <- sampleMetrics handle + writeIORef clock (seconds 120) + stuck <- checkReadiness defaultHealthConfig master [] + stuck.ready `shouldBe` False + stuck.processors.stuck `shouldBe` 1 + stuck.processors.paused `shouldBe` 0 + it "reports a genuinely non-progressing handler stuck" $ \master -> do now <- getCurrentTime clock <- newIORef 0 @@ -155,7 +173,7 @@ readiness <- checkReadiness defaultHealthConfig master [] readiness.ready `shouldBe` False readiness.application `shouldBe` ApplicationFailed - readiness.processors `shouldBe` ProcessorHealth {total = 1, healthy = 0, failed = 1, stuck = 0} + readiness.processors `shouldBe` ProcessorHealth {total = 1, healthy = 0, failed = 1, stuck = 0, paused = 0} it "reports draining and stopped masters unavailable" $ \master -> do runEff $ markMasterDraining master diff --git a/test/Shibuya/Metrics/JSONSpec.hs b/test/Shibuya/Metrics/JSONSpec.hs --- a/test/Shibuya/Metrics/JSONSpec.hs +++ b/test/Shibuya/Metrics/JSONSpec.hs @@ -7,5 +7,5 @@ spec :: Spec spec = describe "JSON wire contract" $ - it "matches the golden encoding for all four processor states" $ + it "matches the golden encoding for all four processor states and opt-in detail" $ assertGolden "processor-metrics.json.golden" (encode fixtureMetrics <> "\n") diff --git a/test/Shibuya/Metrics/OutboundQueueSpec.hs b/test/Shibuya/Metrics/OutboundQueueSpec.hs new file mode 100644 --- /dev/null +++ b/test/Shibuya/Metrics/OutboundQueueSpec.hs @@ -0,0 +1,67 @@ +module Shibuya.Metrics.OutboundQueueSpec (spec) where + +import Control.Monad (forM_) +import Data.ByteString qualified as BS +import Shibuya.Metrics.Outbound qualified as Q +import Test.Hspec (Spec, describe, it, shouldBe, shouldSatisfy) + +spec :: Spec +spec = describe "bounded outbound queue" $ do + it "preserves frame and byte bounds and exact counters over exhaustive transition traces" $ do + let actions = + [ Add Q.Required 1, + Add Q.Required 7, + Add Q.Update 3, + Add Q.Update 12, + Add Q.Ephemeral 2, + Pop, + Discard + ] + step q action = case action of + Add keep size -> snd $ Q.enqueue (Q.Limits 3 10) [value keep size] q + Pop -> maybe q snd $ Q.pop q + Discard -> Q.discardUpdates q + walk depth q = do + Q.queuedFrames q `shouldBe` length (Q.entries q) + Q.queuedBytes q `shouldBe` sum (map (BS.length . Q.payload) $ Q.entries q) + Q.queuedFrames q `shouldSatisfy` (<= 3) + Q.queuedBytes q `shouldSatisfy` (<= 10) + if depth == (0 :: Int) then pure () else forM_ actions (walk (depth - 1) . step q) + walk 5 Q.empty + + it "evicts updates oldest-first while retaining required order" $ do + let a = value Q.Update 1 + b = value Q.Required 2 + c = value Q.Update 3 + (_, q) = Q.enqueue (Q.Limits 3 10) [a, b, c] Q.empty + (outcome, next) = Q.enqueue (Q.Limits 3 10) [value Q.Required 4] q + outcome `shouldBe` Q.DroppedUpdates + map (BS.length . Q.payload) (Q.entries next) `shouldBe` [2, 3, 4] + + it "refuses required saturation and oversized snapshots without dropping required frames" $ do + let (_, q) = Q.enqueue (Q.Limits 2 10) [value Q.Required 4, value Q.Required 4] Q.empty + Q.enqueue (Q.Limits 2 10) [value Q.Required 1] q `shouldBe` (Q.MustClose, q) + Q.enqueue (Q.Limits 2 10) [value Q.Required 11] Q.empty `shouldBe` (Q.MustClose, Q.empty) + + it "admits recovery batches wholly and consecutively after earlier required frames" $ do + let (_, q) = Q.enqueue (Q.Limits 3 10) [value Q.Required 1, value Q.Update 3] Q.empty + batch = [value Q.Required 4, value Q.Required 5] + (result, next) = Q.enqueue (Q.Limits 3 10) batch (Q.discardUpdates q) + result `shouldBe` Q.Enqueued + Q.entries next `shouldBe` value Q.Required 1 : batch + Q.enqueue (Q.Limits 2 10) batch q `shouldBe` (Q.MustClose, q) + + it "marks rejected oversized updates for recovery, but not lost pongs" $ do + fst (Q.enqueue (Q.Limits 2 10) [value Q.Update 11] Q.empty) `shouldBe` Q.DroppedUpdates + fst (Q.enqueue (Q.Limits 2 10) [value Q.Ephemeral 11] Q.empty) `shouldBe` Q.Enqueued + + it "removes unsent old-generation updates without deleting required snapshots" $ do + let snapshot = Q.frame 2 Q.Required "snapshot" + update = Q.frame 1 Q.Update "old" + (_, q) = Q.enqueue (Q.Limits 2 30) [update, snapshot] Q.empty + Q.entries (Q.discardUpdates q) `shouldBe` [snapshot] + +value :: Q.Retention -> Int -> Q.Frame +value keep size = Q.frame 0 keep (BS.replicate size 97) + +data Action = Add Q.Retention Int | Pop | Discard diff --git a/test/Shibuya/Metrics/ServerSpec.hs b/test/Shibuya/Metrics/ServerSpec.hs --- a/test/Shibuya/Metrics/ServerSpec.hs +++ b/test/Shibuya/Metrics/ServerSpec.hs @@ -1,19 +1,28 @@ module Shibuya.Metrics.ServerSpec (spec) where import Control.Exception (SomeException, throwIO) -import Data.Aeson (decode, object, (.=)) +import Data.Aeson (Value (..), decode, object, (.=)) +import Data.Aeson.KeyMap qualified as KM import Data.ByteString (ByteString) import Data.ByteString.Lazy qualified as LBS +import Data.Char (toLower) +import Data.List (isInfixOf) +import Effectful (runEff) import Network.HTTP.Types (Status, hContentType, status200, status404, status503) -import Network.Wai (Application) -import Network.Wai.Test (SResponse (..)) +import Network.Wai (Application, requestHeaders) +import Network.Wai.Test (SResponse (..), defaultRequest, request, runSession, setPath) import Shibuya.App (Master) import Shibuya.Core.Metrics (ProcessorId (..)) +import Shibuya.Core.Metrics qualified as Metrics +import Shibuya.Core.Types (Cursor (..)) +import Shibuya.Internal.Runner.Master (registerProcessor) import Shibuya.Metrics.Config (MetricsServerConfig (..), defaultConfig) +import Shibuya.Metrics.Error (errorResponse) import Shibuya.Metrics.Health (DependencyCheck, DependencyStatus (..)) import Shibuya.Metrics.Server (combinedApp, startMetricsServer) import Shibuya.Metrics.TestSupport - ( getResponse, + ( fixedTime, + getResponse, registerIdleProcessor, withMaster, ) @@ -26,11 +35,65 @@ it "defaults to a loopback-only listener" $ \_ -> defaultConfig.host `shouldBe` "127.0.0.1" + it "validates frame, byte, and send budgets before opening a socket" $ \master -> + mapM_ + (\config -> startMetricsServer config master `shouldThrow` anyException) + [ defaultConfig {wsMaxQueuedFrames = 0}, + defaultConfig {wsMaxQueuedFrames = 1}, + defaultConfig {wsMaxQueuedBytes = 0}, + defaultConfig {wsSendTimeoutUs = 0} + ] + it "rejects nonpositive WebSocket resource limits" $ \master -> startMetricsServer defaultConfig {wsMaxSubscriptions = 0} master `shouldThrow` (anyException :: SomeException -> Bool) + describe "structured error response" $ do + it "protects reserved keys while retaining context and content type" $ \_ -> do + let app _ respond = + respond $ + errorResponse + status404 + "Not found" + "not_found" + ["error" .= ("override" :: String), "code" .= ("override" :: String), "processor" .= ("missing" :: String)] + response <- getResponse app "/" + response.simpleStatus `shouldBe` status404 + lookup hContentType response.simpleHeaders `shouldBe` Just "application/json" + decode response.simpleBody + `shouldBe` Just + ( object + ["error" .= ("Not found" :: String), "code" .= ("not_found" :: String), "processor" .= ("missing" :: String)] + ) + describe "combinedApp HTTP routes" $ do + it "exposes progress only after a detailed cursor acknowledgement" $ \master -> do + h <- Metrics.newMetricsHandleWith (Metrics.MetricsDetailed Metrics.defaultDetailOptions) 1 fixedTime + runEff $ registerProcessor master (ProcessorId "progressing") h + _ <- registerIdleProcessor master (ProcessorId "idle") + Metrics.recordAcknowledgedCursor h Nothing (Just (CursorInt 4210)) + app <- appFor defaultConfig master [] + yes <- getResponse app "/metrics/progressing" + no <- getResponse app "/metrics/idle" + let hasProgress body = case decode body of + Just (Object obj) -> KM.member "progress" obj + _ -> False + hasProgress yes.simpleBody `shouldBe` True + hasProgress no.simpleBody `shouldBe` False + + it "preserves headerless defaults on every read route with an Origin" $ \master -> do + _ <- registerIdleProcessor master (ProcessorId "known") + app <- appFor defaultConfig master [] + mapM_ + ( \path -> do + plain <- getResponse app path + origin <- runSession (request $ setPath defaultRequest {requestHeaders = [("Origin", "https://ops.example.com")]} path) app + origin.simpleHeaders `shouldBe` plain.simpleHeaders + origin.simpleHeaders `shouldSatisfy` (not . any (\(name, _) -> "access-control-" `isInfixOf` map toLower (show name))) + lookup "Vary" origin.simpleHeaders `shouldBe` Nothing + ) + ["/metrics", "/metrics/known", "/metrics/missing", "/metrics/prometheus", "/health", "/health/live", "/health/ready", "/ws", "/unknown"] + it "serves every enabled JSON, health, Prometheus, and WebSocket path" $ \master -> do _ <- registerIdleProcessor master (ProcessorId "known") app <- appFor defaultConfig master [] @@ -45,7 +108,7 @@ wsResponse <- getResponse app "/ws" wsResponse.simpleStatus `shouldBe` status404 decode wsResponse.simpleBody - `shouldBe` Just (object ["error" .= ("WebSocket endpoint - use ws:// protocol" :: String)]) + `shouldBe` Just (object ["error" .= ("WebSocket endpoint - use ws:// protocol" :: String), "code" .= ("websocket_upgrade_required" :: String)]) it "returns the published JSON error for an unknown processor" $ \master -> do app <- appFor defaultConfig master [] @@ -55,7 +118,8 @@ `shouldBe` Just ( object [ "error" .= ("Processor not found" :: String), - "processor" .= ("missing" :: String) + "processor" .= ("missing" :: String), + "code" .= ("processor_not_found" :: String) ] ) @@ -64,7 +128,7 @@ response <- getResponse app "/unknown" response.simpleStatus `shouldBe` status404 decode response.simpleBody - `shouldBe` Just (object ["error" .= ("Not found" :: String)]) + `shouldBe` Just (object ["error" .= ("Not found" :: String), "code" .= ("not_found" :: String)]) it "returns 404 for every JSON route when JSON endpoints are disabled" $ \master -> do app <- appFor defaultConfig {enableJSON = False} master [] @@ -81,7 +145,7 @@ response <- getResponse app "/ws" response.simpleStatus `shouldBe` status404 decode response.simpleBody - `shouldBe` Just (object ["error" .= ("Not found" :: String)]) + `shouldBe` Just (object ["error" .= ("Not found" :: String), "code" .= ("not_found" :: String)]) it "returns 503 from readiness and detailed health for an unhealthy dependency" $ \master -> do app <- appFor defaultConfig master [failingDependency] diff --git a/test/Shibuya/Metrics/TestSupport.hs b/test/Shibuya/Metrics/TestSupport.hs --- a/test/Shibuya/Metrics/TestSupport.hs +++ b/test/Shibuya/Metrics/TestSupport.hs @@ -17,6 +17,7 @@ import Control.Monad (replicateM_) import Data.ByteString (ByteString) import Data.ByteString.Lazy qualified as LBS +import Data.IORef (newIORef, readIORef, writeIORef) import Data.Map.Strict qualified as Map import Data.Time.Calendar (fromGregorian) import Data.Time.Clock (UTCTime (..)) @@ -40,6 +41,8 @@ incrementReceived, newMetricsHandle, ) +import Shibuya.Core.Metrics qualified as Metrics +import Shibuya.Core.Types (Cursor (..), MessageId (..)) import Shibuya.Internal.Runner.Master ( markMasterRunning, registerProcessor, @@ -54,8 +57,20 @@ fixtureMetrics :: MetricsMap fixtureMetrics = Map.fromList - [ (ProcessorId "failed", fixtureProcessor (Failed "boom" (at 90)) 20), + [ ( ProcessorId "detailed", + (fixtureProcessor (Processing (InFlightInfo 1 4) (at 40) (at 40)) 40) + { latency = Just (Metrics.LatencySummary 3 0.42 [Metrics.LatencyBucket bound (if bound < 0.1 then 0 else if bound < 0.25 then 1 else 3) | bound <- Metrics.latencyBucketBoundsSeconds]), + oldestInFlight = Just (Metrics.OldestInFlight (MessageId "m-7") (at 40) 12.5) + } + ), + ( ProcessorId "progressing", + (fixtureProcessor (Processing (InFlightInfo 1 4) (at 40) (at 40)) 50) + { progress = Just (Metrics.ProgressSnapshot (Just (CursorInt 4210)) (Map.fromList [("p-0", CursorInt 4210), ("p-1", CursorText "0198f2f3")]) False) + } + ), + (ProcessorId "failed", fixtureProcessor (Failed "boom" (at 90)) 20), (ProcessorId "idle", fixtureProcessor Idle 0), + (ProcessorId "paused", fixtureProcessor (Paused (InFlightInfo 1 2) (at 45) (at 44)) 60), (ProcessorId "processing", fixtureProcessor (Processing (InFlightInfo 2 4) (at 60) (at 75)) 10), (ProcessorId "stopped", fixtureProcessor Stopped 30) ] @@ -68,6 +83,9 @@ { state, stats = StreamStats (offset + 1) (offset + 2) (offset + 3), batch = BatchStats (offset + 4) (offset + 5) (offset + 6) (offset + 7) (offset + 8) (offset + 9), + latency = Nothing, + oldestInFlight = Nothing, + progress = Nothing, startedAt = fixedTime } @@ -95,6 +113,15 @@ registerPrometheusFixtures :: Master -> IO () registerPrometheusFixtures master = do + clock <- newIORef 0 + detailed <- Metrics.newMetricsHandleWithClockAndDetail (readIORef clock) (Metrics.MetricsDetailed Metrics.defaultDetailOptions) 1 fixedTime + runEff $ registerProcessor master (ProcessorId "detailed") detailed + first <- Metrics.claimInFlightSlot detailed (MessageId "m-1") + writeIORef clock 3_000_000 + Metrics.releaseInFlightSlot detailed first + writeIORef clock 5_000_000 + _ <- Metrics.claimInFlightSlot detailed (MessageId "m-2") + writeIORef clock 17_500_000 _ <- registerIdleProcessor master (ProcessorId "idle") processing <- registerIdleProcessor master (ProcessorId "processing") @@ -105,6 +132,10 @@ incrementReceived failed _ <- beginProcessing failed 1 finishProcessing failed (Left "fixture failure") + + paused <- registerIdleProcessor master (ProcessorId "paused") + _ <- beginProcessing paused 2 + Metrics.markPaused paused fixedTime stopped <- registerIdleProcessor master (ProcessorId "stopped") replicateM_ 4 $ incrementReceived stopped diff --git a/test/Shibuya/Metrics/TypesSpec.hs b/test/Shibuya/Metrics/TypesSpec.hs --- a/test/Shibuya/Metrics/TypesSpec.hs +++ b/test/Shibuya/Metrics/TypesSpec.hs @@ -1,9 +1,11 @@ module Shibuya.Metrics.TypesSpec (spec) where -import Data.Aeson (FromJSON, ToJSON, Value, eitherDecode, encode, object, toJSON, (.=)) +import Data.Aeson (FromJSON, ToJSON, Value, eitherDecode, encode, object, toJSON, withObject, (.:), (.=)) +import Data.Aeson.Types (parseEither) import Data.Map.Strict qualified as Map -import Shibuya.Core.Metrics (ProcessorId (..)) -import Shibuya.Metrics.TestSupport (fixtureMetrics) +import Data.Text (Text) +import Shibuya.Core.Metrics (InFlightInfo (..), ProcessorId (..), ProcessorState (..)) +import Shibuya.Metrics.TestSupport (fixedTime, fixtureMetrics) import Shibuya.Metrics.Types ( ClientMessage (..), ProcessorTerminalStatus (..), @@ -13,6 +15,7 @@ spec :: Spec spec = do + describe "paused state JSON" $ messageCase (Paused (InFlightInfo 1 2) fixedTime fixedTime) (object ["status" .= ("paused" :: String), "inFlight" .= (1 :: Int), "maxConcurrency" .= (2 :: Int), "pausedAt" .= fixedTime, "lastProgress" .= fixedTime]) describe "ClientMessage JSON contract" $ do messageCase SubscribeAll $ object ["type" .= ("subscribe_all" :: String)] messageCase (Subscribe [ProcessorId "alpha"]) $ @@ -51,7 +54,23 @@ "error" .= ("boom" :: String), "messageId" .= (Just "message-1" :: Maybe String) ] + messageCase (ServerError "invalid_message" "bad input") $ + object ["type" .= ("error" :: String), "code" .= ("invalid_message" :: String), "message" .= ("bad input" :: String)] messageCase Goodbye $ object ["type" .= ("goodbye" :: String)] + + it "decodes a legacy update with no detail members" $ do + let legacy = ProcessorUpdate (ProcessorId "processing") processing + eitherDecode "{\"type\":\"update\",\"processor\":\"processing\",\"metrics\":{\"batch\":{\"batchedMessages\":15,\"batchesEmitted\":14,\"flushTriggered\":19,\"partialFailures\":16,\"sizeTriggered\":17,\"timeoutTriggered\":18},\"startedAt\":\"2026-09-20T03:25:45Z\",\"state\":{\"inFlight\":2,\"lastActivity\":\"2026-09-20T00:01:00Z\",\"lastProgress\":\"2026-09-20T00:01:15Z\",\"maxConcurrency\":4,\"status\":\"processing\"},\"stats\":{\"failed\":13,\"processed\":12,\"received\":11}}}" `shouldBe` Right legacy + + it "characterizes strict legacy decoders and extension-tolerant consumers" $ do + let err = ServerError "queue_overflow" "resnapshot" + legacy = withObject "LegacyServerMessage" $ \v -> do + tag <- v .: "type" + if (tag :: Text) `elem` ["snapshot", "update", "terminal", "pong", "goodbye"] + then pure tag + else fail "unknown legacy tag" + parseEither legacy (toJSON err) `shouldSatisfy` isLeft + eitherDecode @Value (encode err) `shouldBe` Right (toJSON err) it "rejects an unknown server message tag" $ (eitherDecode "{\"type\":\"unknown\"}" :: Either String ServerMessage) diff --git a/test/Shibuya/Metrics/WebSocketSpec.hs b/test/Shibuya/Metrics/WebSocketSpec.hs --- a/test/Shibuya/Metrics/WebSocketSpec.hs +++ b/test/Shibuya/Metrics/WebSocketSpec.hs @@ -1,15 +1,19 @@ module Shibuya.Metrics.WebSocketSpec (spec) where -import Control.Concurrent (MVar, newEmptyMVar, putMVar, takeMVar) +import Control.Concurrent (MVar, newEmptyMVar, putMVar, takeMVar, threadDelay) import Control.Concurrent.Async (async, wait) -import Control.Concurrent.STM (atomically, check, readTVar) +import Control.Concurrent.STM (TQueue, TVar, atomically, check, newTQueueIO, newTVarIO, orElse, readTQueue, readTVar, writeTQueue, writeTVar) import Control.Exception (SomeException, throwIO, try) -import Data.Aeson (eitherDecode, encode) +import Control.Monad (replicateM_, when) +import Data.Aeson (Value, eitherDecode, encode, object, (.=)) import Data.ByteString.Lazy (ByteString) -import Data.IORef (atomicModifyIORef', newIORef) +import Data.ByteString.Lazy qualified as LBS +import Data.IORef (atomicModifyIORef', newIORef, readIORef, writeIORef) import Data.Map.Strict qualified as Map +import Data.String (fromString) import Effectful (runEff) import Network.Wai.Handler.Warp qualified as Warp +import Network.Wai.Handler.WebSockets qualified as WaiWS import Network.WebSockets qualified as WS import Shibuya.App (Master, getAllMetricsIO) import Shibuya.Core.Metrics @@ -21,9 +25,12 @@ incrementReceived, newMetricsHandleWithClock, ) +import Shibuya.Core.Metrics qualified as Metrics +import Shibuya.Core.Types (Cursor (..), MessageId (..)) import Shibuya.Internal.Runner.Master (markProcessorFailedIO, registerProcessor, unregisterProcessor) import Shibuya.Metrics.Config (MetricsServerConfig (..), defaultConfig) -import Shibuya.Metrics.Server (combinedApp) +import Shibuya.Metrics.Outbound qualified as Out +import Shibuya.Metrics.Server (combinedApp, stopMetricsServer, withMetricsServer) import Shibuya.Metrics.TestSupport ( fixedTime, registerIdleProcessor, @@ -35,9 +42,12 @@ ServerMessage (..), ) import Shibuya.Metrics.WebSocket - ( WebSocketState (..), + ( WebSocketHooks (..), + WebSocketState (..), + defaultWebSocketHooks, newWebSocketState, shutdownWebSockets, + websocketAppWithHooks, ) import System.Timeout (timeout) import Test.Hspec @@ -48,12 +58,278 @@ expectationFailure, it, shouldBe, + shouldSatisfy, shouldThrow, ) spec :: Spec spec = around withMaster $ do + describe "cross-origin WebSocket upgrades" $ do + it "accepts an allowed Origin with an initial snapshot" $ \master -> + withServer originConfig master $ \port -> + clientWithOrigin port [("Origin", "https://ops.example.com")] $ \conn -> do + expected <- MetricsSnapshot <$> currentSnapshot master + receiveServer conn `shouldReturn` expected + + it "refuses forbidden or multiple Origins with HTTP 403 before acquiring a slot" $ \master -> do + wsState <- newWebSocketState 1 + let app = combinedApp originConfig master wsState [] + Warp.testWithApplication (pure app) $ \port -> do + mapM_ + ( \headers -> do + clientWithOrigin port headers receiveServer + `shouldThrow` isForbiddenHandshake + atomically (readTVar wsState.connectionCount) `shouldReturn` 0 + ) + [[("Origin", "https://evil.example.com")], [("Origin", "null")], [("Origin", "")], [("Origin", "https://ops.example.com"), ("origin", "https://ops.example.com")]] + clientWithOrigin port [("Origin", "https://ops.example.com")] $ \conn -> do + expected <- MetricsSnapshot <$> currentSnapshot master + receiveServer conn `shouldReturn` expected + + it "allows clients without Origin when the policy is enabled" $ \master -> + withServer originConfig master $ \port -> + clientWithOrigin port [] $ \conn -> do + expected <- MetricsSnapshot <$> currentSnapshot master + receiveServer conn `shouldReturn` expected + + it "allows any Origin when the policy is disabled" $ \master -> + withServer defaultConfig master $ \port -> + clientWithOrigin port [("Origin", "https://evil.example.com")] $ \conn -> do + expected <- MetricsSnapshot <$> currentSnapshot master + receiveServer conn `shouldReturn` expected + + describe "detailed metrics updates" $ do + it "publishes a cursor change without counter changes" $ \master -> do + handle <- Metrics.newMetricsHandleWith (Metrics.MetricsDetailed Metrics.defaultDetailOptions) 1 fixedTime + runEff $ registerProcessor master (ProcessorId "progress") handle + withServer defaultConfig {wsPushIntervalUs = 10_000} master $ \port -> WS.runClient "127.0.0.1" port "/ws" $ \conn -> do + _ <- receiveServer conn + Metrics.recordAcknowledgedCursor handle Nothing (Just (CursorInt 1)) + update <- timeout 1_000_000 (receiveServer conn) + case update of + Just (ProcessorUpdate _ pm) -> fmap (.latest) pm.progress `shouldBe` Just (Just (CursorInt 1)) + other -> expectationFailure (show other) + + it "suppresses age-only changes and publishes release with latency" $ \master -> do + clock <- newIORef 0 + handle <- + Metrics.newMetricsHandleWithClockAndDetail + (readIORef clock) + (Metrics.MetricsDetailed Metrics.defaultDetailOptions) + 1 + fixedTime + runEff $ registerProcessor master (ProcessorId "detail") handle + withServer fastConfig master $ \port -> WS.runClient "127.0.0.1" port "/ws" $ \conn -> do + _ <- receiveServer conn + slot <- Metrics.claimInFlightSlot handle (MessageId "blocked") + _ <- Metrics.beginProcessing handle 1 + timeout 1_000_000 (receiveServer conn) >>= \case + Just (ProcessorUpdate _ pm) -> fmap (.messageId) pm.oldestInFlight `shouldBe` Just (MessageId "blocked") + other -> expectationFailure (show other) + writeIORef clock 1_000_000_000 + timeout 50_000 (receiveServer conn) `shouldReturn` Nothing + writeIORef clock 2_000_000_000 + timeout 50_000 (receiveServer conn) `shouldReturn` Nothing + Metrics.finishProcessing handle (Right Metrics.CountProcessed) + Metrics.releaseInFlightSlot handle slot + timeout 1_000_000 (receiveServer conn) >>= \case + Just (ProcessorUpdate _ pm) -> do + pm.oldestInFlight `shouldBe` Nothing + fmap (.count) pm.latency `shouldBe` Just 1 + other -> expectationFailure (show other) + + describe "controlled outbound delivery" $ do + it "keeps sampling during a blocked write and sends an adjacent recovery pair" $ \master -> do + handle <- registerIdleProcessor master (ProcessorId "alpha") + withControlledServer fastConfig {wsMaxQueuedFrames = 2} master $ \port control -> + WS.runClient "127.0.0.1" port "/ws" $ \conn -> do + _ <- receiveServer conn + atomically $ writeTVar control.gate True + incrementReceived handle + _ <- tick control + awaitBlocked control + replicateM_ 3 $ incrementReceived handle >> tick control + recovery <- tick control + Out.queuedFrames recovery `shouldBe` 2 + map (eitherDecode @ServerMessage . LBS.fromStrict . Out.payload) (Out.entries recovery) + `shouldSatisfy` (\case [Right (ServerError "queue_overflow" _), Right (MetricsSnapshot _)] -> True; _ -> False) + atomically $ writeTVar control.gate False + _ <- receiveServer conn -- the single already-popped update + receiveServer conn `shouldReturn` ServerError "queue_overflow" "Updates dropped; replacing metrics snapshot" + expected <- MetricsSnapshot <$> currentSnapshot master + receiveServer conn `shouldReturn` expected + WS.sendTextData conn $ encode Ping + receiveServer conn `shouldReturn` Pong + + it "closes a malformed-frame flood with bounded required retention and slot recovery" $ \master -> + withControlledServer fastConfig {wsMaxQueuedFrames = 2, wsSendTimeoutUs = 50_000} master $ \port control -> do + WS.runClient "127.0.0.1" port "/ws" $ \conn -> do + _ <- receiveServer conn + atomically $ writeTVar control.gate True + WS.sendTextData conn $ encode Ping + awaitBlocked control + replicateM_ 8 $ WS.sendTextData conn ("{" :: ByteString) + awaitReleased control.state + atomically (readTVar control.state.connectionCount) `shouldReturn` 0 + + it "refuses an initial snapshot larger than the byte budget and restores its slot" $ \master -> + withControlledServer fastConfig {wsMaxQueuedBytes = 1} master $ \port control -> do + WS.runClient + "127.0.0.1" + port + "/ws" + ( \conn -> + (WS.receiveData conn :: IO ByteString) `shouldThrow` isCapacityClose + ) + awaitReleased control.state + + it "flushes shutdown within one deadline even while a write is blocked" $ \master -> + withControlledServer fastConfig {wsSendTimeoutUs = 50_000} master $ \port control -> + WS.runClient "127.0.0.1" port "/ws" $ \conn -> do + _ <- receiveServer conn + atomically $ writeTVar control.gate True + WS.sendTextData conn $ encode Ping + awaitBlocked control + shutdownWebSockets control.state + awaitReleased control.state + + it "orders shutdown goodbye after admitted errors" $ \master -> + withControlledServer fastConfig {wsMaxQueuedFrames = 3} master $ \port control -> + WS.runClient "127.0.0.1" port "/ws" $ \conn -> do + _ <- receiveServer conn + atomically $ writeTVar control.gate True + WS.sendTextData conn $ encode Ping + awaitBlocked control + atomically $ readTQueue control.received + replicateM_ 2 $ do + WS.sendTextData conn ("{\"type\":\"bogus\"}" :: ByteString) + atomically $ readTQueue control.received + shutdownWebSockets control.state + timeout 1_000_000 (atomically $ readTQueue control.produced) >>= \case + Just q -> Out.queuedFrames q `shouldBe` 3 + Nothing -> expectationFailure "shutdown did not queue goodbye" + atomically $ writeTVar control.gate False + receiveServer conn `shouldReturn` Pong + replicateM_ 2 $ receiveServer conn `shouldReturn` ServerError "invalid_message" "Error in $: Unknown message type: bogus" + receiveServer conn `shouldReturn` Goodbye + (WS.receiveData conn :: IO ByteString) `shouldThrow` isNormalClose + + it "removes unsent old updates at an unsubscribe snapshot boundary" $ \master -> do + handle <- registerIdleProcessor master (ProcessorId "alpha") + withControlledServer fastConfig master $ \port control -> + WS.runClient "127.0.0.1" port "/ws" $ \conn -> do + _ <- receiveServer conn + atomically $ writeTVar control.gate True + incrementReceived handle + _ <- tick control + awaitBlocked control + incrementReceived handle + _ <- tick control + WS.sendTextData conn $ encode $ Unsubscribe [ProcessorId "alpha"] + timeout 1_000_000 (atomically $ readTQueue control.received) `shouldReturn` Just () + WS.sendTextData conn $ encode Ping + atomically $ writeTVar control.gate False + _ <- receiveServer conn + receiveServer conn `shouldReturn` MetricsSnapshot Map.empty + receiveServer conn `shouldReturn` Pong + _ <- tick control + timeout 30_000 (receiveServer conn) `shouldReturn` Nothing + + it "stops ordinary deltas when terminal admission drops an update" $ \master -> do + alpha <- registerIdleProcessor master (ProcessorId "alpha") + beta <- registerIdleProcessor master (ProcessorId "beta") + withControlledServer fastConfig {wsMaxQueuedFrames = 3} master $ \port control -> + WS.runClient "127.0.0.1" port "/ws" $ \conn -> do + _ <- receiveServer conn + atomically $ writeTVar control.gate True + incrementReceived alpha + _ <- tick control + awaitBlocked control + incrementReceived alpha + _ <- tick control + incrementReceived beta + _ <- tick control + incrementReceived alpha + _ <- tick control + markProcessorFailedIO master (ProcessorId "alpha") "boom" Nothing + runEff $ unregisterProcessor master (ProcessorId "alpha") + incrementReceived beta + q <- tick control + map Out.retention (Out.entries q) `shouldBe` [Out.Update, Out.Update, Out.Required] + _ <- tick control + atomically $ writeTVar control.gate False + _ <- receiveServer conn + receiveServer conn `shouldReturn` ProcessorTerminal (ProcessorId "alpha") (TerminalFailed "boom" Nothing) + receiveServer conn `shouldReturn` ServerError "queue_overflow" "Updates dropped; replacing metrics snapshot" + expected <- MetricsSnapshot <$> currentSnapshot master + receiveServer conn `shouldReturn` expected + + it "preserves once-only terminal delivery before overflow recovery" $ \master -> do + handle <- registerIdleProcessor master (ProcessorId "alpha") + withControlledServer fastConfig {wsMaxQueuedFrames = 3} master $ \port control -> + WS.runClient "127.0.0.1" port "/ws" $ \conn -> do + _ <- receiveServer conn + atomically $ writeTVar control.gate True + incrementReceived handle + _ <- tick control + awaitBlocked control + replicateM_ 4 $ incrementReceived handle >> tick control + markProcessorFailedIO master (ProcessorId "alpha") "boom" Nothing + runEff $ unregisterProcessor master (ProcessorId "alpha") + q <- tick control + Out.queuedFrames q `shouldBe` 3 + atomically $ writeTVar control.gate False + _ <- receiveServer conn + receiveServer conn `shouldReturn` ProcessorTerminal (ProcessorId "alpha") (TerminalFailed "boom" Nothing) + receiveServer conn `shouldReturn` ServerError "queue_overflow" "Updates dropped; replacing metrics snapshot" + receiveServer conn `shouldReturn` MetricsSnapshot Map.empty + _ <- tick control + timeout 30_000 (receiveServer conn) `shouldReturn` Nothing + + it "bounds stopMetricsServer while a real peer leaves a large snapshot unread" $ \master -> do + mapM_ (\n -> registerIdleProcessor master (ProcessorId $ "large-" <> fromString (show n))) [1 .. 30_000 :: Int] + wsState <- newWebSocketState 1 + freePort <- Warp.testWithApplication (pure $ combinedApp defaultConfig master wsState []) pure + let config = defaultConfig {port = freePort, wsMaxQueuedBytes = 32_000_000, wsSendTimeoutUs = 100_000} + withMetricsServer config master $ \server -> do + threadDelay 30_000 -- listener startup; no delay is used to trigger queue overflow + WS.runClient "127.0.0.1" freePort "/ws" $ \_ -> do + threadDelay 100_000 + timeout 500_000 (stopMetricsServer server) `shouldReturn` Just () + + it "produces metrics during a continuous ping stream" $ \master -> do + handle <- registerIdleProcessor master (ProcessorId "alpha") + withServer fastConfig master $ \port -> + WS.runClient "127.0.0.1" port "/ws" $ \conn -> do + _ <- receiveServer conn + incrementReceived handle + let findUpdate = do + WS.sendTextData conn $ encode Ping + receiveServer conn >>= \case + ProcessorUpdate _ _ -> pure () + Pong -> findUpdate + other -> expectationFailure $ "unexpected frame: " <> show other + timeout 500_000 findUpdate `shouldReturn` Just () + describe "WebSocket wire protocol" $ do + it "reports invalid client messages and remains usable" $ \master -> + withServer defaultConfig master $ \port -> + WS.runClient "127.0.0.1" port "/ws" $ \conn -> do + _ <- receiveServer conn + WS.sendTextData conn ("{\"type\":\"bogus\"}" :: ByteString) + reply <- timeout 1_000_000 (WS.receiveData conn :: IO ByteString) + fmap (eitherDecode @Value) reply + `shouldBe` Just + ( Right $ + object + [ "type" .= ("error" :: String), + "code" .= ("invalid_message" :: String), + "message" .= ("Error in $: Unknown message type: bogus" :: String) + ] + ) + WS.sendTextData conn $ encode Ping + receiveServer conn `shouldReturn` Pong + it "sends an initial snapshot" $ \master -> do _ <- registerIdleProcessor master (ProcessorId "alpha") withServer defaultConfig master $ \port -> @@ -79,6 +355,16 @@ MetricsSnapshot metrics -> Map.keys metrics `shouldBe` [ProcessorId "alpha"] other -> expectationFailure $ "expected selective snapshot, got " <> show other + it "establishes an unsubscribe snapshot before subsequent replies" $ \master -> do + _ <- registerIdleProcessor master (ProcessorId "alpha") + withServer defaultConfig master $ \port -> + WS.runClient "127.0.0.1" port "/ws" $ \conn -> do + _ <- receiveServer conn + WS.sendTextData conn $ encode $ Unsubscribe [ProcessorId "alpha"] + WS.sendTextData conn $ encode Ping + receiveServer conn `shouldReturn` MetricsSnapshot Map.empty + receiveServer conn `shouldReturn` Pong + it "answers ping with pong" $ \master -> withServer defaultConfig master $ \port -> WS.runClient "127.0.0.1" port "/ws" $ \conn -> do @@ -125,16 +411,25 @@ WS.runClient "127.0.0.1" port "/ws" $ \conn -> do _ <- receiveServer conn WS.sendTextData conn $ encode $ Subscribe [ProcessorId "one", ProcessorId "two", ProcessorId "three"] - (WS.receiveData conn :: IO ByteString) - `shouldThrow` (anyException :: SomeException -> Bool) + payload <- timeout 1_000_000 (WS.receiveData conn :: IO ByteString) + fmap (eitherDecode @Value) payload + `shouldBe` Just + ( Right $ + object + [ "type" .= ("error" :: String), + "code" .= ("subscription_limit_exceeded" :: String), + "message" .= ("WebSocket processor subscription limit exceeded (maximum 2)" :: String) + ] + ) + (WS.receiveData conn :: IO ByteString) `shouldThrow` isPolicyClose it "also bounds subscribe-all exclusion state" $ \master -> withServer fastConfig {wsMaxSubscriptions = 1} master $ \port -> WS.runClient "127.0.0.1" port "/ws" $ \conn -> do _ <- receiveServer conn WS.sendTextData conn $ encode $ Unsubscribe [ProcessorId "one", ProcessorId "two"] - (WS.receiveData conn :: IO ByteString) - `shouldThrow` (anyException :: SomeException -> Bool) + _ <- WS.receiveData conn :: IO ByteString + (WS.receiveData conn :: IO ByteString) `shouldThrow` isPolicyClose it "excludes processors unsubscribed from subscribe-all" $ \master -> do alpha <- registerIdleProcessor master (ProcessorId "alpha") @@ -143,6 +438,7 @@ WS.runClient "127.0.0.1" port "/ws" $ \conn -> do _ <- receiveServer conn WS.sendTextData conn $ encode $ Unsubscribe [ProcessorId "alpha"] + _ <- receiveServer conn -- unsubscribe establishes the new generation snapshot incrementReceived alpha incrementReceived beta receiveServer conn >>= \case @@ -235,3 +531,87 @@ shouldReturn :: (Eq a, Show a) => IO a -> a -> IO () shouldReturn action expected = action >>= (`shouldBe` expected) + +originConfig :: MetricsServerConfig +originConfig = defaultConfig {corsAllowedOrigins = ["https://ops.example.com"], wsMaxConnections = 1} + +clientWithOrigin :: Int -> WS.Headers -> WS.ClientApp a -> IO a +clientWithOrigin port = WS.runClientWith "127.0.0.1" port "/ws" WS.defaultConnectionOptions + +isForbiddenHandshake :: WS.HandshakeException -> Bool +isForbiddenHandshake (WS.MalformedResponse response _) = + WS.responseCode response == 403 + && lookup "Content-Type" (WS.responseHeaders response) == Just "application/json" +isForbiddenHandshake _ = False + +isPolicyClose :: WS.ConnectionException -> Bool +isPolicyClose (WS.CloseRequest 1008 _) = True +isPolicyClose _ = False + +data Controlled = Controlled + { gate :: TVar Bool, + started :: TQueue (), + ticks :: TQueue (), + produced :: TQueue Out.OutboundQueue, + received :: TQueue (), + state :: WebSocketState + } + +withControlledServer :: MetricsServerConfig -> Master -> (Int -> Controlled -> IO a) -> IO a +withControlledServer config master action = do + control <- + Controlled + <$> newTVarIO False + <*> newTQueueIO + <*> newTQueueIO + <*> newTQueueIO + <*> newTQueueIO + <*> newWebSocketState config.wsMaxConnections + let hooks = + defaultWebSocketHooks + { beforeWrite = \_ -> do + atomically $ readTVar control.gate >>= \blocked -> when blocked (writeTQueue control.started ()) + atomically $ readTVar control.gate >>= check . not, + afterReceive = \q -> do + Out.queuedFrames q `shouldSatisfy` (<= config.wsMaxQueuedFrames) + Out.queuedBytes q `shouldSatisfy` (<= config.wsMaxQueuedBytes) + atomically $ writeTQueue control.received (), + afterProduce = \q -> do + Out.queuedFrames q `shouldSatisfy` (<= config.wsMaxQueuedFrames) + Out.queuedBytes q `shouldSatisfy` (<= config.wsMaxQueuedBytes) + atomically $ writeTQueue control.produced q, + waitForPush = \_ state -> + atomically $ + (readTVar state.shutdownRequested >>= \requested -> check requested >> pure True) + `orElse` (readTQueue control.ticks >> pure False) + } + app = + WaiWS.websocketsOr + WS.defaultConnectionOptions + (websocketAppWithHooks hooks config master control.state) + (combinedApp config {enableWebSocket = False} master control.state []) + Warp.testWithApplication (pure app) $ \port -> action port control + +tick :: Controlled -> IO Out.OutboundQueue +tick control = do + atomically $ writeTQueue control.ticks () + timeout 1_000_000 (atomically $ readTQueue control.produced) >>= \case + Just queue -> do + Out.queuedFrames queue `shouldSatisfy` (<= 64) + Out.queuedBytes queue `shouldSatisfy` (<= 1_048_576) + pure queue + Nothing -> expectationFailure "producer did not complete while writer was blocked" >> fail "producer stalled" + +awaitBlocked :: Controlled -> IO () +awaitBlocked control = timeout 1_000_000 (atomically $ readTQueue control.started) `shouldReturn` Just () + +awaitReleased :: WebSocketState -> IO () +awaitReleased state = timeout 500_000 (atomically $ readTVar state.connectionCount >>= check . (== 0)) `shouldReturn` Just () + +isCapacityClose :: WS.ConnectionException -> Bool +isCapacityClose (WS.CloseRequest 1013 _) = True +isCapacityClose _ = False + +isNormalClose :: WS.ConnectionException -> Bool +isNormalClose (WS.CloseRequest 1000 _) = True +isNormalClose _ = False diff --git a/test/golden/processor-metrics.json.golden b/test/golden/processor-metrics.json.golden --- a/test/golden/processor-metrics.json.golden +++ b/test/golden/processor-metrics.json.golden @@ -1,1 +1,1 @@ -{"failed":{"batch":{"batchedMessages":25,"batchesEmitted":24,"flushTriggered":29,"partialFailures":26,"sizeTriggered":27,"timeoutTriggered":28},"startedAt":"2026-09-20T03:25:45Z","state":{"error":"boom","status":"failed","timestamp":"2026-09-20T00:01:30Z"},"stats":{"failed":23,"processed":22,"received":21}},"idle":{"batch":{"batchedMessages":5,"batchesEmitted":4,"flushTriggered":9,"partialFailures":6,"sizeTriggered":7,"timeoutTriggered":8},"startedAt":"2026-09-20T03:25:45Z","state":{"status":"idle"},"stats":{"failed":3,"processed":2,"received":1}},"processing":{"batch":{"batchedMessages":15,"batchesEmitted":14,"flushTriggered":19,"partialFailures":16,"sizeTriggered":17,"timeoutTriggered":18},"startedAt":"2026-09-20T03:25:45Z","state":{"inFlight":2,"lastActivity":"2026-09-20T00:01:00Z","lastProgress":"2026-09-20T00:01:15Z","maxConcurrency":4,"status":"processing"},"stats":{"failed":13,"processed":12,"received":11}},"stopped":{"batch":{"batchedMessages":35,"batchesEmitted":34,"flushTriggered":39,"partialFailures":36,"sizeTriggered":37,"timeoutTriggered":38},"startedAt":"2026-09-20T03:25:45Z","state":{"status":"stopped"},"stats":{"failed":33,"processed":32,"received":31}}} +{"detailed":{"batch":{"batchedMessages":45,"batchesEmitted":44,"flushTriggered":49,"partialFailures":46,"sizeTriggered":47,"timeoutTriggered":48},"latency":{"buckets":[{"count":0,"le":1.0e-4},{"count":0,"le":2.5e-4},{"count":0,"le":5.0e-4},{"count":0,"le":1.0e-3},{"count":0,"le":2.5e-3},{"count":0,"le":5.0e-3},{"count":0,"le":1.0e-2},{"count":0,"le":2.5e-2},{"count":0,"le":5.0e-2},{"count":1,"le":0.1},{"count":3,"le":0.25},{"count":3,"le":0.5},{"count":3,"le":1},{"count":3,"le":2.5},{"count":3,"le":5},{"count":3,"le":10},{"count":3,"le":30},{"count":3,"le":60}],"count":3,"sumSeconds":0.42},"oldestInFlight":{"ageSeconds":12.5,"messageId":"m-7","startedAt":"2026-09-20T00:00:40Z"},"startedAt":"2026-09-20T03:25:45Z","state":{"inFlight":1,"lastActivity":"2026-09-20T00:00:40Z","lastProgress":"2026-09-20T00:00:40Z","maxConcurrency":4,"status":"processing"},"stats":{"failed":43,"processed":42,"received":41}},"failed":{"batch":{"batchedMessages":25,"batchesEmitted":24,"flushTriggered":29,"partialFailures":26,"sizeTriggered":27,"timeoutTriggered":28},"startedAt":"2026-09-20T03:25:45Z","state":{"error":"boom","status":"failed","timestamp":"2026-09-20T00:01:30Z"},"stats":{"failed":23,"processed":22,"received":21}},"idle":{"batch":{"batchedMessages":5,"batchesEmitted":4,"flushTriggered":9,"partialFailures":6,"sizeTriggered":7,"timeoutTriggered":8},"startedAt":"2026-09-20T03:25:45Z","state":{"status":"idle"},"stats":{"failed":3,"processed":2,"received":1}},"paused":{"batch":{"batchedMessages":65,"batchesEmitted":64,"flushTriggered":69,"partialFailures":66,"sizeTriggered":67,"timeoutTriggered":68},"startedAt":"2026-09-20T03:25:45Z","state":{"inFlight":1,"lastProgress":"2026-09-20T00:00:44Z","maxConcurrency":2,"pausedAt":"2026-09-20T00:00:45Z","status":"paused"},"stats":{"failed":63,"processed":62,"received":61}},"processing":{"batch":{"batchedMessages":15,"batchesEmitted":14,"flushTriggered":19,"partialFailures":16,"sizeTriggered":17,"timeoutTriggered":18},"startedAt":"2026-09-20T03:25:45Z","state":{"inFlight":2,"lastActivity":"2026-09-20T00:01:00Z","lastProgress":"2026-09-20T00:01:15Z","maxConcurrency":4,"status":"processing"},"stats":{"failed":13,"processed":12,"received":11}},"progressing":{"batch":{"batchedMessages":55,"batchesEmitted":54,"flushTriggered":59,"partialFailures":56,"sizeTriggered":57,"timeoutTriggered":58},"progress":{"latest":{"kind":"integer","value":"4210"},"partitions":{"p-0":{"kind":"integer","value":"4210"},"p-1":{"kind":"text","value":"0198f2f3"}},"truncated":false},"startedAt":"2026-09-20T03:25:45Z","state":{"inFlight":1,"lastActivity":"2026-09-20T00:00:40Z","lastProgress":"2026-09-20T00:00:40Z","maxConcurrency":4,"status":"processing"},"stats":{"failed":53,"processed":52,"received":51}},"stopped":{"batch":{"batchedMessages":35,"batchesEmitted":34,"flushTriggered":39,"partialFailures":36,"sizeTriggered":37,"timeoutTriggered":38},"startedAt":"2026-09-20T03:25:45Z","state":{"status":"stopped"},"stats":{"failed":33,"processed":32,"received":31}}} diff --git a/test/golden/prometheus.golden b/test/golden/prometheus.golden --- a/test/golden/prometheus.golden +++ b/test/golden/prometheus.golden @@ -1,30 +1,66 @@ # HELP shibuya_messages_received_total Total messages received by processor # TYPE shibuya_messages_received_total counter +shibuya_messages_received_total{processor="detailed"} 0.0 shibuya_messages_received_total{processor="failed"} 1.0 shibuya_messages_received_total{processor="idle"} 0.0 +shibuya_messages_received_total{processor="paused"} 0.0 shibuya_messages_received_total{processor="processing"} 2.0 shibuya_messages_received_total{processor="stopped"} 4.0 # HELP shibuya_messages_processed_total Total messages successfully processed # TYPE shibuya_messages_processed_total counter +shibuya_messages_processed_total{processor="detailed"} 0.0 shibuya_messages_processed_total{processor="failed"} 0.0 shibuya_messages_processed_total{processor="idle"} 0.0 +shibuya_messages_processed_total{processor="paused"} 0.0 shibuya_messages_processed_total{processor="processing"} 0.0 shibuya_messages_processed_total{processor="stopped"} 1.0 # HELP shibuya_messages_failed_total Total messages that failed processing # TYPE shibuya_messages_failed_total counter +shibuya_messages_failed_total{processor="detailed"} 0.0 shibuya_messages_failed_total{processor="failed"} 1.0 shibuya_messages_failed_total{processor="idle"} 0.0 +shibuya_messages_failed_total{processor="paused"} 0.0 shibuya_messages_failed_total{processor="processing"} 0.0 shibuya_messages_failed_total{processor="stopped"} 0.0 -# HELP shibuya_processor_state Current processor state (1=idle, 2=processing, 3=failed, 4=stopped) +# HELP shibuya_processor_state Current processor state (1=idle, 2=processing, 3=failed, 4=stopped, 5=paused) # TYPE shibuya_processor_state gauge +shibuya_processor_state{processor="detailed"} 1.0 shibuya_processor_state{processor="failed"} 3.0 shibuya_processor_state{processor="idle"} 1.0 +shibuya_processor_state{processor="paused"} 5.0 shibuya_processor_state{processor="processing"} 2.0 shibuya_processor_state{processor="stopped"} 4.0 # HELP shibuya_processor_in_flight Number of messages currently being processed # TYPE shibuya_processor_in_flight gauge +shibuya_processor_in_flight{processor="detailed"} 0.0 shibuya_processor_in_flight{processor="failed"} 0.0 shibuya_processor_in_flight{processor="idle"} 0.0 +shibuya_processor_in_flight{processor="paused"} 1.0 shibuya_processor_in_flight{processor="processing"} 1.0 shibuya_processor_in_flight{processor="stopped"} 0.0 +# HELP shibuya_message_processing_seconds Processing attempt duration including finalization retries (one observation per message or emitted batch) +# TYPE shibuya_message_processing_seconds histogram +shibuya_message_processing_seconds_bucket{processor="detailed",le="0.0001"} 0 +shibuya_message_processing_seconds_bucket{processor="detailed",le="0.00025"} 0 +shibuya_message_processing_seconds_bucket{processor="detailed",le="0.0005"} 0 +shibuya_message_processing_seconds_bucket{processor="detailed",le="0.001"} 0 +shibuya_message_processing_seconds_bucket{processor="detailed",le="0.0025"} 0 +shibuya_message_processing_seconds_bucket{processor="detailed",le="0.005"} 1 +shibuya_message_processing_seconds_bucket{processor="detailed",le="0.01"} 1 +shibuya_message_processing_seconds_bucket{processor="detailed",le="0.025"} 1 +shibuya_message_processing_seconds_bucket{processor="detailed",le="0.05"} 1 +shibuya_message_processing_seconds_bucket{processor="detailed",le="0.1"} 1 +shibuya_message_processing_seconds_bucket{processor="detailed",le="0.25"} 1 +shibuya_message_processing_seconds_bucket{processor="detailed",le="0.5"} 1 +shibuya_message_processing_seconds_bucket{processor="detailed",le="1"} 1 +shibuya_message_processing_seconds_bucket{processor="detailed",le="2.5"} 1 +shibuya_message_processing_seconds_bucket{processor="detailed",le="5"} 1 +shibuya_message_processing_seconds_bucket{processor="detailed",le="10"} 1 +shibuya_message_processing_seconds_bucket{processor="detailed",le="30"} 1 +shibuya_message_processing_seconds_bucket{processor="detailed",le="60"} 1 +shibuya_message_processing_seconds_bucket{processor="detailed",le="+Inf"} 1 +shibuya_message_processing_seconds_sum{processor="detailed"} 3.0e-3 +shibuya_message_processing_seconds_count{processor="detailed"} 1 +# HELP shibuya_processor_oldest_in_flight_seconds Age of oldest in-flight processing attempt +# TYPE shibuya_processor_oldest_in_flight_seconds gauge +shibuya_processor_oldest_in_flight_seconds{processor="detailed"} 1.25e-2