Routing.hs

-- | Cluster-aware routing: automatic server discovery and failover for Neo4j causal clusters.
module Database.Bolty.Routing
  ( -- * Routing table
    getRoutingTable
  , RoutingTable(..)
    -- * Access mode
  , AccessMode(..)
    -- * Routing pool
  , RoutingPool(..)
  , RoutingPoolConfig(..)
  , defaultRoutingPoolConfig
  , createRoutingPool
  , destroyRoutingPool
  , withRoutingConnection
  , acquireRoutingConnection
  , withRoutingTransaction
  , invalidateRoutingTable
    -- * Internals (exported for testing)
  , parseAddress
  ) where

import           Control.Concurrent             (threadDelay)
import           Control.Concurrent.MVar        (MVar, newMVar, withMVar)
import           Control.Exception              (SomeException, throwIO, try, fromException,
                                                 onException)
import           Control.Monad                  (when)
import           Data.HashMap.Lazy              (HashMap)
import qualified Data.HashMap.Lazy              as H
import           Data.IORef                     (IORef, newIORef, readIORef, writeIORef, atomicModifyIORef')
import           Data.Kind                      (Type)
import           Data.Text                      (Text)
import           Data.List                      (isInfixOf)
import qualified Data.Text                      as T
import qualified Data.Vector                    as V
import           Data.Word                      (Word16, Word64)
import           GHC.Clock                      (getMonotonicTimeNSec)
import           GHC.Stack                      (HasCallStack)

import           Database.Bolty.Connection.Type
import qualified Database.Bolty.Connection.Pipe as P
import           Database.Bolty.Message.Request (Request(..), Route(..), RouteExtra(..), Begin(..))
import           Database.Bolty.Message.Response (Response(..), Failure(..), RoutingTable(..), parseRoutingTable)
import           Database.Bolty.Pool


-- | Access mode for routing: determines whether to use reader or writer servers.
type AccessMode :: Type
data AccessMode = ReadAccess | WriteAccess
  deriving stock (Show, Eq)


-- | Fetch a routing table from the server. The connection must be in Ready state.
getRoutingTable :: HasCallStack => Connection -> Maybe Text -> IO RoutingTable
getRoutingTable conn dbName = do
  P.requireStateIO conn [Ready] "ROUTE"
  P.flushIO conn $ RRoute Route
    { routing   = H.empty
    , bookmarks = V.empty
    , extra     = RouteExtra{db = dbName, imp_user = Nothing}
    }
  response <- P.fetchIO conn
  case response of
    RSuccess meta ->
      case parseRoutingTable meta of
        Right rt -> pure rt
        Left err -> throwIO $ RoutingTableError err
    RFailure Failure{code, message} -> do
      P.setState conn Failed
      P.reset conn
      throwIO $ ResponseErrorFailure code message
    _ -> do
      P.reset conn
      throwIO $ WrongMessageFormat "Unexpected response to ROUTE"


-- | Configuration for a routing-aware connection pool.
type RoutingPoolConfig :: Type
data RoutingPoolConfig = RoutingPoolConfig
  { poolConfig    :: PoolConfig
  -- ^ Per-address pool configuration.
  , routingDb     :: Maybe Text
  -- ^ Database to route for (Nothing = default database).
  , refreshBuffer :: Int
  -- ^ Seconds before TTL expiry to proactively refresh the routing table.
  }

-- | Default routing pool configuration.
defaultRoutingPoolConfig :: RoutingPoolConfig
defaultRoutingPoolConfig = RoutingPoolConfig
  { poolConfig    = defaultPoolConfig
  , routingDb     = Nothing
  , refreshBuffer = 10
  }


-- | Cached routing table with monotonic expiry time.
type CachedRoutingTable :: Type
data CachedRoutingTable = CachedRoutingTable
  { cachedTable  :: !RoutingTable
  , expiresAtNs  :: !Word64
  }

-- | A routing-aware connection pool that directs connections based on access mode.
type RoutingPool :: Type
data RoutingPool = RoutingPool
  { rpConfig       :: !ValidatedConfig
  , rpPoolConfig   :: !PoolConfig
  , rpRoutingDb    :: !(Maybe Text)
  , rpRefreshBuf   :: !Int
  , rpCacheRef     :: !(IORef (Maybe CachedRoutingTable))
  , rpPoolsRef     :: !(IORef (HashMap Text BoltPool))
  , rpRefreshLock  :: !(MVar ())
  , rpCounter      :: !(IORef Int)
  }


-- | Create a routing-aware connection pool. Connects to the seed address,
-- fetches the initial routing table, and sets up per-address pools.
createRoutingPool :: ValidatedConfig -> RoutingPoolConfig -> IO RoutingPool
createRoutingPool cfg RoutingPoolConfig{..} = do
  let routingCfg = setRouting cfg Routing
  cacheRef <- newIORef Nothing
  poolsRef <- newIORef H.empty
  lock     <- newMVar ()
  counter  <- newIORef 0
  let rp = RoutingPool
        { rpConfig     = routingCfg
        , rpPoolConfig = poolConfig
        , rpRoutingDb  = routingDb
        , rpRefreshBuf = refreshBuffer
        , rpCacheRef   = cacheRef
        , rpPoolsRef   = poolsRef
        , rpRefreshLock = lock
        , rpCounter    = counter
        }
  -- Bootstrap: fetch routing table from seed
  _ <- refreshRoutingTable rp
  pure rp


-- | Destroy all per-address pools in the routing pool.
destroyRoutingPool :: RoutingPool -> IO ()
destroyRoutingPool RoutingPool{rpPoolsRef} = do
  pools <- readIORef rpPoolsRef
  mapM_ destroyPool pools
  writeIORef rpPoolsRef H.empty


-- | Invalidate the cached routing table, forcing a refresh on the next operation.
-- Use this when a routing error (e.g. NotALeader) indicates the table is stale.
invalidateRoutingTable :: RoutingPool -> IO ()
invalidateRoutingTable RoutingPool{rpCacheRef} = writeIORef rpCacheRef Nothing


-- | Pick the next address from a vector using round-robin.
roundRobin :: RoutingPool -> V.Vector Text -> IO Text
roundRobin RoutingPool{rpCounter} addrs = do
  idx <- atomicModifyIORef' rpCounter $ \n -> (n + 1, n)
  pure $ addrs V.! (idx `mod` V.length addrs)


-- | Acquire a connection routed by access mode, run an action, then release.
-- On connection failure, tries the next address in round-robin order
-- until all addresses are exhausted.
withRoutingConnection :: HasCallStack => RoutingPool -> AccessMode -> (Connection -> IO a) -> IO a
withRoutingConnection rp mode action = do
  rt <- getOrRefreshTable rp
  let addrs = case mode of
        ReadAccess  -> readers rt
        WriteAccess -> writers rt
  when (V.null addrs) $
    throwIO $ RoutingTableError $ "No servers available for " <> T.pack (show mode)
  tryAddresses rp addrs (V.length addrs) Nothing action


-- | Acquire a routed connection by access mode. Returns a
-- 'CheckedOutConnection' that must be released by the caller.
-- Tries addresses in round-robin order, failing over on unavailable servers.
acquireRoutingConnection :: HasCallStack => RoutingPool -> AccessMode -> IO CheckedOutConnection
acquireRoutingConnection rp mode = do
  rt <- getOrRefreshTable rp
  let addrs = case mode of
        ReadAccess  -> readers rt
        WriteAccess -> writers rt
  when (V.null addrs) $
    throwIO $ RoutingTableError $ "No servers available for " <> T.pack (show mode)
  tryAcquireAddresses rp addrs (V.length addrs) Nothing


-- | Try addresses in round-robin order for acquire, failing over on connection errors.
tryAcquireAddresses :: HasCallStack
                    => RoutingPool -> V.Vector Text -> Int -> Maybe SomeException -> IO CheckedOutConnection
tryAcquireAddresses _rp _addrs 0 (Just lastErr) = throwIO lastErr
tryAcquireAddresses _rp _addrs 0 Nothing =
  throwIO $ RoutingTableError "All servers unavailable"
tryAcquireAddresses rp addrs remaining _lastErr = do
  addr <- roundRobin rp addrs
  pool <- getOrCreatePool rp addr
  result <- try $ acquireConnection pool
  case result of
    Right coc -> pure coc
    Left (e :: SomeException)
      | remaining > 1, isServerUnavailable e ->
          tryAcquireAddresses rp addrs (remaining - 1) (Just e)
      | otherwise -> throwIO e


-- | Try addresses in round-robin order, failing over on connection errors.
tryAddresses :: HasCallStack
             => RoutingPool -> V.Vector Text -> Int -> Maybe SomeException -> (Connection -> IO a) -> IO a
tryAddresses _rp _addrs 0 (Just lastErr) _action = throwIO lastErr
tryAddresses _rp _addrs 0 Nothing _action =
  throwIO $ RoutingTableError "All servers unavailable"
tryAddresses rp addrs remaining _lastErr action = do
  addr <- roundRobin rp addrs
  pool <- getOrCreatePool rp addr
  result <- try $ withConnection pool action
  case result of
    Right x -> pure x
    Left (e :: SomeException)
      | remaining > 1, isServerUnavailable e ->
          tryAddresses rp addrs (remaining - 1) (Just e) action
      | otherwise -> throwIO e


-- | Check if an exception indicates the server is unavailable (connection refused,
-- broken pipe, no healthy connection in pool, etc.) as opposed to a query-level error.
isServerUnavailable :: SomeException -> Bool
isServerUnavailable e = case fromException e :: Maybe Error of
  Just (NonboltyError _) -> True   -- IO exceptions (connection refused, broken pipe, etc.)
  _ -> "no healthy connection" `isInfixOf` show e


-- | Run a retrying transaction routed by access mode.
-- Re-acquires routing table and connection on each retry attempt, so that
-- routing errors (NotALeader) and transient errors trigger fresh routing.
withRoutingTransaction :: HasCallStack => RoutingPool -> AccessMode -> (Connection -> IO a) -> IO a
withRoutingTransaction rp mode action =
  let rc = retryConfig (rpPoolConfig rp)
      maxR = maxRetries rc
      initD = initialDelay rc
      maxD = maxDelay rc
  in go maxR initD maxD
  where
    go 0 _ _    = attempt
    go n delay maxD = do
      result <- try attempt
      case result of
        Right x -> pure x
        Left (e :: SomeException) -> case fromException e :: Maybe Error of
          Just err
            | isTransient err -> do
                threadDelay delay
                go (n - 1) (min maxD (delay * 2)) maxD
            | isRoutingError err -> do
                invalidateRoutingTable rp
                threadDelay delay
                go (n - 1) (min maxD (delay * 2)) maxD
          _ -> throwIO e

    modeChar = case mode of
      ReadAccess  -> 'r'
      WriteAccess -> 'w'

    attempt = do
      rt <- getOrRefreshTable rp
      let addrs = case mode of
            ReadAccess  -> readers rt
            WriteAccess -> writers rt
      when (V.null addrs) $
        throwIO $ RoutingTableError $ "No servers available for " <> T.pack (show mode)
      addr <- roundRobin rp addrs
      pool <- getOrCreatePool rp addr
      withConnection pool $ \conn -> do
        P.beginTx conn $ Begin V.empty Nothing H.empty modeChar Nothing Nothing
        result <- action conn `onException` P.tryRollback conn
        _ <- P.commitTx conn
        pure result


-- | Get the cached routing table, refreshing if expired.
getOrRefreshTable :: RoutingPool -> IO RoutingTable
getOrRefreshTable rp@RoutingPool{rpCacheRef, rpRefreshBuf} = do
  cached <- readIORef rpCacheRef
  now <- getMonotonicTimeNSec
  let bufferNs = fromIntegral rpRefreshBuf * 1_000_000_000
  case cached of
    Just CachedRoutingTable{cachedTable, expiresAtNs}
      | now + bufferNs < expiresAtNs -> pure cachedTable
    _ -> refreshRoutingTable rp


-- | Refresh the routing table. Tries known router addresses from the cached
-- table first, then falls back to the seed address from the initial config.
refreshRoutingTable :: RoutingPool -> IO RoutingTable
refreshRoutingTable rp@RoutingPool{rpConfig, rpCacheRef, rpRefreshLock} =
  withMVar rpRefreshLock $ \_ -> do
    -- Double-check after acquiring lock (another thread may have refreshed)
    cached <- readIORef rpCacheRef
    now <- getMonotonicTimeNSec
    let bufferNs = fromIntegral (rpRefreshBuf rp) * 1_000_000_000
    case cached of
      Just CachedRoutingTable{cachedTable, expiresAtNs}
        | now + bufferNs < expiresAtNs -> pure cachedTable
      _ -> do
        -- Build list of addresses to try: known routers first, seed last
        let knownRouters = case cached of
              Just CachedRoutingTable{cachedTable} -> V.toList (routers cachedTable)
              Nothing -> []
        let seedAddr = seedAddress rpConfig
        let allAddrs = dedupAddrs (knownRouters <> [seedAddr])
        tryRefreshFrom rp allAddrs now


-- | Try to refresh the routing table from a list of router addresses.
-- Tries each in order, falling through on failure.
tryRefreshFrom :: RoutingPool -> [Text] -> Word64 -> IO RoutingTable
tryRefreshFrom _rp [] _now =
  throwIO $ RoutingTableError "Could not reach any router to refresh routing table"
tryRefreshFrom rp@RoutingPool{rpConfig, rpRoutingDb, rpCacheRef} (addr:rest) now = do
  let (h, p) = parseAddress addr
  let cfg = setHostPort rpConfig h p
  result <- try $ do
    conn <- P.connect cfg
    rt <- getRoutingTable conn rpRoutingDb
    P.close conn
    pure rt
  case result of
    Right rt -> do
      let expiresAt = now + fromIntegral (ttl rt) * 1_000_000_000
      writeIORef rpCacheRef $ Just CachedRoutingTable{cachedTable = rt, expiresAtNs = expiresAt}
      pure rt
    Left (_ :: SomeException)
      | not (null rest) -> tryRefreshFrom rp rest now
      | otherwise -> throwIO $ RoutingTableError "Could not reach any router to refresh routing table"


-- | Remove duplicate addresses while preserving order.
dedupAddrs :: [Text] -> [Text]
dedupAddrs = go []
  where
    go :: [Text] -> [Text] -> [Text]
    go _seen [] = []
    go seen (x:xs)
      | x `elem` seen = go seen xs
      | otherwise      = x : go (x:seen) xs


-- | Get or create a pool for a given address.
getOrCreatePool :: RoutingPool -> Text -> IO BoltPool
getOrCreatePool RoutingPool{rpConfig, rpPoolConfig, rpPoolsRef} addr = do
  pools <- readIORef rpPoolsRef
  case H.lookup addr pools of
    Just pool -> pure pool
    Nothing   -> do
      let (h, p) = parseAddress addr
      let cfg = setHostPort rpConfig h p
      pool <- createPool cfg rpPoolConfig
      atomicModifyIORef' rpPoolsRef $ \ps -> (H.insert addr pool ps, ())
      pure pool


-- | Parse "host:port" into (host, port). Falls back to default port 7687.
parseAddress :: Text -> (Text, Word16)
parseAddress addr =
  case T.splitOn ":" addr of
    [h, p] -> case reads (T.unpack p) of
                [(port, "")] -> (h, port)
                _            -> (h, 7687)
    _       -> (addr, 7687)


-- | Set the routing field on a ValidatedConfig (avoids ambiguous record update).
setRouting :: ValidatedConfig -> Routing -> ValidatedConfig
setRouting ValidatedConfig{host, port, scheme, use_tls, versions, timeout, user_agent, queryLogger, notificationHandler} r =
  ValidatedConfig{host, port, scheme, use_tls, versions, timeout, routing = r, user_agent, queryLogger, notificationHandler}


-- | Set host and port on a ValidatedConfig (avoids ambiguous record update).
setHostPort :: ValidatedConfig -> Text -> Word16 -> ValidatedConfig
setHostPort ValidatedConfig{scheme, use_tls, versions, timeout, routing, user_agent, queryLogger, notificationHandler} h p =
  ValidatedConfig{host = h, port = p, scheme, use_tls, versions, timeout, routing, user_agent, queryLogger, notificationHandler}


-- | Extract "host:port" from a ValidatedConfig (avoids ambiguous record field access).
seedAddress :: ValidatedConfig -> Text
seedAddress ValidatedConfig{host = h, port = p} = h <> ":" <> T.pack (show p)