{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MultilineStrings #-}
module Main where
import Control.Concurrent.Async (mapConcurrently_)
import Control.Lens ((^.))
import Data.Aeson qualified as Aeson
import Data.Functor.Contravariant ((>$<))
import Data.Generics.Labels ()
import Data.IORef
import Data.Int (Int32, Int64)
import Data.Text (Text)
import Data.Text qualified as T
import Data.Time.Clock (UTCTime, diffUTCTime, getCurrentTime)
import Data.UUID (UUID)
import Data.UUID.V7 qualified as V7
import Data.Vector qualified as V
import EphemeralPg qualified as Pg
import GHC.Generics (Generic)
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 Hasql.Transaction qualified as Tx
import Hasql.Transaction.Sessions qualified as TxSessions
import Kiroku.Store
import Kiroku.Store.SQL qualified as SQL
import Kiroku.Test.Fixtures.CategoryScaling (categoryScalingFixtureSql, categoryScalingHead)
import Kiroku.Test.Postgres (ephemeralConfig, migrateTestDatabase, withMigratedTestDatabase, withSharedMigratedPostgres)
import Test.Tasty.Bench
data RawAppendParams = RawAppendParams
{ eventId :: !UUID,
eventType :: !Text,
causationId :: !(Maybe UUID),
correlationId :: !(Maybe UUID),
payload :: !Aeson.Value,
metadata :: !(Maybe Aeson.Value),
createdAt :: !UTCTime,
streamName :: !Text
}
deriving stock (Generic)
data RawProductionAppendParams = RawProductionAppendParams
{ productionEventIds :: !(V.Vector UUID),
productionEventTypes :: !(V.Vector Text),
productionCausationIds :: !(V.Vector (Maybe UUID)),
productionCorrelationIds :: !(V.Vector (Maybe UUID)),
productionPayloads :: !(V.Vector Aeson.Value),
productionMetadatas :: !(V.Vector (Maybe Aeson.Value)),
productionCreatedAts :: !(V.Vector UTCTime),
productionStreamName :: !Text
}
deriving stock (Generic)
type RawAppendResult = (Int64, Int64, Int64)
rawAppendParamsEncoder :: E.Params RawAppendParams
rawAppendParamsEncoder =
((^. #eventId) >$< E.param (E.nonNullable E.uuid))
<> ((^. #eventType) >$< E.param (E.nonNullable E.text))
<> ((^. #causationId) >$< E.param (E.nullable E.uuid))
<> ((^. #correlationId) >$< E.param (E.nullable E.uuid))
<> ((^. #payload) >$< E.param (E.nonNullable E.jsonb))
<> ((^. #metadata) >$< E.param (E.nullable E.jsonb))
<> ((^. #createdAt) >$< E.param (E.nonNullable E.timestamptz))
<> ((^. #streamName) >$< E.param (E.nonNullable E.text))
rawProductionAppendParamsEncoder :: E.Params RawProductionAppendParams
rawProductionAppendParamsEncoder =
((^. #productionEventIds) >$< E.param (E.nonNullable (E.foldableArray (E.nonNullable E.uuid))))
<> ((^. #productionEventTypes) >$< E.param (E.nonNullable (E.foldableArray (E.nonNullable E.text))))
<> ((^. #productionCausationIds) >$< E.param (E.nonNullable (E.foldableArray (E.nullable E.uuid))))
<> ((^. #productionCorrelationIds) >$< E.param (E.nonNullable (E.foldableArray (E.nullable E.uuid))))
<> ((^. #productionPayloads) >$< E.param (E.nonNullable (E.foldableArray (E.nonNullable E.jsonb))))
<> ((^. #productionMetadatas) >$< E.param (E.nonNullable (E.foldableArray (E.nullable E.jsonb))))
<> ((^. #productionCreatedAts) >$< E.param (E.nonNullable (E.foldableArray (E.nonNullable E.timestamptz))))
<> ((^. #productionStreamName) >$< E.param (E.nonNullable E.text))
rawAppendResultDecoder :: D.Result (Maybe RawAppendResult)
rawAppendResultDecoder =
D.rowMaybe $
(,,)
<$> D.column (D.nonNullable D.int8)
<*> D.column (D.nonNullable D.int8)
<*> D.column (D.nonNullable D.int8)
rawScalarAppendAnyVersion :: Statement RawAppendParams (Maybe RawAppendResult)
rawScalarAppendAnyVersion =
preparable
rawScalarAppendAnyVersionSQL
rawAppendParamsEncoder
rawAppendResultDecoder
rawProductionAppendAnyVersion :: Statement RawProductionAppendParams (Maybe RawAppendResult)
rawProductionAppendAnyVersion =
preparable
rawProductionAppendAnyVersionSQL
rawProductionAppendParamsEncoder
rawAppendResultDecoder
rawScalarAppendAnyVersionSQL :: Text
rawScalarAppendAnyVersionSQL =
"""
WITH
new_event AS (
SELECT $1::uuid AS event_id,
$2::text AS event_type,
$3::uuid AS causation_id,
$4::uuid AS correlation_id,
$5::jsonb AS data,
$6::jsonb AS metadata,
$7::timestamptz AS created_at
),
stream_upsert AS (
INSERT INTO streams (stream_name, stream_version)
VALUES ($8, 1)
ON CONFLICT (stream_name)
DO UPDATE SET stream_version = streams.stream_version + 1
WHERE streams.deleted_at IS NULL
RETURNING stream_id, category, stream_version - 1 AS initial_version
),
inserted_events AS (
INSERT INTO events (event_id, event_type, causation_id, correlation_id, data, metadata, created_at)
SELECT event_id, event_type, causation_id, correlation_id, data, metadata, created_at
FROM new_event
WHERE EXISTS (SELECT 1 FROM stream_upsert)
),
source_links AS (
INSERT INTO stream_events (event_id, stream_id, stream_version, original_stream_id, original_stream_version)
SELECT ne.event_id, su.stream_id, su.initial_version + 1, su.stream_id, su.initial_version + 1
FROM new_event ne
CROSS JOIN stream_upsert su
),
all_update AS (
UPDATE streams
SET stream_version = stream_version + 1
WHERE stream_id = 0
AND EXISTS (SELECT 1 FROM stream_upsert)
RETURNING stream_version - 1 AS initial_global_version
),
all_links AS (
INSERT INTO stream_events (event_id, stream_id, stream_version, original_stream_id, original_stream_version, category)
SELECT ne.event_id, 0, au.initial_global_version + 1, su.stream_id, su.initial_version + 1, su.category
FROM new_event ne
CROSS JOIN all_update au
CROSS JOIN stream_upsert su
)
SELECT su.stream_id,
su.initial_version + 1,
au.initial_global_version + 1
FROM stream_upsert su
CROSS JOIN all_update au
"""
rawProductionAppendAnyVersionSQL :: Text
rawProductionAppendAnyVersionSQL =
"""
WITH
new_events AS (
SELECT *
FROM unnest($1::uuid[], $2::text[], $3::uuid[], $4::uuid[], $5::jsonb[], $6::jsonb[], $7::timestamptz[])
WITH ORDINALITY AS t(event_id, event_type, causation_id, correlation_id, data, metadata, created_at, idx)
),
stream_upsert AS (
INSERT INTO streams (stream_name, stream_version)
VALUES ($8, (SELECT count(*) FROM new_events))
ON CONFLICT (stream_name)
DO UPDATE SET stream_version = streams.stream_version + (SELECT count(*) FROM new_events)
WHERE streams.deleted_at IS NULL
RETURNING stream_id, category, stream_version - (SELECT count(*) FROM new_events) AS initial_version
),
inserted_events AS (
INSERT INTO events (event_id, event_type, causation_id, correlation_id, data, metadata, created_at)
SELECT event_id, event_type, causation_id, correlation_id, data, metadata, created_at
FROM new_events
WHERE EXISTS (SELECT 1 FROM stream_upsert)
ORDER BY idx
),
source_links AS (
INSERT INTO stream_events (event_id, stream_id, stream_version, original_stream_id, original_stream_version)
SELECT ne.event_id, su.stream_id, su.initial_version + ne.idx, su.stream_id, su.initial_version + ne.idx
FROM new_events ne
CROSS JOIN stream_upsert su
),
all_update AS (
UPDATE streams
SET stream_version = stream_version + (SELECT count(*) FROM new_events)
WHERE stream_id = 0
AND EXISTS (SELECT 1 FROM stream_upsert)
RETURNING stream_version - (SELECT count(*) FROM new_events) AS initial_global_version
),
all_links AS (
INSERT INTO stream_events (event_id, stream_id, stream_version, original_stream_id, original_stream_version, category)
SELECT ne.event_id, 0, au.initial_global_version + ne.idx, su.stream_id, su.initial_version + ne.idx, su.category
FROM new_events ne
CROSS JOIN all_update au
CROSS JOIN stream_upsert su
)
SELECT su.stream_id,
su.initial_version + (SELECT count(*) FROM new_events),
au.initial_global_version + (SELECT count(*) FROM new_events)
FROM stream_upsert su
CROSS JOIN all_update au
"""
-- Two-round-trip variants: the proof-of-concept that the plan-23 restructure
-- is structurally faster than the production arrays/unnest CTE. The resolve
-- query reads (stream_id, stream_version, deleted_at) by stream_name; one of
-- two append CTEs then runs keyed on the resolved stream_id (existing) or on
-- the stream_name with no ON CONFLICT (new).
data RawResolution = RawResolution
{ resStreamId :: !Int64,
resStreamVersion :: !Int64,
resDeletedAt :: !(Maybe UTCTime)
}
deriving stock (Generic)
rawResolveStreamStmt :: Statement Text (Maybe RawResolution)
rawResolveStreamStmt =
preparable
rawResolveStreamSQL
(E.param (E.nonNullable E.text))
( D.rowMaybe $
RawResolution
<$> D.column (D.nonNullable D.int8)
<*> D.column (D.nonNullable D.int8)
<*> D.column (D.nullable D.timestamptz)
)
rawResolveStreamSQL :: Text
rawResolveStreamSQL =
"""
SELECT stream_id, stream_version, deleted_at
FROM streams
WHERE stream_name = $1
"""
-- | Append params for the two-round-trip "existing" path: stream_id is known.
data RawAppendExistingParams = RawAppendExistingParams
{ existingStreamId :: !Int64,
existingEventId :: !UUID,
existingEventType :: !Text,
existingCausationId :: !(Maybe UUID),
existingCorrelationId :: !(Maybe UUID),
existingPayload :: !Aeson.Value,
existingMetadata :: !(Maybe Aeson.Value),
existingCreatedAt :: !UTCTime
}
deriving stock (Generic)
rawAppendExistingEncoder :: E.Params RawAppendExistingParams
rawAppendExistingEncoder =
((^. #existingStreamId) >$< E.param (E.nonNullable E.int8))
<> ((^. #existingEventId) >$< E.param (E.nonNullable E.uuid))
<> ((^. #existingEventType) >$< E.param (E.nonNullable E.text))
<> ((^. #existingCausationId) >$< E.param (E.nullable E.uuid))
<> ((^. #existingCorrelationId) >$< E.param (E.nullable E.uuid))
<> ((^. #existingPayload) >$< E.param (E.nonNullable E.jsonb))
<> ((^. #existingMetadata) >$< E.param (E.nullable E.jsonb))
<> ((^. #existingCreatedAt) >$< E.param (E.nonNullable E.timestamptz))
rawAppendUpdateExisting :: Statement RawAppendExistingParams (Maybe RawAppendResult)
rawAppendUpdateExisting =
preparable
rawAppendUpdateExistingSQL
rawAppendExistingEncoder
rawAppendResultDecoder
-- | Update an existing stream keyed on integer stream_id. No version check,
-- no soft-delete check, no EXISTS gating, no count(*) — Haskell-side
-- validation has already done all of that.
rawAppendUpdateExistingSQL :: Text
rawAppendUpdateExistingSQL =
"""
WITH
stream_update AS (
UPDATE streams
SET stream_version = stream_version + 1
WHERE stream_id = $1::bigint
RETURNING stream_id, category, stream_version - 1 AS initial_version
),
inserted_event AS (
INSERT INTO events (event_id, event_type, causation_id, correlation_id, data, metadata, created_at)
VALUES ($2::uuid, $3::text, $4::uuid, $5::uuid, $6::jsonb, $7::jsonb, $8::timestamptz)
),
source_link AS (
INSERT INTO stream_events (event_id, stream_id, stream_version, original_stream_id, original_stream_version)
SELECT $2::uuid, su.stream_id, su.initial_version + 1, su.stream_id, su.initial_version + 1
FROM stream_update su
),
all_update AS (
UPDATE streams
SET stream_version = stream_version + 1
WHERE stream_id = 0
RETURNING stream_version - 1 AS initial_global_version
),
all_link AS (
INSERT INTO stream_events (event_id, stream_id, stream_version, original_stream_id, original_stream_version, category)
SELECT $2::uuid, 0, au.initial_global_version + 1, su.stream_id, su.initial_version + 1, su.category
FROM all_update au
CROSS JOIN stream_update su
)
SELECT su.stream_id,
su.initial_version + 1,
au.initial_global_version + 1
FROM stream_update su
CROSS JOIN all_update au
"""
-- | Append params for the two-round-trip "new stream" path: stream_name only.
data RawAppendNewParams = RawAppendNewParams
{ newStreamName :: !Text,
newEventId :: !UUID,
newEventType :: !Text,
newCausationId :: !(Maybe UUID),
newCorrelationId :: !(Maybe UUID),
newPayload :: !Aeson.Value,
newMetadata :: !(Maybe Aeson.Value),
newCreatedAt :: !UTCTime
}
deriving stock (Generic)
rawAppendNewEncoder :: E.Params RawAppendNewParams
rawAppendNewEncoder =
((^. #newStreamName) >$< E.param (E.nonNullable E.text))
<> ((^. #newEventId) >$< E.param (E.nonNullable E.uuid))
<> ((^. #newEventType) >$< E.param (E.nonNullable E.text))
<> ((^. #newCausationId) >$< E.param (E.nullable E.uuid))
<> ((^. #newCorrelationId) >$< E.param (E.nullable E.uuid))
<> ((^. #newPayload) >$< E.param (E.nonNullable E.jsonb))
<> ((^. #newMetadata) >$< E.param (E.nullable E.jsonb))
<> ((^. #newCreatedAt) >$< E.param (E.nonNullable E.timestamptz))
rawAppendCreateNew :: Statement RawAppendNewParams (Maybe RawAppendResult)
rawAppendCreateNew =
preparable
rawAppendCreateNewSQL
rawAppendNewEncoder
rawAppendResultDecoder
-- | Create a new stream and append one event. No ON CONFLICT — Haskell-side
-- validation already proved the stream is absent; a race-loser hits the
-- ix_streams_stream_name unique constraint and surfaces as a usage error.
rawAppendCreateNewSQL :: Text
rawAppendCreateNewSQL =
"""
WITH
stream_insert AS (
INSERT INTO streams (stream_name, stream_version)
VALUES ($1::text, 1)
RETURNING stream_id, category, 0::bigint AS initial_version
),
inserted_event AS (
INSERT INTO events (event_id, event_type, causation_id, correlation_id, data, metadata, created_at)
VALUES ($2::uuid, $3::text, $4::uuid, $5::uuid, $6::jsonb, $7::jsonb, $8::timestamptz)
),
source_link AS (
INSERT INTO stream_events (event_id, stream_id, stream_version, original_stream_id, original_stream_version)
SELECT $2::uuid, si.stream_id, si.initial_version + 1, si.stream_id, si.initial_version + 1
FROM stream_insert si
),
all_update AS (
UPDATE streams
SET stream_version = stream_version + 1
WHERE stream_id = 0
RETURNING stream_version - 1 AS initial_global_version
),
all_link AS (
INSERT INTO stream_events (event_id, stream_id, stream_version, original_stream_id, original_stream_version, category)
SELECT $2::uuid, 0, au.initial_global_version + 1, si.stream_id, si.initial_version + 1, si.category
FROM all_update au
CROSS JOIN stream_insert si
)
SELECT si.stream_id,
si.initial_version + 1,
au.initial_global_version + 1
FROM stream_insert si
CROSS JOIN all_update au
"""
mkRawAppendParams :: Text -> IO RawAppendParams
mkRawAppendParams name = do
eventUuid <- oneUuid
now <- getCurrentTime
pure
RawAppendParams
{ eventId = eventUuid,
eventType = "RawBenchEvent",
causationId = Nothing,
correlationId = Nothing,
payload = Aeson.object [("benchmark", Aeson.Bool True)],
metadata = Nothing,
createdAt = now,
streamName = name
}
mkRawProductionAppendParams :: RawAppendParams -> RawProductionAppendParams
mkRawProductionAppendParams params =
RawProductionAppendParams
{ productionEventIds = V.singleton (params ^. #eventId),
productionEventTypes = V.singleton (params ^. #eventType),
productionCausationIds = V.singleton (params ^. #causationId),
productionCorrelationIds = V.singleton (params ^. #correlationId),
productionPayloads = V.singleton (params ^. #payload),
productionMetadatas = V.singleton (params ^. #metadata),
productionCreatedAts = V.singleton (params ^. #createdAt),
productionStreamName = params ^. #streamName
}
oneUuid :: IO UUID
oneUuid = do
generated <- V7.genUUIDs 1
case generated of
uuid : _ -> pure uuid
[] -> error "oneUuid: UUID generator returned no IDs"
runRawScalarAppendAnyVersionNewStream :: KirokuStore -> IORef Int -> IO ()
runRawScalarAppendAnyVersionNewStream store counter = do
streamId <- atomicModifyIORef' counter (\n -> (n + 1, n))
params <- mkRawAppendParams ("raw-scalar-new-" <> T.pack (show streamId))
result <- Pool.use (store ^. #pool) $ Session.statement params rawScalarAppendAnyVersion
forceRawAppend result
runRawProductionAppendAnyVersionNewStream :: KirokuStore -> IORef Int -> IO ()
runRawProductionAppendAnyVersionNewStream store counter = do
streamId <- atomicModifyIORef' counter (\n -> (n + 1, n))
params <- mkRawProductionAppendParams <$> mkRawAppendParams ("raw-production-new-" <> T.pack (show streamId))
result <- Pool.use (store ^. #pool) $ Session.statement params rawProductionAppendAnyVersion
forceRawAppend result
runRawScalarAppendAnyVersionHotStream :: KirokuStore -> IO ()
runRawScalarAppendAnyVersionHotStream store = do
params <- mkRawAppendParams "raw-scalar-hot"
result <- Pool.use (store ^. #pool) $ Session.statement params rawScalarAppendAnyVersion
forceRawAppend result
runRawProductionAppendAnyVersionHotStream :: KirokuStore -> IO ()
runRawProductionAppendAnyVersionHotStream store = do
params <- mkRawProductionAppendParams <$> mkRawAppendParams "raw-production-hot"
result <- Pool.use (store ^. #pool) $ Session.statement params rawProductionAppendAnyVersion
forceRawAppend result
-- | Two-round-trip hot-stream variant: resolve stream by name, then UPDATE
-- the existing row keyed on integer stream_id. Both statements run on the
-- same pooled connection in a Session, each in its own implicit transaction
-- (no explicit BEGIN/COMMIT). This mirrors eventstore's small-batch path
-- (`EventStore.Streams.Stream.append_to_stream/5` at < 1000 events), which
-- runs two Postgrex.query calls without a transaction wrapper and relies on
-- the stream_events (stream_id, stream_version) unique constraint to detect
-- races. Wrapping in an explicit BEGIN/COMMIT adds two extra round-trips
-- and pessimises this path; that variant is measured separately below.
runRawTwoRoundtripAppendExistingHotStream :: KirokuStore -> IO ()
runRawTwoRoundtripAppendExistingHotStream store = do
base <- mkRawAppendParams "raw-two-roundtrip-hot"
result <- Pool.use (store ^. #pool) $ do
mRes <- Session.statement (base ^. #streamName) rawResolveStreamStmt
case mRes of
Nothing -> pure (Left "two-roundtrip hot bench: stream not pre-created")
Just res -> do
let params = mkRawAppendExistingParams (res ^. #resStreamId) base
fmap Right (Session.statement params rawAppendUpdateExisting)
forceTwoRoundtripResult result
-- | Two-round-trip new-stream variant. Same shape as the existing-stream
-- variant: two implicit-transaction round-trips on a pooled connection,
-- no explicit transaction wrapper.
runRawTwoRoundtripAppendNewStream :: KirokuStore -> IORef Int -> IO ()
runRawTwoRoundtripAppendNewStream store counter = do
streamId <- atomicModifyIORef' counter (\n -> (n + 1, n))
base <- mkRawAppendParams ("raw-two-roundtrip-new-" <> T.pack (show streamId))
result <- Pool.use (store ^. #pool) $ do
mRes <- Session.statement (base ^. #streamName) rawResolveStreamStmt
case mRes of
Just _ -> pure (Left "two-roundtrip new bench: stream already exists")
Nothing -> do
let params = mkRawAppendNewParams base
fmap Right (Session.statement params rawAppendCreateNew)
forceTwoRoundtripResult result
-- | Two-round-trip hot-stream variant wrapped in an explicit
-- read-committed transaction. Provided as a separate benchmark to quantify
-- the BEGIN/COMMIT round-trip overhead so we can decide whether the live
-- append path can safely drop the transaction wrapper.
runRawTwoRoundtripAppendExistingHotStreamTx :: KirokuStore -> IO ()
runRawTwoRoundtripAppendExistingHotStreamTx store = do
base <- mkRawAppendParams "raw-two-roundtrip-hot"
let txn = do
mRes <- Tx.statement (base ^. #streamName) rawResolveStreamStmt
case mRes of
Nothing -> pure (Left "two-roundtrip hot tx bench: stream not pre-created")
Just res -> do
let params = mkRawAppendExistingParams (res ^. #resStreamId) base
fmap Right (Tx.statement params rawAppendUpdateExisting)
result <-
Pool.use (store ^. #pool) $
TxSessions.transaction TxSessions.ReadCommitted TxSessions.Write txn
forceTwoRoundtripResult result
-- | Two-round-trip new-stream variant wrapped in an explicit transaction.
runRawTwoRoundtripAppendNewStreamTx :: KirokuStore -> IORef Int -> IO ()
runRawTwoRoundtripAppendNewStreamTx store counter = do
streamId <- atomicModifyIORef' counter (\n -> (n + 1, n))
base <- mkRawAppendParams ("raw-two-roundtrip-new-tx-" <> T.pack (show streamId))
let txn = do
mRes <- Tx.statement (base ^. #streamName) rawResolveStreamStmt
case mRes of
Just _ -> pure (Left "two-roundtrip new tx bench: stream already exists")
Nothing -> do
let params = mkRawAppendNewParams base
fmap Right (Tx.statement params rawAppendCreateNew)
result <-
Pool.use (store ^. #pool) $
TxSessions.transaction TxSessions.ReadCommitted TxSessions.Write txn
forceTwoRoundtripResult result
mkRawAppendExistingParams :: Int64 -> RawAppendParams -> RawAppendExistingParams
mkRawAppendExistingParams sid base =
RawAppendExistingParams
{ existingStreamId = sid,
existingEventId = base ^. #eventId,
existingEventType = base ^. #eventType,
existingCausationId = base ^. #causationId,
existingCorrelationId = base ^. #correlationId,
existingPayload = base ^. #payload,
existingMetadata = base ^. #metadata,
existingCreatedAt = base ^. #createdAt
}
mkRawAppendNewParams :: RawAppendParams -> RawAppendNewParams
mkRawAppendNewParams base =
RawAppendNewParams
{ newStreamName = base ^. #streamName,
newEventId = base ^. #eventId,
newEventType = base ^. #eventType,
newCausationId = base ^. #causationId,
newCorrelationId = base ^. #correlationId,
newPayload = base ^. #payload,
newMetadata = base ^. #metadata,
newCreatedAt = base ^. #createdAt
}
forceTwoRoundtripResult ::
Either Pool.UsageError (Either String (Maybe RawAppendResult)) -> IO ()
forceTwoRoundtripResult (Left e) =
error ("Two-round-trip bench pool error: " <> show e)
forceTwoRoundtripResult (Right (Left msg)) =
error ("Two-round-trip bench precondition failed: " <> msg)
forceTwoRoundtripResult (Right (Right inner)) =
forceRawAppend (Right inner)
-- | Run @writers@ concurrent appenders, each performing @ops@ appends to
-- its own unique stream. Used by the structured concurrent-writer
-- benchmarks (EP-6 F19); replaces the wall-clock @mapConcurrently_@
-- measurement that previously lived inline in @main@.
--
-- The @runCounter@ ref is bumped once per call so stream names are unique
-- across the many iterations tasty-bench runs.
runConcurrentWriters :: KirokuStore -> IORef Int -> Int -> Int -> IO ()
runConcurrentWriters store runCounter writers ops = do
runId <- atomicModifyIORef' runCounter (\m -> (m + 1, m))
mapConcurrently_
(\tid -> mapM_ (appendOne tid runId) [1 .. ops])
[1 .. writers]
where
appendOne :: Int -> Int -> Int -> IO ()
appendOne tid runId i = do
let sn = StreamName ("conc-" <> T.pack (show runId) <> "-" <> T.pack (show tid) <> "-" <> T.pack (show i))
r <- runStoreIO store $ appendToStream sn AnyVersion [makeEvent "ConcEvent"]
forceAppend r
-- | Exercise the hot stream named in the focused reliability audit.
runHotInvoicePayment :: KirokuStore -> Int -> IO ()
runHotInvoicePayment store ops =
mapM_
( \i -> do
r <- runStoreIO store $ appendToStream (StreamName "invoice-payment") AnyVersion [makeEvent ("InvoicePayment" <> T.pack (show i))]
forceAppend r
)
[1 .. ops]
-- | Exercise appendMultiStream against existing streams.
runAppendMultiStream :: KirokuStore -> IO ()
runAppendMultiStream store = do
r <-
runStoreIO store $
appendMultiStream
[ (StreamName "bench-multi-a", AnyVersion, [makeEvent "MultiA"]),
(StreamName "bench-multi-b", AnyVersion, [makeEvent "MultiB"]),
(StreamName "bench-multi-c", AnyVersion, [makeEvent "MultiC"])
]
forceAppendList r
-- | Exercise subscription catch-up over a compact category-local backlog.
runSubscriptionCatchup :: KirokuStore -> IORef Int -> IO ()
runSubscriptionCatchup store runCounter = do
runId <- atomicModifyIORef' runCounter (\m -> (m + 1, m))
let cat = "benchsub" <> T.pack (show runId)
sn = StreamName (cat <> "-stream")
subName = SubscriptionName ("bench-sub-" <> T.pack (show runId))
events = map (\i -> makeEvent ("SubCatchup" <> T.pack (show i))) [1 .. 100 :: Int]
r <- runStoreIO store $ appendToStream sn NoStream events
forceAppend r
seenRef <- newIORef (0 :: Int)
let handler _ = do
n <- atomicModifyIORef' seenRef (\m -> let m' = m + 1 in (m', m'))
pure $ if n >= 100 then Stop else Continue
cfg =
SubscriptionConfig
{ name = subName,
target = Category (CategoryName cat),
handler = handler,
batchSize = defaultBatchSize,
queueCapacity = 16,
overflowPolicy = DropSubscription,
consumerGroup = Nothing,
consumerGroupGuard = False,
missingCheckpointPolicy = FromBeginning,
targetBindingPolicy = AdoptUnbound,
undecodableHandler = Nothing,
handlerStallWarnAfter = Nothing,
retryPolicy = defaultRetryPolicy,
eventTypeFilter = AllEventTypes,
selector = Nothing
}
handle <- subscribe store cfg
result <- wait handle
case result of
Right () -> pure ()
Left e -> error ("Subscription catch-up benchmark failed: " <> show e)
-- | Force evaluation of an append result or fail the benchmark.
forceAppend :: Either StoreError AppendResult -> IO ()
forceAppend (Right r) = (r ^. #streamVersion) `seq` (r ^. #globalPosition) `seq` pure ()
forceAppend (Left e) = error ("Benchmark append failed: " <> show e)
-- | Force evaluation of multi-stream append results or fail the benchmark.
forceAppendList :: Either StoreError [AppendResult] -> IO ()
forceAppendList (Right rs) = mapM_ (\r -> (r ^. #streamVersion) `seq` (r ^. #globalPosition) `seq` pure ()) rs
forceAppendList (Left e) = error ("Benchmark appendMultiStream failed: " <> show e)
-- | Force evaluation of a raw append shape result or fail the benchmark.
forceRawAppend :: Either Pool.UsageError (Maybe RawAppendResult) -> IO ()
forceRawAppend (Right (Just (!streamId, !streamVersion, !globalPosition))) =
streamId `seq` streamVersion `seq` globalPosition `seq` pure ()
forceRawAppend (Right Nothing) = error "Raw append benchmark produced no result"
forceRawAppend (Left e) = error ("Raw append benchmark failed: " <> show e)
-- | Force evaluation of a read result or fail the benchmark.
forceRead :: Either StoreError (V.Vector RecordedEvent) -> IO ()
forceRead (Right v) = V.length v `seq` pure ()
forceRead (Left e) = error ("Benchmark read failed: " <> show e)
seedCheckpointInventoryStmt :: Statement Int32 ()
seedCheckpointInventoryStmt =
preparable
"""
INSERT INTO subscriptions (
subscription_name,
consumer_group_member,
last_seen,
updated_at
)
SELECT 'inventory-' || lpad(n::text, 8, '0'),
0,
0,
now()
FROM generate_series(1, $1::int4) AS n
"""
(E.param (E.nonNullable E.int4))
D.noResult
seedCheckpointInventory :: KirokuStore -> Int32 -> IO ()
seedCheckpointInventory store rowCount = do
result <- Pool.use (store ^. #pool) $ Session.statement rowCount seedCheckpointInventoryStmt
case result of
Left err -> error ("Checkpoint inventory benchmark setup failed: " <> show err)
Right () -> pure ()
forceCheckpointInventory :: Either StoreError SubscriptionCheckpointInventory -> IO ()
forceCheckpointInventory (Left err) = error ("Checkpoint inventory benchmark failed: " <> show err)
forceCheckpointInventory (Right (SubscriptionCheckpointInventory (GlobalPosition storePos) rows)) = do
let !checksum = V.foldl' forceCheckpoint (fromIntegral storePos) rows
checksum `seq` pure ()
where
forceCheckpoint !acc (SubscriptionCheckpoint (SubscriptionName name) member (GlobalPosition position) updatedAt) =
updatedAt `seq`
acc
+ T.length name
+ fromIntegral member
+ fromIntegral position
runCheckpointInventoryBenchmark :: KirokuStore -> IO ()
runCheckpointInventoryBenchmark store =
runStoreIO store subscriptionCheckpointInventory >>= forceCheckpointInventory
withInventoryBenchmarkStores :: (KirokuStore -> KirokuStore -> IO a) -> IO a
withInventoryBenchmarkStores action =
withSharedMigratedPostgres $
withMigratedTestDatabase $ \inventory100Connection ->
withMigratedTestDatabase $ \inventory10000Connection ->
withStore (defaultConnectionSettings inventory100Connection) $ \inventory100Store ->
withStore (defaultConnectionSettings inventory10000Connection) $ \inventory10000Store -> do
seedCheckpointInventory inventory100Store 100
seedCheckpointInventory inventory10000Store 10_000
action inventory100Store inventory10000Store
-- | A store seeded with 'categoryScalingFixtureSql': categories @performance@
-- (200 streams), @idle@ (20,000 single-event streams), and @noise@, head at
-- 'categoryScalingHead'. Must run inside 'withSharedMigratedPostgres'.
withCategoryScalingStore :: (KirokuStore -> IO a) -> IO a
withCategoryScalingStore action =
withMigratedTestDatabase $ \connection ->
withStore (defaultConnectionSettings connection) $ \store -> do
seeded <- Pool.use (store ^. #pool) (Session.script categoryScalingFixtureSql)
case seeded of
Left err -> error ("Category-scaling benchmark setup failed: " <> show err)
Right () -> action store
-- | Run the plain category read 10 times at one cursor and limit.
runCategoryPolls :: KirokuStore -> Text -> Int64 -> Int32 -> IO ()
runCategoryPolls store category cursor limit =
mapM_
( \_ ->
Pool.use (store ^. #pool) (Session.statement (cursor, category, limit) SQL.readCategoryForwardStmt)
>>= forceCategoryPoll
)
[1 .. 10 :: Int]
-- | Run the consumer-group category read 10 times for one member.
runGroupCategoryPolls :: KirokuStore -> Text -> Int64 -> Int32 -> Int32 -> Int32 -> IO ()
runGroupCategoryPolls store category cursor member size limit =
mapM_
( \_ ->
Pool.use (store ^. #pool) (Session.statement (cursor, category, member, size, limit) SQL.readCategoryForwardConsumerGroupStmt)
>>= forceCategoryPoll
)
[1 .. 10 :: Int]
forceCategoryPoll :: Either Pool.UsageError (V.Vector RecordedEvent) -> IO ()
forceCategoryPoll (Right events) = V.length events `seq` pure ()
forceCategoryPoll (Left err) = error ("Category-scaling benchmark read failed: " <> show err)
main :: IO ()
main = do
-- Start ephemeral PostgreSQL once for all benchmarks
config <- ephemeralConfig
result <- Pg.withCachedConfig config Pg.defaultCacheConfig $ \db -> do
migrateTestDatabase (Pg.connectionString db)
let settings = defaultConnectionSettings (Pg.connectionString db)
withStore settings $ \store -> do
-- Counter for unique stream names across benchmarks
counter <- newIORef (0 :: Int)
let nextStream :: Text -> IO StreamName
nextStream prefix = do
n <- atomicModifyIORef' counter (\n -> (n + 1, n))
pure (StreamName (prefix <> "-" <> T.pack (show n)))
-- Pre-populate streams for category benchmarks (B10)
-- 100 categories × 10 streams × 100 events = 100K events
putStrLn "\n--- Pre-populating category data (100 cats × 10 streams × 100 events) ---"
catT0 <- getCurrentTime
mapM_
( \cat -> do
mapM_
( \s -> do
let sn = StreamName ("cat" <> T.pack (show cat) <> "-" <> T.pack (show s))
let evts = map (\i -> makeEvent ("E" <> T.pack (show i))) [1 .. 100 :: Int]
r' <- runStoreIO store $ appendToStream sn NoStream evts
forceAppend r'
)
[1 .. 10 :: Int]
)
[1 .. 100 :: Int]
catT1 <- getCurrentTime
let catElapsed = realToFrac (diffUTCTime catT1 catT0) :: Double
putStrLn $ " Setup time: " <> show catElapsed <> "s (100K events)"
-- Pre-populate streams for read benchmarks
-- B4: Single stream with 1000 events
let readStreamName = StreamName "bench-read-stream"
let readEvents = map (\i -> makeEvent ("E" <> T.pack (show i))) [1 .. 1000 :: Int]
r <- runStoreIO store $ appendToStream readStreamName NoStream readEvents
forceAppend r
-- B5: 10 streams with 100 events each for $all reads (1000 total)
mapM_
( \s -> do
let sn = StreamName ("bench-all-" <> T.pack (show s))
let evts = map (\i -> makeEvent ("E" <> T.pack (show i))) [1 .. 100 :: Int]
r' <- runStoreIO store $ appendToStream sn NoStream evts
forceAppend r'
)
[1 .. 10 :: Int]
-- Pre-create fixed streams for appendMultiStream benchmark iterations.
mapM_
( \sn -> do
r' <- runStoreIO store $ appendToStream sn NoStream [makeEvent "Init"]
forceAppend r'
)
[StreamName "bench-multi-a", StreamName "bench-multi-b", StreamName "bench-multi-c"]
-- Pre-create the hot stream targeted by the two-round-trip raw
-- shape variant. rawAppendUpdateExisting requires the row to
-- exist (no ON CONFLICT), so without this warmup the first
-- iteration would fail. The scalar-singleton and production
-- variants already create their hot stream on first iteration
-- via ON CONFLICT — pre-creating them here would change the
-- existing benchmark contract, so we only seed the two-round-trip
-- target.
do
r' <- runStoreIO store $ appendToStream (StreamName "raw-two-roundtrip-hot") AnyVersion [makeEvent "Init"]
forceAppend r'
-- B9: Pool saturation benchmark (64 concurrent writers, 100 appends each)
putStrLn "\n--- B9: Pool saturation (64 writers × 100 appends, pool size 10) ---"
satCounter <- newIORef (0 :: Int)
let nextSatStream :: Int -> Int -> StreamName
nextSatStream tid i = StreamName ("sat-" <> T.pack (show tid) <> "-" <> T.pack (show i))
t0 <- getCurrentTime
mapConcurrently_
( \tid -> do
mapM_
( \i -> do
let sn = nextSatStream tid i
r' <- runStoreIO store $ appendToStream sn AnyVersion [makeEvent "SatEvent"]
forceAppend r'
atomicModifyIORef' satCounter (\n -> (n + 1, ()))
)
[1 .. 100 :: Int]
)
[1 .. 64 :: Int]
t1 <- getCurrentTime
totalOps <- readIORef satCounter
let elapsed = realToFrac (diffUTCTime t1 t0) :: Double
let throughput = fromIntegral totalOps / elapsed
let avgLatency = elapsed / fromIntegral totalOps * 1000 -- ms
putStrLn $ " Total appends: " <> show totalOps
putStrLn $ " Elapsed: " <> show elapsed <> "s"
putStrLn $ " Throughput: " <> show (round throughput :: Int) <> " ops/s"
putStrLn $ " Avg latency: " <> show avgLatency <> " ms"
putStrLn "---"
-- Counter shared by the concurrent-writer benchmarks so each
-- iteration uses a fresh stream-name run-id.
concCounter <- newIORef (0 :: Int)
rawCounter <- newIORef (0 :: Int)
subCounter <- newIORef (0 :: Int)
withInventoryBenchmarkStores $ \inventory100Store inventory10000Store ->
withCategoryScalingStore $ \scalingStore ->
defaultMain
[ bgroup
"append"
[ bgroup
"single-event"
[ bench "NoStream (new stream)" $ whnfIO $ do
sn <- nextStream "bench-single"
r' <- runStoreIO store $ appendToStream sn NoStream [makeEvent "BenchEvent"]
forceAppend r',
bench "AnyVersion (new stream)" $ whnfIO $ do
sn <- nextStream "bench-any"
r' <- runStoreIO store $ appendToStream sn AnyVersion [makeEvent "BenchEvent"]
forceAppend r'
],
bgroup
"batch-10"
[ bench "NoStream" $ whnfIO $ do
sn <- nextStream "bench-b10"
let events = map (\i -> makeEvent ("E" <> T.pack (show i))) [1 .. 10 :: Int]
r' <- runStoreIO store $ appendToStream sn NoStream events
forceAppend r'
],
bgroup
"batch-100"
[ bench "NoStream" $ whnfIO $ do
sn <- nextStream "bench-b100"
let events = map (\i -> makeEvent ("E" <> T.pack (show i))) [1 .. 100 :: Int]
r' <- runStoreIO store $ appendToStream sn NoStream events
forceAppend r'
],
bgroup
"sequential"
[ bench "10 appends to same stream" $ whnfIO $ do
sn <- nextStream "bench-seq"
r0 <- runStoreIO store $ appendToStream sn NoStream [makeEvent "Init"]
forceAppend r0
let res0 = case r0 of
Right ok -> ok
Left e -> error ("Sequential append failed: " <> show e)
let go _ 0 = pure ()
go v n = do
r' <- runStoreIO store $ appendToStream sn (ExactVersion v) [makeEvent "Seq"]
case r' of
Right res -> go (res ^. #streamVersion) (n - 1 :: Int)
Left e -> error ("Sequential append failed: " <> show e)
go (res0 ^. #streamVersion) 9
]
],
bgroup
"raw-append-shape"
[ bgroup
"AnyVersion"
[ bench "scalar singleton (new stream)" $
whnfIO $
runRawScalarAppendAnyVersionNewStream store rawCounter,
bench "production arrays/unnest (new stream)" $
whnfIO $
runRawProductionAppendAnyVersionNewStream store rawCounter,
bench "two-roundtrip (new stream)" $
whnfIO $
runRawTwoRoundtripAppendNewStream store rawCounter,
bench "two-roundtrip + BEGIN/COMMIT (new stream)" $
whnfIO $
runRawTwoRoundtripAppendNewStreamTx store rawCounter,
bench "scalar singleton (hot stream)" $
whnfIO $
runRawScalarAppendAnyVersionHotStream store,
bench "production arrays/unnest (hot stream)" $
whnfIO $
runRawProductionAppendAnyVersionHotStream store,
bench "two-roundtrip (hot stream)" $
whnfIO $
runRawTwoRoundtripAppendExistingHotStream store,
bench "two-roundtrip + BEGIN/COMMIT (hot stream)" $
whnfIO $
runRawTwoRoundtripAppendExistingHotStreamTx store
]
],
bgroup
"read"
[ bench "stream forward (100-event page)" $ whnfIO $ do
r' <- runStoreIO store $ readStreamForward readStreamName (StreamVersion 0) 100
forceRead r',
bench "$all forward (100-event page)" $ whnfIO $ do
r' <- runStoreIO store $ readAllForward (GlobalPosition 0) 100
forceRead r'
],
bgroup
"category"
[ bench "category forward (100-event page)" $ whnfIO $ do
-- Read from cat1 category (has 10 streams × 100 events = 1000 events)
r' <- runStoreIO store $ readCategory (CategoryName "cat1") (GlobalPosition 0) 100
forceRead r',
bench "exhausted-category" $ whnfIO $ do
-- cat1 events are inserted early in setup; a high cursor proves
-- category reads do not scan the rest of $all looking for matches.
r' <- runStoreIO store $ readCategory (CategoryName "cat1") (GlobalPosition 90_000) 100
forceRead r',
bench "$all forward (100-event page baseline)" $ whnfIO $ do
r' <- runStoreIO store $ readAllForward (GlobalPosition 0) 100
forceRead r'
],
-- F19 — Concurrent-writer stress as structured benchmarks.
-- The legacy ad-hoc B9 measurement (still present above
-- for historical comparability) prints throughput and
-- latency once; these bgroup entries surface the same
-- workload through tasty-bench so it participates in the
-- baseline-regression workflow (Justfile bench-regression).
bgroup
"concurrent"
[ bench "8 writers x 10 appends" $ whnfIO $ runConcurrentWriters store concCounter 8 10,
bench "32 writers x 10 appends" $ whnfIO $ runConcurrentWriters store concCounter 32 10
],
bgroup
"reliability-audit"
[ bench "hot invoice-payment 10 AnyVersion appends" $ whnfIO $ runHotInvoicePayment store 10,
bench "appendMultiStream 3 existing streams" $ whnfIO $ runAppendMultiStream store,
bench "subscription category catch-up 100 events" $ whnfIO $ runSubscriptionCatchup store subCounter
],
-- BUG-2: a caught-up category poll must not cost work
-- proportional to the number of streams in the category.
-- Each cell runs the statement 10 times.
bgroup
"category-scaling"
[ bench "plain caught-up poll (200 streams)" $
whnfIO $
runCategoryPolls scalingStore "performance" categoryScalingHead 100,
bench "plain caught-up poll (20000 streams)" $
whnfIO $
runCategoryPolls scalingStore "idle" categoryScalingHead 100,
bench "group caught-up poll (20000 streams)" $
whnfIO $
runGroupCategoryPolls scalingStore "idle" categoryScalingHead 1 2 100,
bench "plain page from 0 (20000 streams)" $
whnfIO $
runCategoryPolls scalingStore "idle" 0 100,
bench "exhausted category (200 streams)" $
-- performance ends at 20,000; 60,000 other rows follow.
whnfIO $
runCategoryPolls scalingStore "performance" 20_000 100
],
bgroup
"subscription-checkpoint-inventory"
[ bench "100 rows" $ whnfIO $ runCheckpointInventoryBenchmark inventory100Store,
bench "10000 rows" $ whnfIO $ runCheckpointInventoryBenchmark inventory10000Store
]
]
case result of
Left err -> error ("Failed to start ephemeral PostgreSQL: " <> show err)
Right () -> pure ()
makeEvent :: Text -> EventData
makeEvent typ =
EventData
{ eventId = Nothing,
eventType = EventType typ,
payload = Aeson.object [("benchmark", Aeson.Bool True)],
metadata = Nothing,
causationId = Nothing,
correlationId = Nothing
}