150 lines
4.6 KiB
Haskell
150 lines
4.6 KiB
Haskell
{-# LANGUAGE RecordWildCards #-}
|
|
|
|
module Aenebris.Backend
|
|
( BackendState(..)
|
|
, RuntimeBackend(..)
|
|
, createRuntimeBackend
|
|
, isHealthy
|
|
, trackConnection
|
|
, getConnectionCount
|
|
, getCurrentWeight
|
|
, transitionToUnhealthy
|
|
, transitionToRecovering
|
|
, transitionToHealthy
|
|
, recordFailure
|
|
, recordSuccess
|
|
) where
|
|
|
|
import Aenebris.Config (Server(..))
|
|
import Control.Concurrent.STM
|
|
import Control.Exception (bracket_)
|
|
import Control.Monad (when)
|
|
import Data.Text (Text)
|
|
import Data.Time.Clock (UTCTime)
|
|
|
|
data BackendState
|
|
= Healthy
|
|
| Unhealthy
|
|
| Recovering
|
|
deriving (Eq, Show)
|
|
|
|
-- | Runtime backend state wrapping config Server
|
|
data RuntimeBackend = RuntimeBackend
|
|
{ rbServerId :: Int -- Unique identifier
|
|
, rbHost :: Text
|
|
, rbWeight :: Int
|
|
-- Runtime state (STM)
|
|
, rbActiveConnections :: TVar Int
|
|
, rbCurrentWeight :: TVar Int
|
|
, rbHealthState :: TVar BackendState
|
|
, rbConsecutiveFailures :: TVar Int
|
|
, rbConsecutiveSuccesses :: TVar Int
|
|
, rbLastHealthCheck :: TVar (Maybe UTCTime)
|
|
, rbTotalRequests :: TVar Int -- For metrics
|
|
, rbTotalFailures :: TVar Int -- For metrics
|
|
}
|
|
|
|
instance Show RuntimeBackend where
|
|
show rb = "RuntimeBackend {id=" ++ show (rbServerId rb) ++
|
|
", host=" ++ show (rbHost rb) ++ "}"
|
|
|
|
instance Eq RuntimeBackend where
|
|
rb1 == rb2 = rbServerId rb1 == rbServerId rb2
|
|
|
|
-- | Runtime backend from config Server
|
|
createRuntimeBackend :: Int -> Server -> IO RuntimeBackend
|
|
createRuntimeBackend serverId Server{..} = do
|
|
atomically $ RuntimeBackend serverId serverHost serverWeight
|
|
<$> newTVar 0 -- activeConnections
|
|
<*> newTVar 0 -- currentWeight (for smooth WRR)
|
|
<*> newTVar Healthy -- healthState
|
|
<*> newTVar 0 -- consecutiveFailures
|
|
<*> newTVar 0 -- consecutiveSuccesses
|
|
<*> newTVar Nothing -- lastHealthCheck
|
|
<*> newTVar 0 -- totalRequests
|
|
<*> newTVar 0 -- totalFailures
|
|
|
|
-- | Check if backend is healthy
|
|
isHealthy :: RuntimeBackend -> STM Bool
|
|
isHealthy rb = (== Healthy) <$> readTVar (rbHealthState rb)
|
|
|
|
-- | Track a connection (increment on start, decrement on end)
|
|
trackConnection :: RuntimeBackend -> IO a -> IO a
|
|
trackConnection rb action =
|
|
bracket_
|
|
(atomically $ do
|
|
modifyTVar' (rbActiveConnections rb) (+1)
|
|
modifyTVar' (rbTotalRequests rb) (+1))
|
|
(atomically $ modifyTVar' (rbActiveConnections rb) (subtract 1))
|
|
action
|
|
|
|
-- | Get current connection count
|
|
getConnectionCount :: RuntimeBackend -> STM Int
|
|
getConnectionCount rb = readTVar (rbActiveConnections rb)
|
|
|
|
-- | Get current weight (for smooth weighted RR)
|
|
getCurrentWeight :: RuntimeBackend -> STM Int
|
|
getCurrentWeight rb = readTVar (rbCurrentWeight rb)
|
|
|
|
-- | State transition: mark as unhealthy
|
|
transitionToUnhealthy :: RuntimeBackend -> STM ()
|
|
transitionToUnhealthy rb = do
|
|
writeTVar (rbHealthState rb) Unhealthy
|
|
writeTVar (rbConsecutiveFailures rb) 0
|
|
writeTVar (rbConsecutiveSuccesses rb) 0
|
|
|
|
-- | State transition: start recovering
|
|
transitionToRecovering :: RuntimeBackend -> STM ()
|
|
transitionToRecovering rb = do
|
|
writeTVar (rbHealthState rb) Recovering
|
|
writeTVar (rbConsecutiveSuccesses rb) 1
|
|
|
|
-- | State transition: mark as healthy
|
|
transitionToHealthy :: RuntimeBackend -> STM ()
|
|
transitionToHealthy rb = do
|
|
writeTVar (rbHealthState rb) Healthy
|
|
writeTVar (rbConsecutiveFailures rb) 0
|
|
writeTVar (rbConsecutiveSuccesses rb) 0
|
|
|
|
-- | Record a health check failure
|
|
recordFailure :: RuntimeBackend -> Int -> STM ()
|
|
recordFailure rb maxFailures = do
|
|
state <- readTVar (rbHealthState rb)
|
|
failures <- readTVar (rbConsecutiveFailures rb)
|
|
|
|
case state of
|
|
Healthy -> do
|
|
let newFailures = failures + 1
|
|
writeTVar (rbConsecutiveFailures rb) newFailures
|
|
when (newFailures >= maxFailures) $
|
|
transitionToUnhealthy rb
|
|
|
|
Recovering -> do
|
|
-- Failed during recovery, back to unhealthy
|
|
transitionToUnhealthy rb
|
|
|
|
Unhealthy ->
|
|
-- Already unhealthy, just record it
|
|
modifyTVar' (rbTotalFailures rb) (+1)
|
|
|
|
-- | Record a health check success
|
|
recordSuccess :: RuntimeBackend -> Int -> STM ()
|
|
recordSuccess rb recoveryAttempts = do
|
|
state <- readTVar (rbHealthState rb)
|
|
successes <- readTVar (rbConsecutiveSuccesses rb)
|
|
|
|
case state of
|
|
Healthy ->
|
|
-- Reset failure counter
|
|
writeTVar (rbConsecutiveFailures rb) 0
|
|
|
|
Unhealthy ->
|
|
-- First success, transition to recovering
|
|
transitionToRecovering rb
|
|
|
|
Recovering -> do
|
|
let newSuccesses = successes + 1
|
|
writeTVar (rbConsecutiveSuccesses rb) newSuccesses
|
|
when (newSuccesses >= recoveryAttempts) $
|
|
transitionToHealthy rb
|