NotifyGuard.hs

module Test.NotifyGuard (spec) where

import Control.Concurrent (threadDelay)
import Control.Concurrent.Async qualified as Async
import Control.Concurrent.STM (TVar, atomically, check, modifyTVar', newTVarIO, readTVar)
import Control.Exception (bracket, throwIO)
import Control.Lens ((^.))
import Data.Aeson qualified as Aeson
import Data.ByteString (ByteString)
import Data.ByteString qualified as BS
import Data.Generics.Labels ()
import Data.Text qualified as T
import Data.Text.Encoding qualified as T
import Data.Vector qualified as V
import Hasql.Connection qualified as Connection
import Hasql.Connection.Settings qualified as Conn
import Hasql.Notifications qualified as Notifications
import Kiroku.Store
import Kiroku.Test.Postgres (withMigratedTestDatabase)
import Test.CompactionHelpers qualified as CompactionHelpers
import Test.Helpers (makeEvent)
import Test.Hspec

spec :: Spec
spec =
  describe "NOTIFY trigger guard" $
    around withNotifyStore $ do
      it "emits one unchanged append payload and no lifecycle or $all payloads" $ \(store, listenerConn) -> do
        payloads <- newTVarIO []

        Notifications.listen listenerConn (Notifications.toPgIdentifier "kiroku.events")
        Async.withAsync
          (Notifications.waitForNotifications (\_ payload -> atomically (modifyTVar' payloads (<> [payload]))) listenerConn)
          $ \_listener -> do
            let firstStream = StreamName "notify-guard-1"
                secondStream = StreamName "notify-guard-2"

            Right first <-
              runStoreIO store $
                appendToStream firstStream NoStream [makeEvent "NotifyGuardFirst" (Aeson.object [])]
            waitForPayload payloads (expectedPayload firstStream first)

            Right _ <- runStoreIO store $ softDeleteStream firstStream
            Right _ <- runStoreIO store $ undeleteStream firstStream

            Right second <-
              runStoreIO store $
                appendToStream secondStream NoStream [makeEvent "NotifyGuardSecond" (Aeson.object [])]
            waitForPayload payloads (expectedPayload secondStream second)

            Right allEvents <- runStoreIO store (readAllForward (GlobalPosition 0) 10)
            input <- CompactionHelpers.inputFor store [V.head allEvents]
            Right (Right (CompactionAppliedNow _)) <- runStoreIO store (applyCompaction (CompactionHelpers.manifest input))

            actual <- atomically (readTVar payloads)
            actual
              `shouldBe` [ expectedPayload firstStream first,
                           expectedPayload secondStream second
                         ]
            map payloadFieldCount actual `shouldBe` [3, 3]
            actual `shouldSatisfy` all (not . BS.isPrefixOf "$all,")

      it "emits the unchanged payload for a 512-byte stream name" $ \(store, listenerConn) -> do
        payloads <- newTVarIO []

        Notifications.listen listenerConn (Notifications.toPgIdentifier "kiroku.events")
        Async.withAsync
          (Notifications.waitForNotifications (\_ payload -> atomically (modifyTVar' payloads (<> [payload]))) listenerConn)
          $ \_listener -> do
            let stream = StreamName (T.replicate 512 "n")

            Right result <-
              runStoreIO store $
                appendToStream stream NoStream [makeEvent "NotifyGuardMaxName" (Aeson.object [])]
            let expected = expectedPayload stream result
            waitForPayload payloads expected

            actual <- atomically (readTVar payloads)
            actual `shouldBe` [expected]
            payloadFieldCount expected `shouldBe` 3

withNotifyStore :: ((KirokuStore, Connection.Connection) -> IO ()) -> IO ()
withNotifyStore action =
  withMigratedTestDatabase $ \connStr ->
    withStore (defaultConnectionSettings connStr) $ \store ->
      withConnection connStr $ \listenerConn ->
        action (store, listenerConn)

withConnection :: T.Text -> (Connection.Connection -> IO a) -> IO a
withConnection connStr =
  bracket acquire Connection.release
  where
    acquire = do
      result <- Connection.acquire (Conn.connectionString connStr)
      case result of
        Left err -> throwIO (userError ("failed to acquire LISTEN connection: " <> show err))
        Right conn -> pure conn

expectedPayload :: StreamName -> AppendResult -> ByteString
expectedPayload (StreamName name) result =
  let StreamId sid = result ^. #streamId
      StreamVersion version = result ^. #streamVersion
   in T.encodeUtf8 $
        T.intercalate
          ","
          [ name,
            T.pack (show sid),
            T.pack (show version)
          ]

payloadFieldCount :: ByteString -> Int
payloadFieldCount = length . BS.split 44

waitForPayload :: TVar [ByteString] -> ByteString -> IO ()
waitForPayload payloads expected = do
  result <-
    Async.race
      (threadDelay 5_000_000)
      ( atomically $ do
          seen <- readTVar payloads
          check (expected `elem` seen)
      )
  case result of
    Left () -> do
      seen <- atomically (readTVar payloads)
      expectationFailure ("timed out waiting for notification " <> show expected <> "; saw " <> show seen)
    Right () -> pure ()