-- This file is part of the Wire Server implementation.
--
-- Copyright (C) 2025 Wire Swiss GmbH <opensource@wire.com>
--
-- This program is free software: you can redistribute it and/or modify it under
-- the terms of the GNU Affero General Public License as published by the Free
-- Software Foundation, either version 3 of the License, or (at your option) any
-- later version.
--
-- This program is distributed in the hope that it will be useful, but WITHOUT
-- ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
-- FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
-- details.
--
-- You should have received a copy of the GNU Affero General Public License along
-- with this program. If not, see <https://www.gnu.org/licenses/>.

module Hasql.Pool.Extended where

import Data.Aeson
import Data.Map as Map
import Data.Misc
import Hasql.Connection qualified
import Hasql.Connection.Settings qualified as HasqlConnSettings
import Hasql.Pool qualified as HasqlPool
import Imports
import PostgresqlConnectionString qualified
import Prometheus
import UnliftIO.IO (getMonotonicTime)
import Util.Options

data PoolConfig = PoolConfig
  { PoolConfig -> Int
size :: Int,
    -- | Configured pool acquisition wait time. hasql-resource-pool only
    -- accepts whole seconds here, so we round up to the nearest second and
    -- pass it through as the pool acquisition timeout.
    PoolConfig -> Duration
acquisitionTimeout :: Duration,
    -- | Controls how long idle connections stay resident in the pool.
    PoolConfig -> Duration
idlenessTimeout :: Duration
  }
  deriving (PoolConfig -> PoolConfig -> Bool
(PoolConfig -> PoolConfig -> Bool)
-> (PoolConfig -> PoolConfig -> Bool) -> Eq PoolConfig
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: PoolConfig -> PoolConfig -> Bool
== :: PoolConfig -> PoolConfig -> Bool
$c/= :: PoolConfig -> PoolConfig -> Bool
/= :: PoolConfig -> PoolConfig -> Bool
Eq, Int -> PoolConfig -> ShowS
[PoolConfig] -> ShowS
PoolConfig -> String
(Int -> PoolConfig -> ShowS)
-> (PoolConfig -> String)
-> ([PoolConfig] -> ShowS)
-> Show PoolConfig
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> PoolConfig -> ShowS
showsPrec :: Int -> PoolConfig -> ShowS
$cshow :: PoolConfig -> String
show :: PoolConfig -> String
$cshowList :: [PoolConfig] -> ShowS
showList :: [PoolConfig] -> ShowS
Show)

instance FromJSON PoolConfig where
  parseJSON :: Value -> Parser PoolConfig
parseJSON = String
-> (Object -> Parser PoolConfig) -> Value -> Parser PoolConfig
forall a. String -> (Object -> Parser a) -> Value -> Parser a
withObject String
"PoolConfig" ((Object -> Parser PoolConfig) -> Value -> Parser PoolConfig)
-> (Object -> Parser PoolConfig) -> Value -> Parser PoolConfig
forall a b. (a -> b) -> a -> b
$ \Object
o ->
    Int -> Duration -> Duration -> PoolConfig
PoolConfig
      (Int -> Duration -> Duration -> PoolConfig)
-> Parser Int -> Parser (Duration -> Duration -> PoolConfig)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Object
o Object -> Key -> Parser Int
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"size"
      Parser (Duration -> Duration -> PoolConfig)
-> Parser Duration -> Parser (Duration -> PoolConfig)
forall a b. Parser (a -> b) -> Parser a -> Parser b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Object
o Object -> Key -> Parser Duration
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"acquisitionTimeout"
      Parser (Duration -> PoolConfig)
-> Parser Duration -> Parser PoolConfig
forall a b. Parser (a -> b) -> Parser a -> Parser b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Object
o Object -> Key -> Parser Duration
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"idlenessTimeout"

data HasqlPoolMetrics = HasqlPoolMetrics
  { HasqlPoolMetrics -> Gauge
readyForUseGauge :: Gauge,
    HasqlPoolMetrics -> Gauge
inUseGauge :: Gauge,
    HasqlPoolMetrics -> Counter
establishedCounter :: Counter,
    HasqlPoolMetrics -> Counter
connectionFailureCounter :: Counter,
    HasqlPoolMetrics -> Counter
acquisitionTimeoutCounter :: Counter,
    HasqlPoolMetrics -> Counter
sessionFailureCounter :: Counter,
    HasqlPoolMetrics -> Counter
sessionCounter :: Counter,
    HasqlPoolMetrics -> Histogram
connectionAcquisitionDuration :: Histogram,
    HasqlPoolMetrics -> Histogram
sessionDuration :: Histogram
  }

data Pool = Pool
  { Pool -> Pool
rawPool :: HasqlPool.Pool,
    Pool -> HasqlPoolMetrics
metrics :: HasqlPoolMetrics,
    -- | Pool acquisition timeout in seconds, rounded up from the configured
    -- duration. This is used by the session runner to bound waiting for an
    -- available connection slot.
    Pool -> Duration
poolAcquisitionTimeout :: Duration
  }

recordHasqlPoolConnectionAcquisition :: HasqlPoolMetrics -> Double -> IO ()
recordHasqlPoolConnectionAcquisition :: HasqlPoolMetrics -> Double -> IO ()
recordHasqlPoolConnectionAcquisition HasqlPoolMetrics
metrics Double
secs =
  Histogram -> Double -> IO ()
forall metric (m :: * -> *).
(Observer metric, MonadMonitor m) =>
metric -> Double -> m ()
forall (m :: * -> *). MonadMonitor m => Histogram -> Double -> m ()
observe HasqlPoolMetrics
metrics.connectionAcquisitionDuration Double
secs

recordHasqlPoolConnectionEstablished :: HasqlPoolMetrics -> IO ()
recordHasqlPoolConnectionEstablished :: HasqlPoolMetrics -> IO ()
recordHasqlPoolConnectionEstablished HasqlPoolMetrics
metrics =
  IO Bool -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ Counter -> Double -> IO Bool
forall (m :: * -> *). MonadMonitor m => Counter -> Double -> m Bool
addCounter HasqlPoolMetrics
metrics.establishedCounter Double
1

recordHasqlPoolConnectionFailure :: HasqlPoolMetrics -> IO ()
recordHasqlPoolConnectionFailure :: HasqlPoolMetrics -> IO ()
recordHasqlPoolConnectionFailure HasqlPoolMetrics
metrics =
  IO Bool -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ Counter -> Double -> IO Bool
forall (m :: * -> *). MonadMonitor m => Counter -> Double -> m Bool
addCounter HasqlPoolMetrics
metrics.connectionFailureCounter Double
1

recordHasqlPoolSessionStarted :: HasqlPoolMetrics -> IO ()
recordHasqlPoolSessionStarted :: HasqlPoolMetrics -> IO ()
recordHasqlPoolSessionStarted HasqlPoolMetrics
metrics =
  IO Bool -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ Counter -> Double -> IO Bool
forall (m :: * -> *). MonadMonitor m => Counter -> Double -> m Bool
addCounter HasqlPoolMetrics
metrics.sessionCounter Double
1

recordHasqlPoolSessionFailure :: HasqlPoolMetrics -> IO ()
recordHasqlPoolSessionFailure :: HasqlPoolMetrics -> IO ()
recordHasqlPoolSessionFailure HasqlPoolMetrics
metrics =
  IO Bool -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ Counter -> Double -> IO Bool
forall (m :: * -> *). MonadMonitor m => Counter -> Double -> m Bool
addCounter HasqlPoolMetrics
metrics.sessionFailureCounter Double
1

recordHasqlPoolSessionDuration :: HasqlPoolMetrics -> Double -> IO ()
recordHasqlPoolSessionDuration :: HasqlPoolMetrics -> Double -> IO ()
recordHasqlPoolSessionDuration HasqlPoolMetrics
metrics Double
secs =
  Histogram -> Double -> IO ()
forall metric (m :: * -> *).
(Observer metric, MonadMonitor m) =>
metric -> Double -> m ()
forall (m :: * -> *). MonadMonitor m => Histogram -> Double -> m ()
observe HasqlPoolMetrics
metrics.sessionDuration Double
secs

recordHasqlPoolAcquisitionTimeout :: HasqlPoolMetrics -> IO ()
recordHasqlPoolAcquisitionTimeout :: HasqlPoolMetrics -> IO ()
recordHasqlPoolAcquisitionTimeout HasqlPoolMetrics
metrics =
  IO Bool -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ Counter -> Double -> IO Bool
forall (m :: * -> *). MonadMonitor m => Counter -> Double -> m Bool
addCounter HasqlPoolMetrics
metrics.acquisitionTimeoutCounter Double
1

recordHasqlPoolStats :: Pool -> IO ()
recordHasqlPoolStats :: Pool -> IO ()
recordHasqlPoolStats Pool
pool = do
  -- hasql-resource-pool does not expose per-acquire/release callbacks, so
  -- these gauges are refreshed from the pool's current total connections stats instead.
  poolStats <- Pool -> IO Stats
HasqlPool.stats Pool
pool.rawPool
  setGauge pool.metrics.readyForUseGauge (fromIntegral poolStats.available)
  setGauge pool.metrics.inUseGauge (fromIntegral poolStats.currentUsage)

startHasqlPoolStatsReporter :: Pool -> IO ()
startHasqlPoolStatsReporter :: Pool -> IO ()
startHasqlPoolStatsReporter Pool
pool = IO ThreadId -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO ThreadId -> IO ()) -> IO ThreadId -> IO ()
forall a b. (a -> b) -> a -> b
$ IO () -> IO ThreadId
forall (m :: * -> *). MonadUnliftIO m => m () -> m ThreadId
forkIO (IO () -> IO ThreadId) -> IO () -> IO ThreadId
forall a b. (a -> b) -> a -> b
$ IO () -> IO ()
forall (f :: * -> *) a b. Applicative f => f a -> f b
forever (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
  Pool -> IO ()
recordHasqlPoolStats Pool
pool
  Int -> IO ()
forall (m :: * -> *). MonadIO m => Int -> m ()
threadDelay (Int
5 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1_000_000) -- 5s

-- | Creates a pool from postgres config params.
--
-- 'acquisitionTimeout' is mapped to the pool acquisition timeout,
-- 'idlenessTimeout' controls how long idle connections stay resident.
initPostgresPool :: PoolConfig -> Map Text Text -> Maybe FilePathSecrets -> IO Pool
initPostgresPool :: PoolConfig -> Map Text Text -> Maybe FilePathSecrets -> IO Pool
initPostgresPool PoolConfig
config Map Text Text
pgConfig Maybe FilePathSecrets
mFpSecrets = do
  mPw <- Maybe FilePathSecrets
-> (FilePathSecrets -> IO Text) -> IO (Maybe Text)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for Maybe FilePathSecrets
mFpSecrets FilePathSecrets -> IO Text
forall (m :: * -> *) a.
(MonadIO m, FromJSON a) =>
FilePathSecrets -> m a
initCredentials
  let pgSettings =
        Text -> Settings
HasqlConnSettings.connectionString (ConnectionString -> Text
PostgresqlConnectionString.toUrl (ConnectionString -> Text) -> ConnectionString -> Text
forall a b. (a -> b) -> a -> b
$ Map Text Text -> ConnectionString
PostgresqlConnectionString.fromKeyValueParams Map Text Text
pgConfig)
          Settings -> Settings -> Settings
forall a. Semigroup a => a -> a -> a
<> (Text -> Settings) -> Maybe Text -> Settings
forall m a. Monoid m => (a -> m) -> Maybe a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Text -> Settings
HasqlConnSettings.password Maybe Text
mPw
  metrics <- mkHasqlPoolMetrics
  rawPool <-
    HasqlPool.acquireWith
      (instrumentedConnectionGetter metrics (Hasql.Connection.acquire pgSettings))
      ( config.size,
        realToFrac config.idlenessTimeout.duration,
        unusedSettings
      )
  let pool = Pool {Pool
rawPool :: Pool
rawPool :: Pool
rawPool, HasqlPoolMetrics
metrics :: HasqlPoolMetrics
metrics :: HasqlPoolMetrics
metrics, poolAcquisitionTimeout :: Duration
poolAcquisitionTimeout = PoolConfig
config.acquisitionTimeout}
  startHasqlPoolStatsReporter pool
  pure pool
  where
    instrumentedConnectionGetter :: HasqlPoolMetrics -> IO (Either a b) -> IO (Either a b)
instrumentedConnectionGetter HasqlPoolMetrics
metrics IO (Either a b)
getter = do
      started <- IO Double
forall (m :: * -> *). MonadIO m => m Double
getMonotonicTime
      res <- getter
      ended <- getMonotonicTime
      recordHasqlPoolConnectionAcquisition metrics (ended - started)
      case res of
        Right b
_ -> HasqlPoolMetrics -> IO ()
recordHasqlPoolConnectionEstablished HasqlPoolMetrics
metrics
        Left a
_ -> HasqlPoolMetrics -> IO ()
recordHasqlPoolConnectionFailure HasqlPoolMetrics
metrics
      pure res

    mkHasqlPoolMetrics :: IO HasqlPoolMetrics
    mkHasqlPoolMetrics :: IO HasqlPoolMetrics
mkHasqlPoolMetrics =
      Gauge
-> Gauge
-> Counter
-> Counter
-> Counter
-> Counter
-> Counter
-> Histogram
-> Histogram
-> HasqlPoolMetrics
HasqlPoolMetrics
        (Gauge
 -> Gauge
 -> Counter
 -> Counter
 -> Counter
 -> Counter
 -> Counter
 -> Histogram
 -> Histogram
 -> HasqlPoolMetrics)
-> IO Gauge
-> IO
     (Gauge
      -> Counter
      -> Counter
      -> Counter
      -> Counter
      -> Counter
      -> Histogram
      -> Histogram
      -> HasqlPoolMetrics)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Metric Gauge -> IO Gauge
forall (m :: * -> *) s. MonadIO m => Metric s -> m s
register (Info -> Metric Gauge
gauge (Info -> Metric Gauge) -> Info -> Metric Gauge
forall a b. (a -> b) -> a -> b
$ Text -> Text -> Info
Info Text
"wire_hasql_pool_ready_for_use" Text
"Number of hasql pool connections ready for use")
        IO
  (Gauge
   -> Counter
   -> Counter
   -> Counter
   -> Counter
   -> Counter
   -> Histogram
   -> Histogram
   -> HasqlPoolMetrics)
-> IO Gauge
-> IO
     (Counter
      -> Counter
      -> Counter
      -> Counter
      -> Counter
      -> Histogram
      -> Histogram
      -> HasqlPoolMetrics)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Metric Gauge -> IO Gauge
forall (m :: * -> *) s. MonadIO m => Metric s -> m s
register (Info -> Metric Gauge
gauge (Info -> Metric Gauge) -> Info -> Metric Gauge
forall a b. (a -> b) -> a -> b
$ Text -> Text -> Info
Info Text
"wire_hasql_pool_in_use" Text
"Number of hasql pool connections in use")
        IO
  (Counter
   -> Counter
   -> Counter
   -> Counter
   -> Counter
   -> Histogram
   -> Histogram
   -> HasqlPoolMetrics)
-> IO Counter
-> IO
     (Counter
      -> Counter
      -> Counter
      -> Counter
      -> Histogram
      -> Histogram
      -> HasqlPoolMetrics)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Metric Counter -> IO Counter
forall (m :: * -> *) s. MonadIO m => Metric s -> m s
register (Info -> Metric Counter
counter (Info -> Metric Counter) -> Info -> Metric Counter
forall a b. (a -> b) -> a -> b
$ Text -> Text -> Info
Info Text
"wire_hasql_pool_connection_established_count" Text
"Number of established connections")
        IO
  (Counter
   -> Counter
   -> Counter
   -> Counter
   -> Histogram
   -> Histogram
   -> HasqlPoolMetrics)
-> IO Counter
-> IO
     (Counter
      -> Counter
      -> Counter
      -> Histogram
      -> Histogram
      -> HasqlPoolMetrics)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Metric Counter -> IO Counter
forall (m :: * -> *) s. MonadIO m => Metric s -> m s
register (Info -> Metric Counter
counter (Info -> Metric Counter) -> Info -> Metric Counter
forall a b. (a -> b) -> a -> b
$ Text -> Text -> Info
Info Text
"wire_hasql_pool_connection_failure_count" Text
"Number of failed connection acquisition attempts")
        IO
  (Counter
   -> Counter
   -> Counter
   -> Histogram
   -> Histogram
   -> HasqlPoolMetrics)
-> IO Counter
-> IO
     (Counter -> Counter -> Histogram -> Histogram -> HasqlPoolMetrics)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Metric Counter -> IO Counter
forall (m :: * -> *) s. MonadIO m => Metric s -> m s
register (Info -> Metric Counter
counter (Info -> Metric Counter) -> Info -> Metric Counter
forall a b. (a -> b) -> a -> b
$ Text -> Text -> Info
Info Text
"wire_hasql_pool_acquisition_timeout_count" Text
"Number of pool acquisition timeouts")
        IO
  (Counter -> Counter -> Histogram -> Histogram -> HasqlPoolMetrics)
-> IO Counter
-> IO (Counter -> Histogram -> Histogram -> HasqlPoolMetrics)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Metric Counter -> IO Counter
forall (m :: * -> *) s. MonadIO m => Metric s -> m s
register (Info -> Metric Counter
counter (Info -> Metric Counter) -> Info -> Metric Counter
forall a b. (a -> b) -> a -> b
$ Text -> Text -> Info
Info Text
"wire_hasql_pool_session_failure_count" Text
"Number of times a session has failed")
        IO (Counter -> Histogram -> Histogram -> HasqlPoolMetrics)
-> IO Counter -> IO (Histogram -> Histogram -> HasqlPoolMetrics)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Metric Counter -> IO Counter
forall (m :: * -> *) s. MonadIO m => Metric s -> m s
register (Info -> Metric Counter
counter (Info -> Metric Counter) -> Info -> Metric Counter
forall a b. (a -> b) -> a -> b
$ Text -> Text -> Info
Info Text
"wire_hasql_pool_session_count" Text
"Number of times a session was created")
        IO (Histogram -> Histogram -> HasqlPoolMetrics)
-> IO Histogram -> IO (Histogram -> HasqlPoolMetrics)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Metric Histogram -> IO Histogram
forall (m :: * -> *) s. MonadIO m => Metric s -> m s
register (Info -> [Double] -> Metric Histogram
histogram (Text -> Text -> Info
Info Text
"wire_hasql_pool_connection_acquisition_seconds" Text
"Time spent establishing new PostgreSQL connections") [Double]
defaultBuckets)
        IO (Histogram -> HasqlPoolMetrics)
-> IO Histogram -> IO HasqlPoolMetrics
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Metric Histogram -> IO Histogram
forall (m :: * -> *) s. MonadIO m => Metric s -> m s
register (Info -> [Double] -> Metric Histogram
histogram (Text -> Text -> Info
Info Text
"wire_hasql_pool_session_seconds" Text
"Time spent using PostgreSQL sessions") [Double]
defaultBuckets)

    unusedSettings :: ConnectionSettings
unusedSettings =
      -- The custom getter above performs the actual connection establishment.
      -- The API forces us to pass this record, but it is actually not used in acquireWith
      HasqlPool.ConnectionSettings
        { host :: Text
host = Text
"",
          port :: Word16
port = Word16
5432,
          user :: Text
user = Text
"",
          password :: Text
password = Text
"",
          dbName :: Text
dbName = Text
"",
          connAcqTimeout :: Word16
connAcqTimeout = Word16
0,
          txIdleTimeout :: TimeoutSetting
txIdleTimeout = Word16 -> TimeUnit -> TimeoutSetting
HasqlPool.TimeoutSetting Word16
0 TimeUnit
HasqlPool.Seconds,
          stmtTimeout :: TimeoutSetting
stmtTimeout = Word16 -> TimeUnit -> TimeoutSetting
HasqlPool.TimeoutSetting Word16
0 TimeUnit
HasqlPool.Seconds,
          sslMode :: Text
sslMode = Text
"prefer",
          sslRootCert :: Text
sslRootCert = Text
""
        }