packages feed

shibuya-metrics (empty) → 0.1.0.0

raw patch · 10 files changed

+1087/−0 lines, 10 files

Files

+ CHANGELOG.md view
@@ -0,0 +1,12 @@+# Changelog++## 0.1.0.0 — 2026-02-24++Initial release.++### New Features++- HTTP/JSON metrics endpoint+- Prometheus metrics endpoint+- WebSocket streaming metrics endpoint+- Kubernetes-compatible health check endpoints
+ shibuya-metrics.cabal view
@@ -0,0 +1,62 @@+cabal-version: 3.14+name: shibuya-metrics+version: 0.1.0.0+synopsis: Metrics web server for Shibuya queue processing framework+description:+  Provides HTTP/JSON, Prometheus, and WebSocket endpoints for+  exposing Shibuya processor metrics.++author: Nadeem Bitar+maintainer: nadeem@gmail.com+license: MIT+build-type: Simple+category: Concurrency+extra-doc-files: CHANGELOG.md++common warnings+  ghc-options: -Wall++library+  import: warnings+  exposed-modules:+    Shibuya.Metrics+    Shibuya.Metrics.Config+    Shibuya.Metrics.Health+    Shibuya.Metrics.JSON+    Shibuya.Metrics.Prometheus+    Shibuya.Metrics.Server+    Shibuya.Metrics.Types+    Shibuya.Metrics.WebSocket++  default-extensions:+    DeriveAnyClass+    DerivingStrategies+    DuplicateRecordFields+    LambdaCase+    NoFieldSelectors+    OverloadedLabels+    OverloadedRecordDot+    OverloadedStrings+    QuasiQuotes+    RecordWildCards++  build-depends:+    aeson ^>=2.2,+    async ^>=2.2,+    base ^>=4.21.0.0,+    bytestring ^>=0.12,+    containers ^>=0.7,+    http-types ^>=0.12,+    prometheus-client ^>=1.1,+    shibuya-core ^>=0.1.0.0,+    stm ^>=2.5,+    text ^>=2.1,+    time ^>=1.14,+    unliftio ^>=0.2,+    wai ^>=3.2,+    wai-websockets ^>=3.0,+    warp ^>=3.4,+    websockets ^>=0.13,++  hs-source-dirs: src+  default-language: GHC2024
+ src/Shibuya/Metrics.hs view
@@ -0,0 +1,88 @@+-- | Metrics web server for Shibuya queue processing framework.+--+-- This module provides HTTP/JSON, Prometheus, and WebSocket endpoints+-- for exposing Shibuya processor metrics.+--+-- == Quick Start+--+-- @+-- import Shibuya.App (runApp, AppHandle)+-- import Shibuya.Metrics+--+-- main :: IO ()+-- main = do+--   -- Start your Shibuya app...+--   Right appHandle <- runEff $ runApp ...+--+--   -- Start metrics server on port 9090+--   withMetricsServer defaultConfig appHandle.master $ \\server -> do+--     putStrLn $ "Metrics server running on port " ++ show server.serverPort+--     -- Your app logic here...+-- @+--+-- == Endpoints+--+-- * @GET /metrics@ - JSON metrics for all processors+-- * @GET /metrics/:id@ - JSON metrics for a specific processor+-- * @GET /health@ - Detailed health status (for debugging)+-- * @GET /health/live@ - Liveness probe (Kubernetes)+-- * @GET /health/ready@ - Readiness probe (Kubernetes)+-- * @GET /metrics/prometheus@ - Prometheus-format metrics+-- * @WS /ws@ - WebSocket for real-time updates+--+-- == WebSocket Protocol+--+-- Clients can send:+--+-- * @{"type": "subscribe_all"}@ - Subscribe to all processor updates+-- * @{"type": "subscribe", "processors": ["id1", "id2"]}@ - Subscribe to specific processors+-- * @{"type": "unsubscribe", "processors": ["id1"]}@ - Unsubscribe from processors+-- * @{"type": "ping"}@ - Keepalive ping+--+-- Server sends:+--+-- * @{"type": "snapshot", "metrics": {...}}@ - Full metrics snapshot+-- * @{"type": "update", "processor": "id", "metrics": {...}}@ - Single processor update+-- * @{"type": "pong"}@ - Response to ping+-- * @{"type": "goodbye"}@ - Server shutting down+module Shibuya.Metrics+  ( -- * Server Lifecycle+    startMetricsServer,+    startMetricsServerWithDeps,+    stopMetricsServer,+    withMetricsServer,++    -- * Configuration+    MetricsServerConfig (..),+    defaultConfig,++    -- * Server Handle+    MetricsServer (..),++    -- * Health Check Types+    DependencyCheck,+    DependencyStatus (..),+    LivenessStatus (..),+    ReadinessStatus (..),+    ProcessorHealth (..),+    HealthConfig (..),+    defaultHealthConfig,++    -- * WebSocket Protocol Types+    ClientMessage (..),+    ServerMessage (..),+  )+where++import Shibuya.Metrics.Config (MetricsServerConfig (..), defaultConfig)+import Shibuya.Metrics.Health+  ( DependencyCheck,+    DependencyStatus (..),+    HealthConfig (..),+    LivenessStatus (..),+    ProcessorHealth (..),+    ReadinessStatus (..),+    defaultHealthConfig,+  )+import Shibuya.Metrics.Server (startMetricsServer, startMetricsServerWithDeps, stopMetricsServer, withMetricsServer)+import Shibuya.Metrics.Types (ClientMessage (..), MetricsServer (..), ServerMessage (..))
+ src/Shibuya/Metrics/Config.hs view
@@ -0,0 +1,44 @@+-- | Configuration for the metrics web server.+module Shibuya.Metrics.Config+  ( MetricsServerConfig (..),+    defaultConfig,+  )+where++import Data.Time.Clock (NominalDiffTime)+import GHC.Generics (Generic)++-- | Configuration for the metrics web server.+data MetricsServerConfig = MetricsServerConfig+  { -- | Port to listen on (default: 9090)+    port :: !Int,+    -- | Enable JSON endpoints (default: True)+    enableJSON :: !Bool,+    -- | Enable Prometheus endpoint (default: True)+    enablePrometheus :: !Bool,+    -- | Enable WebSocket endpoint (default: True)+    enableWebSocket :: !Bool,+    -- | WebSocket push interval in microseconds (default: 100_000 = 100ms)+    wsPushIntervalUs :: !Int,+    -- | Maximum WebSocket connections (default: 100)+    wsMaxConnections :: !Int,+    -- | Timeout for liveness check in microseconds (default: 1_000_000 = 1s)+    livenessTimeoutMicros :: !Int,+    -- | How long a processor can be in Processing state before considered stuck (default: 60s)+    stuckThreshold :: !NominalDiffTime+  }+  deriving stock (Eq, Show, Generic)++-- | Default configuration.+defaultConfig :: MetricsServerConfig+defaultConfig =+  MetricsServerConfig+    { port = 9090,+      enableJSON = True,+      enablePrometheus = True,+      enableWebSocket = True,+      wsPushIntervalUs = 100_000, -- 100ms+      wsMaxConnections = 100,+      livenessTimeoutMicros = 1_000_000, -- 1 second+      stuckThreshold = 60 -- 60 seconds+    }
+ src/Shibuya/Metrics/Health.hs view
@@ -0,0 +1,219 @@+-- | Health check types and logic for Kubernetes-compatible probes.+--+-- Provides:+-- * Liveness probe: Is the system running?+-- * Readiness probe: Is the system ready to handle traffic?+-- * Detailed health status for debugging+module Shibuya.Metrics.Health+  ( -- * Health Status Types+    LivenessStatus (..),+    ReadinessStatus (..),+    ProcessorHealth (..),+    DependencyStatus (..),++    -- * Health Check Configuration+    HealthConfig (..),+    defaultHealthConfig,++    -- * Health Check Functions+    checkLiveness,+    checkReadiness,+    checkDetailedHealth,++    -- * Dependency Checks+    DependencyCheck,+  )+where++import Data.Aeson (ToJSON (..), object, (.=))+import Data.Map.Strict qualified as Map+import Data.Maybe (isJust)+import Data.Text (Text)+import Data.Time.Clock (NominalDiffTime, UTCTime, diffUTCTime, getCurrentTime)+import Shibuya.Runner.Master (Master, getAllMetricsIO)+import Shibuya.Runner.Metrics+  ( MetricsMap,+    ProcessorMetrics (..),+    ProcessorState (..),+  )+import System.Timeout (timeout)++--------------------------------------------------------------------------------+-- Configuration+--------------------------------------------------------------------------------++-- | Configuration for health checks.+data HealthConfig = HealthConfig+  { -- | Timeout for liveness check (microseconds)+    livenessTimeoutMicros :: !Int,+    -- | How long a processor can be in Processing state before considered stuck+    stuckThreshold :: !NominalDiffTime+  }+  deriving stock (Eq, Show)++-- | Default health configuration.+-- Liveness timeout: 1 second+-- Stuck threshold: 60 seconds+defaultHealthConfig :: HealthConfig+defaultHealthConfig =+  HealthConfig+    { livenessTimeoutMicros = 1_000_000,+      stuckThreshold = 60+    }++--------------------------------------------------------------------------------+-- Health Status Types+--------------------------------------------------------------------------------++-- | Liveness status for Kubernetes liveness probe.+-- A simple "am I running?" check.+data LivenessStatus = LivenessStatus+  { alive :: !Bool+  }+  deriving stock (Eq, Show)++instance ToJSON LivenessStatus where+  toJSON status =+    object+      [ "alive" .= status.alive+      ]++-- | Readiness status for Kubernetes readiness probe.+-- Indicates whether the system is ready to handle traffic.+data ReadinessStatus = ReadinessStatus+  { ready :: !Bool,+    processors :: !ProcessorHealth,+    dependencies :: ![DependencyStatus]+  }+  deriving stock (Eq, Show)++instance ToJSON ReadinessStatus where+  toJSON status =+    object+      [ "ready" .= status.ready,+        "processors" .= status.processors,+        "dependencies" .= status.dependencies+      ]++-- | Summary of processor health across all processors.+data ProcessorHealth = ProcessorHealth+  { total :: !Int,+    healthy :: !Int,+    failed :: !Int,+    stuck :: !Int+  }+  deriving stock (Eq, Show)++instance ToJSON ProcessorHealth where+  toJSON ph =+    object+      [ "total" .= ph.total,+        "healthy" .= ph.healthy,+        "failed" .= ph.failed,+        "stuck" .= ph.stuck+      ]++-- | Status of an external dependency.+data DependencyStatus = DependencyStatus+  { name :: !Text,+    healthy :: !Bool,+    latencyMs :: !(Maybe Int),+    errorMsg :: !(Maybe Text)+  }+  deriving stock (Eq, Show)++instance ToJSON DependencyStatus where+  toJSON ds =+    object+      [ "name" .= ds.name,+        "healthy" .= ds.healthy,+        "latencyMs" .= ds.latencyMs,+        "error" .= ds.errorMsg+      ]++-- | A dependency check is an IO action that returns the dependency's status.+type DependencyCheck = IO DependencyStatus++--------------------------------------------------------------------------------+-- Health Check Functions+--------------------------------------------------------------------------------++-- | Check liveness - is the master responding?+-- This is a fast check suitable for Kubernetes liveness probes.+checkLiveness :: HealthConfig -> Master -> IO LivenessStatus+checkLiveness config master = do+  -- Try to query metrics with timeout+  result <- timeout config.livenessTimeoutMicros $ getAllMetricsIO master+  pure $ LivenessStatus {alive = isJust result}++-- | Check readiness - are all processors healthy and dependencies available?+-- This is suitable for Kubernetes readiness probes.+checkReadiness ::+  HealthConfig ->+  Master ->+  [DependencyCheck] ->+  IO ReadinessStatus+checkReadiness config master depChecks = do+  now <- getCurrentTime+  metrics <- getAllMetricsIO master+  let procHealth = analyzeProcessorHealth config now metrics+  depStatus <- sequence depChecks++  let allDepsHealthy = all (.healthy) depStatus+      noFailedProcessors = procHealth.failed == 0+      noStuckProcessors = procHealth.stuck == 0+      isReady = allDepsHealthy && noFailedProcessors && noStuckProcessors++  pure+    ReadinessStatus+      { ready = isReady,+        processors = procHealth,+        dependencies = depStatus+      }++-- | Get detailed health status for debugging.+-- Returns the full metrics along with health analysis.+checkDetailedHealth ::+  HealthConfig ->+  Master ->+  [DependencyCheck] ->+  IO (ReadinessStatus, MetricsMap)+checkDetailedHealth config master depChecks = do+  readiness <- checkReadiness config master depChecks+  metrics <- getAllMetricsIO master+  pure (readiness, metrics)++--------------------------------------------------------------------------------+-- Internal Helpers+--------------------------------------------------------------------------------++-- | Analyze processor health from metrics.+analyzeProcessorHealth :: HealthConfig -> UTCTime -> MetricsMap -> ProcessorHealth+analyzeProcessorHealth config now metrics =+  let processors = Map.elems metrics+      total = length processors+      (healthy, failed, stuck) = foldr (categorize config now) (0, 0, 0) processors+   in ProcessorHealth+        { total = total,+          healthy = healthy,+          failed = failed,+          stuck = stuck+        }++-- | Categorize a processor as healthy, failed, or stuck.+categorize ::+  HealthConfig ->+  UTCTime ->+  ProcessorMetrics ->+  (Int, Int, Int) ->+  (Int, Int, Int)+categorize config now pm (h, f, s) =+  case pm.state of+    Idle -> (h + 1, f, s)+    Stopped -> (h, f, s) -- Stopped is neither healthy nor failed+    Failed _ _ -> (h, f + 1, s)+    Processing _ lastActivity ->+      let timeSinceActivity = diffUTCTime now lastActivity+       in if timeSinceActivity > config.stuckThreshold+            then (h, f, s + 1) -- Stuck+            else (h + 1, f, s) -- Healthy (actively processing)
+ src/Shibuya/Metrics/JSON.hs view
@@ -0,0 +1,101 @@+-- | JSON HTTP endpoints for metrics and health checks.+module Shibuya.Metrics.JSON+  ( jsonApp,+    jsonAppWithHealth,+  )+where++import Data.Aeson (encode, object, (.=))+import Data.ByteString.Lazy qualified as LBS+import Data.Text (Text)+import Network.HTTP.Types+  ( Status,+    hContentType,+    status200,+    status404,+    status503,+  )+import Network.Wai (Application, Response, pathInfo, responseLBS)+import Shibuya.Metrics.Health+  ( DependencyCheck,+    HealthConfig,+    LivenessStatus (..),+    ReadinessStatus (..),+    checkDetailedHealth,+    checkLiveness,+    checkReadiness,+    defaultHealthConfig,+  )+import Shibuya.Runner.Master (Master, getAllMetricsIO, getProcessorMetricsIO)+import Shibuya.Runner.Metrics (ProcessorId (..))++-- | WAI application for JSON endpoints (without dependency checks).+-- Handles:+--   GET /metrics       - all processor metrics+--   GET /metrics/:id   - single processor metrics+--   GET /health        - detailed health status+--   GET /health/live   - liveness probe+--   GET /health/ready  - readiness probe+jsonApp :: Master -> Application+jsonApp master = jsonAppWithHealth defaultHealthConfig master []++-- | WAI application for JSON endpoints with dependency checks.+jsonAppWithHealth :: HealthConfig -> Master -> [DependencyCheck] -> Application+jsonAppWithHealth config master depChecks req respond = do+  let path = pathInfo req+  response <- routeRequest config master depChecks path+  respond response++routeRequest :: HealthConfig -> Master -> [DependencyCheck] -> [Text] -> IO Response+routeRequest config master depChecks path = case path of+  -- GET /metrics - all processor metrics+  ["metrics"] -> do+    metrics <- getAllMetricsIO master+    pure $ jsonResponse status200 $ encode metrics+  -- GET /metrics/:id - single processor metrics+  ["metrics", procIdText] -> do+    let procId = ProcessorId procIdText+    mMetrics <- getProcessorMetricsIO master procId+    case mMetrics of+      Just m -> pure $ jsonResponse status200 $ encode m+      Nothing ->+        pure $+          jsonResponse status404 $+            encode $+              object+                [ "error" .= ("Processor not found" :: Text),+                  "processor" .= procIdText+                ]+  -- GET /health/live - liveness probe (fast, simple)+  ["health", "live"] -> do+    liveness <- checkLiveness config master+    let LivenessStatus {alive} = liveness+        status = if alive then status200 else status503+    pure $ jsonResponse status $ encode liveness+  -- GET /health/ready - readiness probe (checks processors and dependencies)+  ["health", "ready"] -> do+    readiness <- checkReadiness config master depChecks+    let ReadinessStatus {ready} = readiness+        status = if ready then status200 else status503+    pure $ jsonResponse status $ encode readiness+  -- GET /health - detailed health status (for debugging)+  ["health"] -> do+    (readiness, metrics) <- checkDetailedHealth config master depChecks+    let ReadinessStatus {ready} = readiness+        status = if ready then status200 else status503+    pure $+      jsonResponse status $+        encode $+          object+            [ "status" .= readiness,+              "processors" .= metrics+            ]+  -- Not found+  _ ->+    pure $+      jsonResponse status404 $+        encode $+          object ["error" .= ("Not found" :: Text)]++jsonResponse :: Status -> LBS.ByteString -> Response+jsonResponse status = responseLBS status [(hContentType, "application/json")]
+ src/Shibuya/Metrics/Prometheus.hs view
@@ -0,0 +1,107 @@+-- | Prometheus metrics endpoint.+module Shibuya.Metrics.Prometheus+  ( prometheusApp,+  )+where++import Data.ByteString.Builder qualified as Builder+import Data.Map.Strict qualified as Map+import Data.Text (Text)+import Data.Text qualified as Text+import Data.Text.Encoding qualified as Text+import Network.HTTP.Types (hContentType, status200)+import Network.Wai (Application, responseLBS)+import Shibuya.Runner.Master (Master, getAllMetricsIO)+import Shibuya.Runner.Metrics+  ( InFlightInfo (..),+    MetricsMap,+    ProcessorId (..),+    ProcessorMetrics (..),+    ProcessorState (..),+    StreamStats (..),+  )++-- | WAI application for Prometheus metrics endpoint.+-- Handles: GET /metrics/prometheus+prometheusApp :: Master -> Application+prometheusApp master _req respond = do+  metrics <- getAllMetricsIO master+  let body = renderPrometheusMetrics metrics+  respond $+    responseLBS+      status200+      [(hContentType, "text/plain; version=0.0.4; charset=utf-8")]+      (Builder.toLazyByteString body)++-- | Render metrics in Prometheus text format.+renderPrometheusMetrics :: MetricsMap -> Builder.Builder+renderPrometheusMetrics metrics =+  mconcat+    [ renderHelp "shibuya_messages_received_total" "Total messages received by processor",+      renderType "shibuya_messages_received_total" "counter",+      renderGaugeMetrics "shibuya_messages_received_total" (fromIntegral . (.stats.received)) metrics,+      renderHelp "shibuya_messages_processed_total" "Total messages successfully processed",+      renderType "shibuya_messages_processed_total" "counter",+      renderGaugeMetrics "shibuya_messages_processed_total" (fromIntegral . (.stats.processed)) metrics,+      renderHelp "shibuya_messages_dropped_total" "Total messages dropped due to backpressure",+      renderType "shibuya_messages_dropped_total" "counter",+      renderGaugeMetrics "shibuya_messages_dropped_total" (fromIntegral . (.stats.dropped)) metrics,+      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)",+      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+    ]++renderHelp :: Text -> Text -> Builder.Builder+renderHelp name help =+  Builder.byteString "# HELP "+    <> Builder.byteString (Text.encodeUtf8 name)+    <> Builder.char8 ' '+    <> Builder.byteString (Text.encodeUtf8 help)+    <> Builder.char8 '\n'++renderType :: Text -> Text -> Builder.Builder+renderType name typ =+  Builder.byteString "# TYPE "+    <> Builder.byteString (Text.encodeUtf8 name)+    <> Builder.char8 ' '+    <> Builder.byteString (Text.encodeUtf8 typ)+    <> Builder.char8 '\n'++renderGaugeMetrics :: Text -> (ProcessorMetrics -> Double) -> MetricsMap -> Builder.Builder+renderGaugeMetrics name getValue = Map.foldMapWithKey (renderMetric name getValue)++renderMetric :: Text -> (ProcessorMetrics -> Double) -> ProcessorId -> ProcessorMetrics -> Builder.Builder+renderMetric name getValue (ProcessorId pid) pm =+  Builder.byteString (Text.encodeUtf8 name)+    <> Builder.byteString "{processor=\""+    <> Builder.byteString (Text.encodeUtf8 $ escapeLabel pid)+    <> Builder.byteString "\"} "+    <> Builder.doubleDec (getValue pm)+    <> Builder.char8 '\n'++-- | Escape special characters in Prometheus label values.+escapeLabel :: Text -> Text+escapeLabel = Text.concatMap escapeChar+  where+    escapeChar '\\' = "\\\\"+    escapeChar '"' = "\\\""+    escapeChar '\n' = "\\n"+    escapeChar c = Text.singleton c++-- | Convert processor state to integer for Prometheus.+stateToInt :: ProcessorState -> Int+stateToInt Idle = 1+stateToInt (Processing _ _) = 2+stateToInt (Failed _ _) = 3+stateToInt Stopped = 4++-- | Get in-flight count from processor state.+inFlightCount :: ProcessorState -> Int+inFlightCount (Processing info _) = info.inFlight+inFlightCount _ = 0
+ src/Shibuya/Metrics/Server.hs view
@@ -0,0 +1,133 @@+-- | Unified metrics web server combining JSON, Prometheus, and WebSocket endpoints.+module Shibuya.Metrics.Server+  ( -- * Server Lifecycle+    startMetricsServer,+    startMetricsServerWithDeps,+    stopMetricsServer,+    withMetricsServer,++    -- * Re-exports+    MetricsServer (..),+    MetricsServerConfig (..),+    defaultConfig,+    DependencyCheck,+  )+where++import Control.Concurrent.Async (async, cancel)+import Control.Exception (bracket)+import Data.Aeson (encode, object, (.=))+import Data.Text (Text)+import Network.HTTP.Types (hContentType, status404)+import Network.Wai (Application, Response, pathInfo, responseLBS)+import Network.Wai.Handler.Warp qualified as Warp+import Network.Wai.Handler.WebSockets qualified as WaiWS+import Network.WebSockets qualified as WS+import Shibuya.Metrics.Config (MetricsServerConfig (..), defaultConfig)+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, websocketApp)+import Shibuya.Runner.Master (Master)++-- | Start the metrics server without dependency checks.+-- Returns a handle that can be used to stop the server.+startMetricsServer :: MetricsServerConfig -> Master -> IO MetricsServer+startMetricsServer config master = startMetricsServerWithDeps config master []++-- | Start the metrics server with dependency checks for health endpoints.+-- Returns a handle that can be used to stop the server.+startMetricsServerWithDeps ::+  MetricsServerConfig ->+  Master ->+  [DependencyCheck] ->+  IO MetricsServer+startMetricsServerWithDeps config master depChecks = do+  wsState <- newWebSocketState config.wsMaxConnections+  let app = combinedApp config master wsState depChecks+      settings =+        Warp.setPort config.port $+          Warp.setHost+            "*"+            Warp.defaultSettings+  serverAsync <- async $ Warp.runSettings settings app+  pure+    MetricsServer+      { serverThread = serverAsync,+        serverPort = config.port+      }++-- | Stop the metrics server.+stopMetricsServer :: MetricsServer -> IO ()+stopMetricsServer server = cancel server.serverThread++-- | Run an action with a metrics server, ensuring cleanup.+withMetricsServer ::+  MetricsServerConfig ->+  Master ->+  (MetricsServer -> IO a) ->+  IO a+withMetricsServer config master =+  bracket+    (startMetricsServer config master)+    stopMetricsServer++-- | Combined WAI application routing to all endpoints.+combinedApp ::+  MetricsServerConfig ->+  Master ->+  WebSocketState ->+  [DependencyCheck] ->+  Application+combinedApp config master wsState depChecks =+  -- Handle WebSocket upgrade first+  WaiWS.websocketsOr+    WS.defaultConnectionOptions+    (websocketApp config master wsState)+    (httpApp config master depChecks)++-- | HTTP application routing based on path.+httpApp :: MetricsServerConfig -> Master -> [DependencyCheck] -> Application+httpApp config master depChecks req respond = do+  let path = pathInfo req+      healthConfig =+        HealthConfig+          { livenessTimeoutMicros = config.livenessTimeoutMicros,+            stuckThreshold = config.stuckThreshold+          }+      jsonHandler = jsonAppWithHealth healthConfig master depChecks+  case path of+    -- Prometheus endpoint+    ["metrics", "prometheus"]+      | config.enablePrometheus ->+          prometheusApp master req respond+    -- JSON endpoints (metrics and health)+    ["metrics"]+      | config.enableJSON ->+          jsonHandler req respond+    ["metrics", _]+      | config.enableJSON ->+          jsonHandler req respond+    ["health"]+      | config.enableJSON ->+          jsonHandler req respond+    ["health", "live"]+      | config.enableJSON ->+          jsonHandler req respond+    ["health", "ready"]+      | config.enableJSON ->+          jsonHandler req respond+    -- WebSocket path info (for documentation, actual upgrade handled above)+    ["ws"]+      | config.enableWebSocket ->+          respond $ notFoundResponse "WebSocket endpoint - use ws:// protocol"+    -- Not found+    _ -> respond $ notFoundResponse "Not found"++notFoundResponse :: Text -> Response+notFoundResponse msg =+  responseLBS+    status404+    [(hContentType, "application/json")]+    (encode $ object ["error" .= msg])
+ src/Shibuya/Metrics/Types.hs view
@@ -0,0 +1,100 @@+-- | Types for the metrics server, including WebSocket protocol messages.+module Shibuya.Metrics.Types+  ( -- * WebSocket Protocol+    ClientMessage (..),+    ServerMessage (..),++    -- * Server Handle+    MetricsServer (..),+  )+where++import Control.Concurrent.Async (Async)+import Data.Aeson+  ( FromJSON (..),+    ToJSON (..),+    object,+    withObject,+    (.:),+    (.=),+  )+import Data.Text (Text)+import Data.Text qualified as Text+import GHC.Generics (Generic)+import Network.Wai.Handler.Warp (Port)+import Shibuya.Runner.Metrics (MetricsMap, ProcessorId, ProcessorMetrics)++--------------------------------------------------------------------------------+-- WebSocket Protocol+--------------------------------------------------------------------------------++-- | Messages from client to server.+data ClientMessage+  = -- | Subscribe to all processor metrics+    SubscribeAll+  | -- | Subscribe to specific processors+    Subscribe ![ProcessorId]+  | -- | Unsubscribe from specific processors+    Unsubscribe ![ProcessorId]+  | -- | Ping for keepalive+    Ping+  deriving stock (Eq, Show, Generic)++instance ToJSON ClientMessage where+  toJSON SubscribeAll = object ["type" .= ("subscribe_all" :: Text)]+  toJSON (Subscribe pids) = object ["type" .= ("subscribe" :: Text), "processors" .= pids]+  toJSON (Unsubscribe pids) = object ["type" .= ("unsubscribe" :: Text), "processors" .= pids]+  toJSON Ping = object ["type" .= ("ping" :: Text)]++instance FromJSON ClientMessage where+  parseJSON = withObject "ClientMessage" $ \v -> do+    msgType <- v .: "type"+    case msgType :: Text of+      "subscribe_all" -> pure SubscribeAll+      "subscribe" -> Subscribe <$> v .: "processors"+      "unsubscribe" -> Unsubscribe <$> v .: "processors"+      "ping" -> pure Ping+      other -> fail $ "Unknown message type: " <> Text.unpack other++-- | Messages from server to client.+data ServerMessage+  = -- | Full metrics snapshot (sent on subscribe)+    MetricsSnapshot !MetricsMap+  | -- | Update for a single processor (sent on push)+    ProcessorUpdate !ProcessorId !ProcessorMetrics+  | -- | Pong response to ping+    Pong+  | -- | Server is shutting down+    Goodbye+  deriving stock (Eq, Show, Generic)++instance ToJSON ServerMessage where+  toJSON (MetricsSnapshot metrics) =+    object ["type" .= ("snapshot" :: Text), "metrics" .= metrics]+  toJSON (ProcessorUpdate pid pm) =+    object ["type" .= ("update" :: Text), "processor" .= pid, "metrics" .= pm]+  toJSON Pong = object ["type" .= ("pong" :: Text)]+  toJSON Goodbye = object ["type" .= ("goodbye" :: Text)]++instance FromJSON ServerMessage where+  parseJSON = withObject "ServerMessage" $ \v -> do+    msgType <- v .: "type"+    case msgType :: Text of+      "snapshot" -> MetricsSnapshot <$> v .: "metrics"+      "update" -> ProcessorUpdate <$> v .: "processor" <*> v .: "metrics"+      "pong" -> pure Pong+      "goodbye" -> pure Goodbye+      other -> fail $ "Unknown message type: " <> Text.unpack other++--------------------------------------------------------------------------------+-- Server Handle+--------------------------------------------------------------------------------++-- | Handle for a running metrics server.+data MetricsServer = MetricsServer+  { -- | Async handle for the server thread+    serverThread :: !(Async ()),+    -- | Port the server is running on+    serverPort :: !Port+  }+  deriving stock (Generic)
+ src/Shibuya/Metrics/WebSocket.hs view
@@ -0,0 +1,221 @@+-- | WebSocket endpoint for real-time metrics updates.+module Shibuya.Metrics.WebSocket+  ( websocketApp,+    WebSocketState (..),+    newWebSocketState,+  )+where++import Control.Concurrent (threadDelay)+import Control.Concurrent.Async (async, cancel, link)+import Control.Concurrent.STM+  ( STM,+    TVar,+    atomically,+    modifyTVar',+    newTVarIO,+    readTVar,+    writeTVar,+  )+import Control.Exception (finally)+import Control.Monad (forever, when)+import Data.Aeson (decode, encode)+import Data.Map.Strict qualified as Map+import Data.Set (Set)+import Data.Set qualified as Set+import Network.WebSockets qualified as WS+import Shibuya.Metrics.Config (MetricsServerConfig (..))+import Shibuya.Metrics.Types (ClientMessage (..), ServerMessage (..))+import Shibuya.Runner.Master (Master, getAllMetricsIO)+import Shibuya.Runner.Metrics (MetricsMap, ProcessorId (..), ProcessorMetrics)++--------------------------------------------------------------------------------+-- WebSocket State+--------------------------------------------------------------------------------++-- | Shared state for WebSocket connections.+data WebSocketState = WebSocketState+  { -- | Current number of connections+    connectionCount :: !(TVar Int),+    -- | Maximum allowed connections+    maxConnections :: !Int+  }++-- | Create new WebSocket state.+newWebSocketState :: Int -> IO WebSocketState+newWebSocketState maxConns = do+  countVar <- newTVarIO 0+  pure+    WebSocketState+      { connectionCount = countVar,+        maxConnections = maxConns+      }++-- | Try to acquire a connection slot. Returns True if successful.+acquireConnection :: WebSocketState -> STM Bool+acquireConnection wsState = do+  count <- readTVar wsState.connectionCount+  if count < wsState.maxConnections+    then do+      writeTVar wsState.connectionCount (count + 1)+      pure True+    else pure False++-- | Release a connection slot.+releaseConnection :: WebSocketState -> STM ()+releaseConnection wsState =+  modifyTVar' wsState.connectionCount (\c -> max 0 (c - 1))++--------------------------------------------------------------------------------+-- Connection State+--------------------------------------------------------------------------------++-- | State for a single WebSocket connection.+data ConnectionState = ConnectionState+  { -- | Subscribed processors (Nothing = all)+    subscriptions :: !(TVar (Maybe (Set ProcessorId))),+    -- | Last sent metrics for delta detection+    lastMetrics :: !(TVar MetricsMap)+  }++-- | Create new connection state.+newConnectionState :: IO ConnectionState+newConnectionState = do+  subsVar <- newTVarIO Nothing -- Start subscribed to all+  lastVar <- newTVarIO Map.empty+  pure+    ConnectionState+      { subscriptions = subsVar,+        lastMetrics = lastVar+      }++--------------------------------------------------------------------------------+-- WebSocket Application+--------------------------------------------------------------------------------++-- | WebSocket server application.+websocketApp ::+  MetricsServerConfig ->+  Master ->+  WebSocketState ->+  WS.ServerApp+websocketApp config master wsState pending = do+  -- Try to acquire a connection slot+  acquired <- atomically $ acquireConnection wsState+  if not acquired+    then WS.rejectRequest pending "Too many connections"+    else do+      conn <- WS.acceptRequest pending+      -- Set up connection with ping/pong for keepalive+      WS.withPingThread conn 30 (pure ()) $ do+        connState <- newConnectionState+        -- Send initial snapshot+        metrics <- getAllMetricsIO master+        WS.sendTextData conn $ encode $ MetricsSnapshot metrics+        atomically $ writeTVar connState.lastMetrics metrics+        -- Run receive and push loops concurrently+        pushThread <- async $ pushLoop config master connState conn+        link pushThread+        finally+          (receiveLoop master connState conn)+          ( do+              cancel pushThread+              WS.sendTextData conn $ encode Goodbye+              atomically $ releaseConnection wsState+          )++--------------------------------------------------------------------------------+-- Receive Loop+--------------------------------------------------------------------------------++-- | Handle incoming messages from client.+receiveLoop :: Master -> ConnectionState -> WS.Connection -> IO ()+receiveLoop master connState conn = forever $ do+  msg <- WS.receiveData conn+  case decode msg of+    Nothing -> pure () -- Ignore invalid messages+    Just clientMsg -> handleClientMessage master connState conn clientMsg++-- | Handle a client message.+handleClientMessage ::+  Master ->+  ConnectionState ->+  WS.Connection ->+  ClientMessage ->+  IO ()+handleClientMessage master connState conn = \case+  SubscribeAll -> do+    atomically $ writeTVar connState.subscriptions Nothing+    -- Send snapshot of all metrics+    metrics <- getAllMetricsIO master+    WS.sendTextData conn $ encode $ MetricsSnapshot metrics+    atomically $ writeTVar connState.lastMetrics metrics+  Subscribe pids -> do+    atomically $ do+      current <- readTVar connState.subscriptions+      let newSubs = case current of+            Nothing -> Just $ Set.fromList pids+            Just existing -> Just $ existing <> Set.fromList pids+      writeTVar connState.subscriptions newSubs+    -- Send snapshot of subscribed processors+    allMetrics <- getAllMetricsIO master+    let filtered = Map.filterWithKey (\pid _ -> pid `elem` pids) allMetrics+    WS.sendTextData conn $ encode $ MetricsSnapshot filtered+  Unsubscribe pids -> do+    atomically $ do+      current <- readTVar connState.subscriptions+      case current of+        Nothing -> do+          -- Was subscribed to all, now remove these+          -- We need all processor IDs to calculate the new set+          pure () -- Keep as Nothing, will filter in push+        Just existing ->+          writeTVar connState.subscriptions $+            Just $+              Set.difference existing (Set.fromList pids)+  Ping ->+    WS.sendTextData conn $ encode Pong++--------------------------------------------------------------------------------+-- Push Loop+--------------------------------------------------------------------------------++-- | Push metrics updates to client at configured interval.+pushLoop ::+  MetricsServerConfig ->+  Master ->+  ConnectionState ->+  WS.Connection ->+  IO ()+pushLoop config master connState conn = forever $ do+  threadDelay config.wsPushIntervalUs+  -- Get current metrics+  currentMetrics <- getAllMetricsIO master+  -- Get subscription filter+  mSubs <- atomically $ readTVar connState.subscriptions+  -- Get last sent metrics+  lastSent <- atomically $ readTVar connState.lastMetrics+  -- Filter metrics based on subscriptions+  let filteredMetrics = case mSubs of+        Nothing -> currentMetrics+        Just subs -> Map.filterWithKey (\pid _ -> Set.member pid subs) currentMetrics+  -- Send updates for changed processors+  _ <- Map.traverseWithKey (sendIfChanged lastSent conn) filteredMetrics+  -- Update last sent+  atomically $ writeTVar connState.lastMetrics filteredMetrics++-- | 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