kiroku-store-0.12.0.0: bench/CheckpointTargetCost.hs

{-# LANGUAGE MultilineStrings #-}
{-# LANGUAGE OverloadedRecordDot #-}

-- | Small supplementary EP-2 comparison of the changed checkpoint-save path.
-- No acceptance threshold: local ratios retain their uncertainty. The control
-- statement and row layout are from 23a03a1b8b56773a140d0683c4c8d4d34b9d7e36.
module Main where

import Contravariant.Extras (contrazip4)
import Control.Lens ((^.))
import Control.Monad (forM_, replicateM_)
import Data.Generics.Labels ()
import Data.IORef (atomicModifyIORef', newIORef)
import Data.Int (Int32, Int64)
import Data.Text (Text)
import EphemeralPg qualified as Pg
import GHC.Clock (getMonotonicTimeNSec)
import Hasql.Decoders qualified as D
import Hasql.Encoders qualified as E
import Hasql.Pool qualified as Pool
import Hasql.Session qualified as Session
import Hasql.Statement (Statement, preparable)
import Kiroku.Store
import Kiroku.Store.SQL qualified as SQL
import Kiroku.Test.Postgres (ephemeralConfig, migrateTestDatabase)

main :: IO ()
main = do
  original <- ephemeralConfig
  let keys = ["fsync", "synchronous_commit", "full_page_writes", "shared_buffers", "wal_level"]
      durable = original {Pg.postgresSettings = filter (\(key, _) -> key `notElem` keys) original.postgresSettings <> [("fsync", "on"), ("synchronous_commit", "on"), ("full_page_writes", "on"), ("shared_buffers", "128MB"), ("wal_level", "replica")]}
  result <- Pg.withCachedConfig durable Pg.defaultCacheConfig $ \database -> do
    let connection = Pg.connectionString database
    migrateTestDatabase connection
    withStore (defaultConnectionSettings connection) measure
  either (fail . show) pure result

measure :: KirokuStore -> IO ()
measure store = do
  settings <- use store $ Session.statement () (preparable "SELECT current_setting('server_version_num')::int4 >= 180000 AND current_setting('server_version_num')::int4 < 190000 AND current_setting('fsync') = 'on' AND current_setting('synchronous_commit') = 'on' AND current_setting('full_page_writes') = 'on'" E.noParams (D.singleRow (D.column (D.nonNullable D.bool))))
  if settings then pure () else fail "requires durable PostgreSQL 18"
  use store $
    Session.script
      """
      CREATE SCHEMA ep2_control;
      CREATE TABLE ep2_control.subscriptions (
        subscription_id BIGSERIAL PRIMARY KEY,
        subscription_name TEXT NOT NULL,
        stream_name TEXT NOT NULL DEFAULT '$all',
        last_seen BIGINT NOT NULL DEFAULT 0,
        consumer_group_member INT NOT NULL DEFAULT 0,
        consumer_group_size INT NOT NULL DEFAULT 1,
        created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
        updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
      );
      CREATE UNIQUE INDEX ON ep2_control.subscriptions (subscription_name, consumer_group_member);
      """
  controlCounter <- newIORef 0
  candidateCounter <- newIORef 0
  let save candidate = do
        p <- atomicModifyIORef' (if candidate then candidateCounter else controlCounter) (\n -> (n + 1, n + 1))
        let m = fromIntegral (p `mod` 4)
        if candidate
          then use store $ Session.statement ("target-cost", m, p, 4, "performance") SQL.saveCategoryCheckpointMemberStmt
          else use store $ Session.statement ("target-cost", m, p, 4) controlSave
      trial pair candidate = do
        walBefore <- use store (Session.statement () walPosition)
        start <- getMonotonicTimeNSec
        replicateM_ 2000 (save candidate)
        end <- getMonotonicTimeNSec
        walAfter <- use store (Session.statement () walPosition)
        putStrLn $ show pair <> "," <> (if candidate then "candidate" else "control") <> ",2000," <> show (fromIntegral (end - start) / 2000 :: Double) <> "," <> show (fromIntegral (walAfter - walBefore) / 2000 :: Double)
  replicateM_ 200 (save False >> save True)
  putStrLn "pair,arm,saves,ns_per_save,wal_bytes_per_save"
  forM_ [1 .. 3 :: Int] $ \pair ->
    if odd pair
      then trial pair False >> trial pair True
      else trial pair True >> trial pair False
  -- Every completed save advances one durable member; both arms do equal work.
  expected <- use store $ Session.statement () (preparable "SELECT count(*) = 4 AND sum(last_seen) = 24794 FROM ep2_control.subscriptions" E.noParams (D.singleRow (D.column (D.nonNullable D.bool))))
  actual <- use store $ Session.statement () (preparable "SELECT count(*) = 4 AND sum(last_seen) = 24794 AND bool_and(target_kind = 'category' AND target_category = 'performance') FROM kiroku.subscriptions WHERE subscription_name = 'target-cost'" E.noParams (D.singleRow (D.column (D.nonNullable D.bool))))
  if expected && actual then putStrLn "durable equal-work checks: passed" else fail "durable equal-work checks failed"

use :: KirokuStore -> Session.Session a -> IO a
use store session = Pool.use (store ^. #pool) session >>= either (fail . show) pure

walPosition :: Statement () Int64
walPosition = preparable "SELECT (pg_current_wal_insert_lsn() - '0/0'::pg_lsn)::bigint" E.noParams (D.singleRow (D.column (D.nonNullable D.int8)))

controlSave :: Statement (Text, Int32, Int64, Int32) ()
controlSave =
  preparable
    """
    INSERT INTO ep2_control.subscriptions (subscription_name, consumer_group_member, last_seen, updated_at, consumer_group_size)
    VALUES ($1, $2, $3, now(), $4)
    ON CONFLICT (subscription_name, consumer_group_member)
    DO UPDATE SET last_seen = GREATEST(subscriptions.last_seen, EXCLUDED.last_seen), updated_at = now(), consumer_group_size = EXCLUDED.consumer_group_size
    """
    (contrazip4 (E.param (E.nonNullable E.text)) (E.param (E.nonNullable E.int4)) (E.param (E.nonNullable E.int8)) (E.param (E.nonNullable E.int4)))
    D.noResult