NarinfoQuerySpec.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Daemon.NarinfoQuerySpec where

import Cachix.Daemon.NarinfoQuery (NarinfoQueryManager, NarinfoQueryOptions (..), NarinfoResponse, defaultNarinfoQueryOptions)
import Cachix.Daemon.NarinfoQuery qualified as NarinfoQuery
import Control.Concurrent.STM
import Data.Set qualified as Set
import Data.Time (UTCTime, getCurrentTime)
import Hercules.CNix qualified as CNix
import Hercules.CNix.Store (Store, StorePath, withStoreFromURI)
import Katip qualified
import Protolude
import Test.Hspec
import UnliftIO.Async qualified as Async
import UnliftIO.Timeout (timeout)

-- Create a mock StorePath for testing
mockStorePath :: Store -> Int -> IO StorePath
mockStorePath store i = do
  let pathText = "/nix/store/" <> show i <> "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa-mock"
  CNix.parseStorePath store (encodeUtf8 pathText)

-- Test data structure to track batch processor calls
data BatchCall = BatchCall
  { bcPaths :: [StorePath],
    bcTimestamp :: UTCTime
  }
  deriving (Show, Eq)

-- Test data structure to track callback calls
data CallbackCall requestId = CallbackCall
  { ccRequestId :: requestId,
    ccResponse :: NarinfoResponse,
    ccTimestamp :: UTCTime
  }
  deriving (Show, Eq)

-- Mock batch processor that records calls and returns controlled results
createMockBatchProcessor ::
  TVar [BatchCall] -> -- Record of all batch calls
  TVar [(Set StorePath, Set StorePath)] -> -- Queue of (existing, missing) responses
  [StorePath] ->
  IO ([StorePath], [StorePath])
createMockBatchProcessor callsVar responsesVar inputPaths = do
  now <- getCurrentTime
  let call = BatchCall inputPaths now

  atomically $ modifyTVar' callsVar (call :)

  responses <- readTVarIO responsesVar
  case responses of
    [] -> return ([], inputPaths) -- Default: all missing
    ((existing, missing) : rest) -> do
      atomically $ writeTVar responsesVar rest
      let existingList = filter (`Set.member` existing) inputPaths
          missingList = filter (`Set.member` missing) inputPaths
      -- Return all paths in closure (existing paths) and missing paths
      return (existingList, missingList)

-- Mock callback that records all calls
createMockCallback ::
  TVar [CallbackCall requestId] ->
  requestId ->
  NarinfoResponse ->
  IO ()
createMockCallback callsVar requestId response = do
  now <- getCurrentTime
  let call = CallbackCall requestId response now
  atomically $ modifyTVar' callsVar (call :)

-- Test context to pass around
data TestContext requestId = TestContext
  { tcManager :: NarinfoQueryManager requestId,
    tcBatchCalls :: TVar [BatchCall],
    tcCallbackCalls :: TVar [CallbackCall requestId],
    tcResponsesQueue :: TVar [(Set StorePath, Set StorePath)]
  }

-- Helper to start batch processor asynchronously with its own Katip context
startQueryProcessorAsync :: NarinfoQueryManager requestId -> ([StorePath] -> IO ([StorePath], [StorePath])) -> IO ()
startQueryProcessorAsync manager batchProcessor =
  void $ Async.async $ do
    handleScribe <- Katip.mkHandleScribe Katip.ColorIfTerminal stderr (Katip.permitItem Katip.InfoS) Katip.V0
    let makeLogEnv = Katip.registerScribe "stderr" handleScribe Katip.defaultScribeSettings =<< Katip.initLogEnv "test" "test"
    bracket makeLogEnv Katip.closeScribes $ \le ->
      Katip.runKatipContextT le () mempty $
        NarinfoQuery.start manager (liftIO . batchProcessor)

-- Test setup helper that encapsulates common initialization
withTestManager ::
  NarinfoQueryOptions ->
  (TestContext Int -> IO a) ->
  IO a
withTestManager config action = do
  batchCalls <- newTVarIO []
  responsesQueue <- newTVarIO []
  callbackCalls <- newTVarIO []

  let callback = createMockCallback callbackCalls
      batchProcessor = createMockBatchProcessor batchCalls responsesQueue

  manager <- NarinfoQuery.new config callback
  startQueryProcessorAsync manager batchProcessor

  let testContext =
        TestContext
          { tcManager = manager,
            tcBatchCalls = batchCalls,
            tcCallbackCalls = callbackCalls,
            tcResponsesQueue = responsesQueue
          }

  finally (action testContext) (NarinfoQuery.stop manager)

-- | Wait for an STM condition to be satisfied, with a timeout.
-- Fails with an error if the timeout expires before the condition is met.
waitForSTM :: Int -> STM Bool -> IO ()
waitForSTM timeoutMicros condition = do
  result <- timeout timeoutMicros $ atomically $ do
    satisfied <- condition
    unless satisfied retry
  when (isNothing result) $
    expectationFailure "Timeout waiting for STM condition"

-- | Wait for at least @n@ callbacks to be recorded, with a 5s timeout.
waitForCallbacks :: TVar [CallbackCall requestId] -> Int -> IO ()
waitForCallbacks callsVar n =
  waitForSTM 5_000_000 $ do
    cbs <- readTVar callsVar
    return $ length cbs >= n

spec :: Spec
spec = do
  -- Initialize the CNix library
  runIO CNix.init

  describe "Batching" $ do
    it "triggers batch when size threshold is reached" $ withStoreFromURI "dummy://" $ \store -> do
      path1 <- mockStorePath store 1
      path2 <- mockStorePath store 2
      path3 <- mockStorePath store 3
      let config = defaultNarinfoQueryOptions {nqoMaxBatchSize = 2, nqoMaxWaitTime = 10} -- Large timeout, small batch
      withTestManager config $ \TestContext {..} -> do
        atomically $ writeTVar tcResponsesQueue [(Set.fromList [path1, path2, path3], Set.empty)]

        -- Submit both requests - second one reaches size threshold (2 paths total)
        NarinfoQuery.submitRequest tcManager (1 :: Int) [path1]
        NarinfoQuery.submitRequest tcManager (2 :: Int) [path2]

        -- Wait for the batch to be processed and both callbacks to arrive
        waitForCallbacks tcCallbackCalls 2

        -- Verify exactly one batch with both paths
        calls <- readTVarIO tcBatchCalls
        length calls `shouldBe` 1
        let batchPaths = case head calls of
              Just (BatchCall paths _) -> paths
              Nothing -> panic "Expected batch call"
        Set.fromList batchPaths `shouldBe` Set.fromList [path1, path2]

    it "triggers batch when timeout is reached" $ withStoreFromURI "dummy://" $ \store -> do
      path1 <- mockStorePath store 1
      let config = defaultNarinfoQueryOptions {nqoMaxBatchSize = 100, nqoMaxWaitTime = 0.05} -- Small timeout, large batch
      withTestManager config $ \TestContext {..} -> do
        atomically $ writeTVar tcResponsesQueue [(Set.fromList [path1], Set.empty)]

        -- Submit request that won't reach size threshold
        NarinfoQuery.submitRequest tcManager (1 :: Int) [path1]

        -- Wait for timeout-triggered batch and callback
        waitForCallbacks tcCallbackCalls 1

        calls <- readTVarIO tcBatchCalls
        length calls `shouldBe` 1 -- Batch triggered by timeout
    it "processes immediately when timeout is zero" $ withStoreFromURI "dummy://" $ \store -> do
      path1 <- mockStorePath store 1
      let config = defaultNarinfoQueryOptions {nqoMaxBatchSize = 100, nqoMaxWaitTime = 0} -- Immediate mode
      withTestManager config $ \TestContext {..} -> do
        atomically $ writeTVar tcResponsesQueue [(Set.fromList [path1], Set.empty)]

        NarinfoQuery.submitRequest tcManager (1 :: Int) [path1]

        waitForCallbacks tcCallbackCalls 1

        calls <- readTVarIO tcBatchCalls
        length calls `shouldBe` 1 -- Processed immediately
    it "only caches existing paths, not missing ones" $ withStoreFromURI "dummy://" $ \store -> do
      path1 <- mockStorePath store 1
      path2 <- mockStorePath store 2
      path3 <- mockStorePath store 3
      let config = defaultNarinfoQueryOptions {nqoMaxWaitTime = 0}
      withTestManager config $ \TestContext {..} -> do
        let existingPaths = Set.fromList [path1, path2]
            missingPaths = Set.fromList [path3]
        atomically $ writeTVar tcResponsesQueue [(existingPaths, missingPaths)]

        -- Submit all three paths
        NarinfoQuery.submitRequest tcManager (1 :: Int) [path1, path2, path3]

        waitForCallbacks tcCallbackCalls 1

        -- Check what's in cache - only existing paths should be cached
        cached1 <- NarinfoQuery.lookupCache tcManager path1
        cached2 <- NarinfoQuery.lookupCache tcManager path2
        cached3 <- NarinfoQuery.lookupCache tcManager path3

        cached1 `shouldBe` True -- Existing path cached
        cached2 `shouldBe` True -- Existing path cached
        cached3 `shouldBe` False -- Missing path not cached
    it "bypasses batch processor for cached paths" $ withStoreFromURI "dummy://" $ \store -> do
      path1 <- mockStorePath store 1
      path2 <- mockStorePath store 2
      let config = defaultNarinfoQueryOptions {nqoMaxWaitTime = 0}
      withTestManager config $ \TestContext {..} -> do
        atomically $ writeTVar tcResponsesQueue [(Set.fromList [path1], Set.empty), (Set.fromList [path2], Set.empty)]

        -- First request - path1 will be cached
        NarinfoQuery.submitRequest tcManager (1 :: Int) [path1]

        waitForCallbacks tcCallbackCalls 1

        -- Second request - path1 from cache, path2 goes to batch
        NarinfoQuery.submitRequest tcManager (2 :: Int) [path1, path2]

        waitForCallbacks tcCallbackCalls 2

        -- Should have 2 batch calls (one for each unique uncached path)
        calls <- readTVarIO tcBatchCalls
        length calls `shouldBe` 2

        -- Second batch should only contain path2
        let secondBatchPaths = case head calls of
              Just (BatchCall paths _) -> paths
              Nothing -> panic "Expected batch call"
        secondBatchPaths `shouldBe` [path2]

        -- Second response should contain both paths (path1 from cache + path2 from batch)
        callbacks <- readTVarIO tcCallbackCalls
        let secondResponse = case head callbacks of
              Just (CallbackCall _ response _) -> response
              Nothing -> panic "Expected callback call"
        NarinfoQuery.nrAllPaths secondResponse `shouldBe` Set.fromList [path1, path2]

    it "distributes correct paths to each request" $ withStoreFromURI "dummy://" $ \store -> do
      path1 <- mockStorePath store 1
      path2 <- mockStorePath store 2
      path3 <- mockStorePath store 3
      path4 <- mockStorePath store 4
      path5 <- mockStorePath store 5
      path6 <- mockStorePath store 6
      -- Use batch size = 5 (exact unique paths count) so batch triggers on request 2
      -- Request 1 adds 3 paths, request 2 adds 2 more unique (5 total), triggering batch
      let config = defaultNarinfoQueryOptions {nqoMaxBatchSize = 5, nqoMaxWaitTime = 10}
      withTestManager config $ \TestContext {..} -> do
        -- Setup: path1,3,5 exist; path2,4,6 missing
        let existingPaths = Set.fromList [path1, path3, path5]
            missingPaths = Set.fromList [path2, path4, path6]
        atomically $ writeTVar tcResponsesQueue [(existingPaths, missingPaths)]

        -- Request 1: paths 1,2,3 (3 unique paths, below threshold)
        NarinfoQuery.submitRequest tcManager (1 :: Int) [path1, path2, path3]
        -- Request 2: paths 3,4,5 (path 3 overlaps, adds 2 new → 5 total, triggers batch)
        NarinfoQuery.submitRequest tcManager (2 :: Int) [path3, path4, path5]

        -- Wait until both callbacks are received (deterministic, no timing dependency)
        waitForCallbacks tcCallbackCalls 2

        callbacks <- readTVarIO tcCallbackCalls

        -- Find responses by request ID
        let findResponse rid = find (\(CallbackCall r _ _) -> r == rid) callbacks
        Just (CallbackCall _ response1 _) <- return $ findResponse 1
        Just (CallbackCall _ response2 _) <- return $ findResponse 2

        -- Request 1 should get: existing=[1,3], missing=[2]
        NarinfoQuery.nrAllPaths response1 `shouldBe` Set.fromList [path1, path2, path3]
        NarinfoQuery.nrMissingPaths response1 `shouldBe` Set.fromList [path2]

        -- Request 2 should get: existing=[3,5], missing=[4]
        NarinfoQuery.nrAllPaths response2 `shouldBe` Set.fromList [path3, path4, path5]
        NarinfoQuery.nrMissingPaths response2 `shouldBe` Set.fromList [path4]

    it "deduplicates paths across requests in same batch" $ withStoreFromURI "dummy://" $ \store -> do
      path1 <- mockStorePath store 1
      path2 <- mockStorePath store 2
      let config = defaultNarinfoQueryOptions {nqoMaxBatchSize = 3, nqoMaxWaitTime = 0.1}
      withTestManager config $ \TestContext {..} -> do
        atomically $ writeTVar tcResponsesQueue [(Set.fromList [path1, path2], Set.empty)]

        -- Submit overlapping requests that will be batched together
        NarinfoQuery.submitRequest tcManager (1 :: Int) [path1, path2] -- paths 1,2
        NarinfoQuery.submitRequest tcManager (2 :: Int) [path2, path1] -- paths 2,1 (same, different order)
        waitForCallbacks tcCallbackCalls 2

        calls <- readTVarIO tcBatchCalls
        length calls `shouldBe` 1

        -- Batch should contain deduplicated paths
        let batchPaths = case head calls of
              Just (BatchCall paths _) -> paths
              Nothing -> panic "Expected batch call"
        Set.fromList batchPaths `shouldBe` Set.fromList [path1, path2]
        length batchPaths `shouldBe` 2 -- No duplicates