{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MultilineStrings #-}
module Main where
import Control.Concurrent (threadDelay)
import Control.Concurrent.Async (mapConcurrently_)
import Control.Concurrent.Async qualified as Async
import Control.Exception (finally)
import Control.Lens ((^.))
import Control.Monad (replicateM_)
import Data.Aeson qualified as Aeson
import Data.Functor.Contravariant ((>$<))
import Data.Generics.Labels ()
import Data.IORef
import Data.Int (Int64)
import Data.Maybe (isNothing)
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.Pipeline qualified as Pipeline
import Hasql.Pool qualified as Pool
import Hasql.Session qualified as Session
import Hasql.Statement (Statement, preparable, unpreparable)
import Hasql.Transaction qualified as Tx
import Hasql.Transaction.Sessions qualified as TxSessions
import Kiroku.Store
import Kiroku.Store.Effect (buildAppendParams, prepareEvents)
import Kiroku.Store.SQL qualified as SQL
import Kiroku.Test.Postgres (migrateTestDatabase)
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
beginStmt :: Statement () ()
beginStmt = unpreparable "BEGIN" E.noParams D.noResult
commitStmt :: Statement () ()
commitStmt = unpreparable "COMMIT" E.noParams D.noResult
rollbackStmt :: Statement () ()
rollbackStmt = unpreparable "ROLLBACK" E.noParams D.noResult
pipelinedMultiAppend4Streams :: [(StreamName, Text)]
pipelinedMultiAppend4Streams =
namedBenchStreams "pipe4" 4
pipelinedMultiAppend8Streams :: [(StreamName, Text)]
pipelinedMultiAppend8Streams =
namedBenchStreams "pipe8" 8
pipelinedContentionCurrentGroups :: [[(StreamName, Text)]]
pipelinedContentionCurrentGroups =
namedBenchGroups "pipe-current-writer" 4 4
pipelinedContentionPipelinedGroups :: [[(StreamName, Text)]]
pipelinedContentionPipelinedGroups =
namedBenchGroups "pipe-pipelined-writer" 4 4
namedBenchGroups :: Text -> Int -> Int -> [[(StreamName, Text)]]
namedBenchGroups prefix groups streamsPerGroup =
[ namedBenchStreams (prefix <> "-" <> T.pack (show groupId)) streamsPerGroup
| groupId <- [1 .. groups]
]
namedBenchStreams :: Text -> Int -> [(StreamName, Text)]
namedBenchStreams prefix count =
[ ( StreamName (prefix <> "-" <> T.pack (show i))
, prefix <> "-event-" <> T.pack (show i)
)
| i <- [1 .. count]
]
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, 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)
SELECT ne.event_id, 0, au.initial_global_version + 1, su.stream_id, su.initial_version + 1
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, 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)
SELECT ne.event_id, 0, au.initial_global_version + ne.idx, su.stream_id, su.initial_version + ne.idx
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, 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)
SELECT $2::uuid, 0, au.initial_global_version + 1, su.stream_id, su.initial_version + 1
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, 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)
SELECT $2::uuid, 0, au.initial_global_version + 1, si.stream_id, si.initial_version + 1
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
runCurrentMultiAppend :: KirokuStore -> [(StreamName, Text)] -> IO ()
runCurrentMultiAppend store streams = do
r <-
runStoreIO store $
appendMultiStream
[ (sn, AnyVersion, [makeEvent typ])
| (sn, typ) <- streams
]
forceAppendList r
runPipelinedMultiAppend :: KirokuStore -> [(StreamName, Text)] -> IO ()
runPipelinedMultiAppend store streams = do
now <- getCurrentTime
params <-
mapM
( \(StreamName name, typ) -> do
prepared <- prepareEvents [makeEvent typ]
pure (buildAppendParams name now prepared)
)
streams
let names = V.fromList [name | (StreamName name, _) <- streams]
result <-
Pool.use (store ^. #pool) $ do
results <-
Session.pipeline $
Pipeline.statement () beginStmt
*> Pipeline.statement names SQL.lockStreamsForMultiStmt
*> traverse (`Pipeline.statement` SQL.appendAnyVersion) params
Session.statement () $
if any isNothing results
then rollbackStmt
else commitStmt
pure results
forcePipelinedAppendList result
runSingleAppendUnderMultiWriters ::
KirokuStore ->
IORef Int ->
(KirokuStore -> [(StreamName, Text)] -> IO ()) ->
Text ->
[[(StreamName, Text)]] ->
IO ()
runSingleAppendUnderMultiWriters store runCounter multiAppendRunner prefix groups = do
runId <- atomicModifyIORef' runCounter (\m -> (m + 1, m))
workers <- mapM (Async.async . replicateM_ 20 . multiAppendRunner store) groups
let cleanup = mapM_ Async.cancel workers
appendMeasured i = do
let sn = StreamName (prefix <> "-single-" <> T.pack (show runId) <> "-" <> T.pack (show i))
r <- runStoreIO store $ appendToStream sn AnyVersion [makeEvent (prefix <> "-single")]
forceAppend r
(threadDelay 20_000 >> mapM_ appendMeasured [1 .. 10 :: Int] >> mapM_ Async.wait workers) `finally` cleanup
-- | 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 = 100
, queueCapacity = 16
, overflowPolicy = DropSubscription
, consumerGroup = Nothing
, consumerGroupGuard = False
, 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)
forcePipelinedAppendList :: Either Pool.UsageError [Maybe AppendResult] -> IO ()
forcePipelinedAppendList (Right rs) =
mapM_
( \mResult -> case mResult of
Just r -> (r ^. #streamVersion) `seq` (r ^. #globalPosition) `seq` pure ()
Nothing -> error "Pipelined appendMultiStream benchmark produced no result"
)
rs
forcePipelinedAppendList (Left e) = error ("Pipelined appendMultiStream benchmark 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)
main :: IO ()
main = do
-- Start ephemeral PostgreSQL once for all benchmarks
result <- Pg.withCached $ \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"]
<> map fst pipelinedMultiAppend4Streams
<> map fst pipelinedMultiAppend8Streams
<> map fst (concat pipelinedContentionCurrentGroups)
<> map fst (concat pipelinedContentionPipelinedGroups)
)
-- 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)
pipelinedContentionCounter <- newIORef (0 :: Int)
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 Right res0 = r0
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
"pipelined-multi-append"
[ bench "current shape (4 streams)" $
whnfIO $
runCurrentMultiAppend store pipelinedMultiAppend4Streams
, bench "pipelined (4 streams)" $
whnfIO $
runPipelinedMultiAppend store pipelinedMultiAppend4Streams
, bench "current shape (8 streams)" $
whnfIO $
runCurrentMultiAppend store pipelinedMultiAppend8Streams
, bench "pipelined (8 streams)" $
whnfIO $
runPipelinedMultiAppend store pipelinedMultiAppend8Streams
, bench "single-stream under 4 current multi-stream writers" $
whnfIO $
runSingleAppendUnderMultiWriters
store
pipelinedContentionCounter
runCurrentMultiAppend
"pipe-current"
pipelinedContentionCurrentGroups
, bench "single-stream under 4 pipelined multi-stream writers" $
whnfIO $
runSingleAppendUnderMultiWriters
store
pipelinedContentionCounter
runPipelinedMultiAppend
"pipe-pipelined"
pipelinedContentionPipelinedGroups
]
, 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
]
]
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
}