kiroku-store-0.12.0.0: bench/ShibuyaOverhead.hs

-- | Benchmark measuring the overhead of the Shibuya adapter layer
-- compared to bare kiroku subscriptions. Three layers are measured:
--
-- 1. Bare subscribe: handler receives events directly from the subscription worker
-- 2. subscriptionStream: events pass through a TBQueue bridge to a Streamly stream
-- 3. Shibuya adapter: events flow through subscriptionStream → morphInner → Shibuya pipeline
--
-- The bare subscribe handler processes events synchronously (one at a time),
-- while the stream-based approaches decouple the producer from the consumer
-- via a TBQueue, allowing the subscription worker to batch-fetch ahead.
-- This means subscriptionStream and Shibuya can be *faster* for catch-up —
-- the relevant overhead comparison is shibuya vs subscriptionStream.
module Main where

import Control.Concurrent (threadDelay)
import Control.Concurrent.STM (atomically, newTVarIO, readTVar, readTVarIO, registerDelay, writeTVar)
import Control.Concurrent.STM qualified as STM
import Control.Lens ((^.))
import Data.Aeson qualified as Aeson
import Data.Generics.Labels ()
import Data.HashMap.Strict qualified as HashMap
import Data.IORef (atomicModifyIORef', newIORef)
import Data.List (sort)
import Data.Text (Text)
import Data.Text qualified as T
import Data.Time.Clock (diffUTCTime, getCurrentTime)
import Effectful (Eff, IOE, liftIO, runEff, (:>))
import EphemeralPg qualified as Pg
import Kiroku.Store
import Kiroku.Store.Subscription.Stream (subscriptionStream)
import Kiroku.Store.Subscription.Stream qualified as Buffer
import Kiroku.Test.Postgres (ephemeralConfig, migrateTestDatabase)
import Shibuya.Adapter (Adapter (..))
import Shibuya.App (ProcessorId (..), defaultAppConfig, mkProcessor, runApp, stopApp)
import Shibuya.Core.Ack (AckDecision (..))
import Shibuya.Core.AckHandle (AckHandle (..))
import Shibuya.Core.Ingested (DeliveryStatus (..), Ingested (..), Message)
import Shibuya.Core.Types (Cursor (..), Envelope (..), MessageId (..))
import Shibuya.Telemetry.Effect (runTracingNoop)
import Streamly.Data.Fold qualified as Fold
import Streamly.Data.Stream qualified as Stream
import System.IO (hFlush, stdout)
import Text.Printf (printf)

iterations :: Int
iterations = 5

main :: IO ()
main = do
  config <- ephemeralConfig
  result <- Pg.withCachedConfig config Pg.defaultCacheConfig $ \db -> do
    migrateTestDatabase (Pg.connectionString db)
    let settings = defaultConnectionSettings (Pg.connectionString db)
    withStore settings $ \store -> do
      putStrLn "=== Kiroku Shibuya Adapter Overhead Benchmark ==="
      printf "    %d iterations per measurement\n\n" iterations

      mapM_ (runBenchmarkSuite store) [100, 1000, 5000]

  case result of
    Left err -> error ("Failed to start ephemeral PostgreSQL: " <> show err)
    Right () -> pure ()

runBenchmarkSuite :: KirokuStore -> Int -> IO ()
runBenchmarkSuite store n = do
  printf "--- %d events ---\n\n" n

  -- Pre-populate events
  counter <- newIORef (0 :: Int)
  let nextId = atomicModifyIORef' counter (\i -> (i + 1, i))

  let streamPrefix = "bench-" <> T.pack (show n) <> "-"
  let events = map (\i -> makeEvent ("E" <> T.pack (show i))) [1 .. n]
  let chunked = chunksOf 10 events
  mapM_
    ( \(idx, chunk) -> do
        let sn = StreamName (streamPrefix <> T.pack (show idx))
        Right _ <- runStoreIO store $ appendToStream sn AnyVersion chunk
        pure ()
    )
    (zip [(0 :: Int) ..] chunked)
  threadDelay 500_000

  -- Run each benchmark multiple times
  bareTimes <- sequence [benchBareSubscribe store n nextId | _ <- [1 .. iterations]]
  streamTimes <- sequence [benchSubscriptionStream store n nextId | _ <- [1 .. iterations]]
  shibuyaTimes <- sequence [benchShibuyaAdapter store n nextId | _ <- [1 .. iterations]]

  let bareMedian = median bareTimes
      streamMedian = median streamTimes
      shibuyaMedian = median shibuyaTimes

  putRow "  bare subscribe" n bareTimes
  putRow "  subscriptionStream" n streamTimes
  putRow "  shibuya adapter" n shibuyaTimes

  putStrLn ""
  let streamVsBare = (streamMedian - bareMedian) / bareMedian * 100
      shibuyaVsBare = (shibuyaMedian - bareMedian) / bareMedian * 100
      shibuyaVsStream = (shibuyaMedian - streamMedian) / streamMedian * 100
  printf "  stream vs bare:    %s\n" (showPct streamVsBare)
  printf "  shibuya vs bare:   %s\n" (showPct shibuyaVsBare)
  printf "  shibuya vs stream: %s  (adapter + framework overhead)\n" (showPct shibuyaVsStream)
  putStrLn ""

-- | Benchmark: bare subscribe with IO handler
benchBareSubscribe :: KirokuStore -> Int -> IO Int -> IO Double
benchBareSubscribe store n nextId = do
  subId <- nextId
  let subName = SubscriptionName ("bare-bench-" <> T.pack (show subId))
  countVar <- newTVarIO (0 :: Int)

  t0 <- getCurrentTime
  let handler _evt = do
        atomically $ do
          c <- readTVar countVar
          writeTVar countVar (c + 1)
        pure Continue
  let cfg =
        SubscriptionConfig
          { name = subName,
            target = AllStreams,
            handler = handler,
            batchSize = either (error . show) Prelude.id (mkBatchSize 500),
            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
  waitForCount countVar n 30_000_000
  cancel handle
  t1 <- getCurrentTime
  pure (realToFrac (diffUTCTime t1 t0))

-- | Benchmark: subscriptionStream consumed via Streamly fold
benchSubscriptionStream :: KirokuStore -> Int -> IO Int -> IO Double
benchSubscriptionStream store n nextId = do
  subId <- nextId
  let subName = SubscriptionName ("stream-bench-" <> T.pack (show subId))

  let cfg =
        SubscriptionConfig
          { name = subName,
            target = AllStreams,
            handler = \_ -> pure Continue,
            batchSize = either (error . show) Prelude.id (mkBatchSize 500),
            queueCapacity = 16,
            overflowPolicy = DropSubscription,
            consumerGroup = Nothing,
            consumerGroupGuard = False,
            missingCheckpointPolicy = FromBeginning,
            targetBindingPolicy = AdoptUnbound,
            undecodableHandler = Nothing,
            handlerStallWarnAfter = Nothing,
            retryPolicy = defaultRetryPolicy,
            eventTypeFilter = AllEventTypes,
            selector = Nothing
          }

  t0 <- getCurrentTime
  (stream, cancelStream) <- subscriptionStream store cfg (either (error . show) Prelude.id (Buffer.mkStreamBufferSize 256))
  Stream.fold Fold.drain (Stream.take n stream)
  cancelStream
  t1 <- getCurrentTime
  pure (realToFrac (diffUTCTime t1 t0))

-- | Benchmark: Shibuya adapter with runApp
benchShibuyaAdapter :: KirokuStore -> Int -> IO Int -> IO Double
benchShibuyaAdapter store n nextId = do
  subId <- nextId
  let subName = SubscriptionName ("shibuya-bench-" <> T.pack (show subId))
  countVar <- newTVarIO (0 :: Int)

  t0 <- getCurrentTime
  runEff $ runTracingNoop $ do
    let cfg =
          SubscriptionConfig
            { name = subName,
              target = AllStreams,
              handler = \_ -> pure Continue,
              batchSize = either (error . show) Prelude.id (mkBatchSize 500),
              queueCapacity = 16,
              overflowPolicy = DropSubscription,
              consumerGroup = Nothing,
              consumerGroupGuard = False,
              missingCheckpointPolicy = FromBeginning,
              targetBindingPolicy = AdoptUnbound,
              undecodableHandler = Nothing,
              handlerStallWarnAfter = Nothing,
              retryPolicy = defaultRetryPolicy,
              eventTypeFilter = AllEventTypes,
              selector = Nothing
            }

    (ioStream, cancelAction) <- liftIO $ subscriptionStream store cfg (either (error . show) Prelude.id (Buffer.mkStreamBufferSize 256))

    let effStream = Stream.morphInner liftIO ioStream
        ingestedStream = fmap (mkIngested cancelAction) effStream

    let adapter =
          Adapter
            { adapterName = "kiroku-bench",
              source = ingestedStream,
              shutdown = liftIO cancelAction
            }

    let handler :: (IOE :> es) => Message es RecordedEvent -> Eff es AckDecision
        handler _ingested = do
          liftIO $ atomically $ do
            c <- readTVar countVar
            writeTVar countVar (c + 1)
          pure AckOk

    res <- runApp defaultAppConfig [(ProcessorId "bench", mkProcessor adapter handler)]
    case res of
      Left err -> liftIO $ error ("runApp failed: " <> show err)
      Right appHandle -> do
        liftIO $ waitForCount countVar n 30_000_000
        stopApp appHandle

  t1 <- getCurrentTime
  pure (realToFrac (diffUTCTime t1 t0))

-- Helpers

mkIngested :: (IOE :> es) => IO () -> RecordedEvent -> Ingested es RecordedEvent
mkIngested cancelAction event =
  Ingested
    { envelope =
        Envelope
          { messageId = MessageId (T.pack (show (event ^. #globalPosition))),
            cursor = let GlobalPosition pos = event ^. #globalPosition in Just (CursorInt (fromIntegral pos)),
            partition = Nothing,
            enqueuedAt = Just (event ^. #createdAt),
            traceContext = Nothing,
            headers = Nothing,
            attempt = Nothing,
            attributes = HashMap.empty,
            payload = event
          },
      ack =
        AckHandle
          { finalize = \case
              AckHalt _ -> liftIO cancelAction
              _ -> pure ()
          },
      lease = Nothing,
      deliveryStatus = pure DeliveryCurrent
    }

waitForCount :: STM.TVar Int -> Int -> Int -> IO ()
waitForCount countVar target timeoutMicros = do
  timeoutVar <- registerDelay timeoutMicros
  ok <-
    atomically $
      (readTVar countVar >>= \c -> STM.check (c >= target) >> pure True)
        `STM.orElse` (readTVar timeoutVar >>= STM.check >> pure False)
  actual <- readTVarIO countVar
  if ok
    then pure ()
    else error ("Timed out: expected " <> show target <> " events, got " <> show actual)

median :: [Double] -> Double
median xs =
  let sorted = sort xs
      len = length sorted
   in if even len
        then (sorted !! (len `div` 2 - 1) + sorted !! (len `div` 2)) / 2
        else sorted !! (len `div` 2)

putRow :: String -> Int -> [Double] -> IO ()
putRow label n times = do
  let med = median times
      lo = minimum times
      hi = maximum times
      throughput = fromIntegral n / med :: Double
      perEvent = med / fromIntegral n * 1_000_000 :: Double
  printf
    "%s:  median %s  [%s .. %s]  (%d events/s, %s/event)\n"
    label
    (showMs med)
    (showMs lo)
    (showMs hi)
    (round throughput :: Int)
    (showUs perEvent)
  hFlush stdout

showMs :: Double -> String
showMs s
  | ms < 1 = printf "%.1f ms" ms
  | ms < 100 = printf "%.0f ms" ms
  | otherwise = printf "%.0f ms" ms
  where
    ms = s * 1000

showUs :: Double -> String
showUs us
  | us >= 1000 = printf "%.1f ms" (us / 1000)
  | us >= 1 = printf "%.0f μs" us
  | otherwise = printf "%.1f μs" us

showPct :: Double -> String
showPct p
  | p >= 0 = printf "+%.0f%%" p
  | otherwise = printf "%.0f%%" p

chunksOf :: Int -> [a] -> [[a]]
chunksOf _ [] = []
chunksOf k xs = let (h, t) = splitAt k xs in h : chunksOf k t

makeEvent :: Text -> EventData
makeEvent typ =
  EventData
    { eventId = Nothing,
      eventType = EventType typ,
      payload = Aeson.object [("benchmark", Aeson.Bool True)],
      metadata = Nothing,
      causationId = Nothing,
      correlationId = Nothing
    }