PushManager.hs

module Cachix.Daemon.PushManager
  ( newPushManagerEnv,
    runPushManager,
    drainPushManager,
    closePushManager,

    -- * Push strategy
    newPushStrategy,

    -- * Push job
    PushJob (..),
    newPushJob,
    addPushJob,
    addPushJobFromRequest,
    lookupPushJob,
    withPushJob,
    resolvePushJob,
    pendingJobCount,

    -- * Query
    filterPushJobs,
    getFailedPushJobs,
    failPushJob,
    failPendingJobs,

    -- * Store paths
    queueStorePaths,
    removeStorePath,
    queuedStorePathCount,

    -- * Tasks
    handleTask,

    -- * Push events
    pushStarted,
    pushFinished,
    pushStorePathAttempt,
    pushStorePathProgress,
    pushStorePathDone,
    pushStorePathFailed,

    -- * Batch processor
    startBatchProcessor,
    stopBatchProcessor,

    -- * Helpers
    atomicallyWithTimeout,
  )
where

import Cachix.Client.CNix (formatStorePathError, logStorePathWarning, resolveStorePaths)
import Cachix.Client.Command.Push hiding (pushStrategy)
import Cachix.Client.OptionsParser as Client.OptionsParser
  ( PushOptions (..),
  )
import Cachix.Client.Push as Client.Push
import Cachix.Client.Retry (retryAll)
import Cachix.Daemon.NarinfoQuery qualified as NarinfoQuery
import Cachix.Daemon.Protocol qualified as Protocol
import Cachix.Daemon.PushManager.PushJob qualified as PushJob
import Cachix.Daemon.ShutdownLatch qualified as ShutdownLatch
import Cachix.Daemon.TaskQueue
import Cachix.Daemon.Types.Log (Logger)
import Cachix.Daemon.Types.PushEvent (PushEvent (..), PushEventMessage (..), newPushRetryStatus)
import Cachix.Daemon.Types.PushManager
import Cachix.Types.BinaryCache qualified as BinaryCache
import Conduit qualified as C
import Control.Concurrent.Async qualified as Async
import Control.Concurrent.MVar qualified as MVar
import Control.Concurrent.STM.TVar
import Control.Monad.Catch qualified as E
import Control.Monad.IO.Unlift (MonadUnliftIO)
import Control.Monad.Trans.Maybe (MaybeT (..), runMaybeT)
import Control.Retry (RetryStatus, rsIterNumber)
import Data.ByteString qualified as BS
import Data.HashMap.Strict qualified as HashMap
import Data.IORef
import Data.Sequence qualified as Seq
import Data.Set qualified as Set
import Data.Text qualified as T
import Data.Time (UTCTime, diffUTCTime, getCurrentTime, secondsToNominalDiffTime)
import GHC.Clock (getMonotonicTimeNSec)
import Hercules.CNix (StorePath)
import Hercules.CNix.Store (Store, parseStorePath, storePathToPath)
import Katip qualified
import Protolude hiding (toS)
import Protolude.Conv (toS)
import Servant.Auth ()
import Servant.Auth.Client
import Servant.Conduit ()
import UnliftIO.QSem qualified as QSem

newPushManagerEnv :: (MonadIO m) => PushOptions -> NarinfoQuery.NarinfoQueryOptions -> PushParams PushManager () -> OnPushEvent -> Logger -> m PushManagerEnv
newPushManagerEnv pushOptions batchOptions pmPushParams onPushEvent pmLogger = liftIO $ do
  pmPushJobs <- newTVarIO mempty
  pmPendingJobCount <- newTVarIO 0
  pmStorePathIndex <- newTVarIO mempty
  pmJobLock <- MVar.newMVar ()
  pmTaskQueue <- atomically newTaskQueue
  pmTaskSemaphore <- QSem.newQSem (numJobs pushOptions)
  pmLastEventTimestamp <- newTVarIO =<< getCurrentTime
  let pmProgressEmitIntervalNs = 200 * 1000 * 1000
  let pmOnPushEvent id pushEvent = updateTimestampTVar pmLastEventTimestamp >> onPushEvent id pushEvent

  -- Create query manager with callback that queues ProcessQueryResponse tasks
  let batchCallback requestId response = do
        atomically $ writeTask pmTaskQueue $ HandleMissingPathsResponse requestId response
  pmNarinfoQueryManager <- NarinfoQuery.new batchOptions batchCallback
  pmShutdownLatch <- ShutdownLatch.newShutdownLatch

  return $ PushManagerEnv {..}

runPushManager :: (MonadIO m) => PushManagerEnv -> PushManager a -> m a
runPushManager env f = liftIO $ unPushManager f `runReaderT` env

-- | Set the shutdown latch (rejecting new jobs), then wait for all in-flight jobs to complete with an idle timeout.
--
-- Returns True if all jobs completed, False if no progress was observed within the timeout window.
drainPushManager :: TimeoutOptions -> PushManagerEnv -> IO Bool
drainPushManager timeoutOptions PushManagerEnv {..} = do
  ShutdownLatch.initiateShutdown () pmShutdownLatch
  -- Reset the last event timestamp to start off the countdown from here
  updateTimestampTVar pmLastEventTimestamp
  atomicallyWithTimeout timeoutOptions pmLastEventTimestamp $ do
    pendingJobs <- readTVar pmPendingJobCount
    check (pendingJobs <= 0)

-- | Close the task queue so workers see it as closed and exit.
closePushManager :: PushManagerEnv -> IO ()
closePushManager PushManagerEnv {..} =
  atomically $ closeTaskQueue pmTaskQueue

-- | Start the batch processor for narinfo queries
startBatchProcessor :: (MonadUnliftIO m, Katip.KatipContext m) => PushManagerEnv -> m ()
startBatchProcessor env@PushManagerEnv {pmNarinfoQueryManager} = do
  NarinfoQuery.start pmNarinfoQueryManager $ \paths ->
    runPushManager env (processBatchedNarinfo paths)
  where
    -- Process a batch of store paths for narinfo queries
    processBatchedNarinfo :: [StorePath] -> PushManager ([StorePath], [StorePath])
    processBatchedNarinfo storePaths = do
      pushParams <- asks pmPushParams
      queryNarInfoBulk pushParams storePaths

-- | Stop the batch processor
stopBatchProcessor :: (MonadIO m) => PushManagerEnv -> m ()
stopBatchProcessor PushManagerEnv {pmNarinfoQueryManager} = do
  NarinfoQuery.stop pmNarinfoQueryManager

-- Manage push jobs

newPushJob :: (MonadIO m) => Protocol.PushRequest -> m PushJob
newPushJob = PushJob.new

addPushJob :: PushJob -> PushManager Bool
addPushJob pushJob = do
  PushManagerEnv {..} <- ask
  let pushId = PushJob.pushId pushJob

  Katip.logLocM Katip.DebugS $ Katip.ls $ "Queued push job " <> (show pushId :: Text)

  let queueJob = do
        shuttingDown <- ShutdownLatch.isShuttingDownSTM pmShutdownLatch
        if shuttingDown
          then return False
          else do
            modifyTVar' pmPushJobs $ HashMap.insert pushId pushJob
            incrementTVar pmPendingJobCount
            res <- tryWriteTask pmTaskQueue $ QueryMissingPaths pushId
            case res of
              Just True -> return True
              _ -> return False

  didQueue <- liftIO $ atomically queueJob

  unless didQueue $
    Katip.logLocM Katip.WarningS "Failed to queue push job. Queue likely full."

  return didQueue

addPushJobFromRequest :: Protocol.PushRequest -> PushManager (Maybe Protocol.PushRequestId)
addPushJobFromRequest pushRequest = do
  pushJob <- newPushJob pushRequest
  success <- addPushJob pushJob
  return $ if success then Just (PushJob.pushId pushJob) else Nothing

removePushJob :: Protocol.PushRequestId -> PushManager ()
removePushJob pushId = do
  PushManagerEnv {..} <- ask
  liftIO $ atomically $ do
    mpushJob <- HashMap.lookup pushId <$> readTVar pmPushJobs
    for_ mpushJob $ \pushJob -> do
      -- Decrement the job count if this job had not been processed yet
      unless (PushJob.isProcessed pushJob) (decrementTVar pmPendingJobCount)
      modifyTVar' pmPushJobs (HashMap.delete pushId)

lookupPushJob :: Protocol.PushRequestId -> PushManager (Maybe PushJob)
lookupPushJob pushId = do
  pushJobs <- asks pmPushJobs
  liftIO $ HashMap.lookup pushId <$> readTVarIO pushJobs

filterPushJobs :: (PushJob -> Bool) -> PushManager [PushJob]
filterPushJobs f = do
  pushJobs <- asks pmPushJobs
  liftIO $ filter f . HashMap.elems <$> readTVarIO pushJobs

getFailedPushJobs :: PushManager [PushJob]
getFailedPushJobs = filterPushJobs PushJob.isFailed

withPushJob :: Protocol.PushRequestId -> (PushJob -> PushManager ()) -> PushManager ()
withPushJob pushId f =
  maybe handleMissingPushJob f =<< lookupPushJob pushId
  where
    handleMissingPushJob =
      Katip.logLocM Katip.ErrorS $ Katip.ls $ "Push job " <> (show pushId :: Text) <> " not found"

-- | Apply an update to many push jobs atomically. Jobs that are already in a
-- terminal state are left untouched. After the update, a job transitions to
-- a terminal state if either the update made it processed directly or it
-- leaves 'pushQueue' empty. Jobs that reach a terminal state are removed from
-- the store path index so that in-flight uploads stop reporting to them.
-- Returns the jobs that were updated; use 'PushJob.isProcessed' to pick out
-- the ones that finished in this call.
--
-- Mark, completion check, index cleanup, and pending-counter decrement happen
-- in one STM transaction so concurrent workers cannot race past the
-- empty-queue check.
applyPushJobUpdates ::
  (Foldable t) =>
  t Protocol.PushRequestId ->
  (UTCTime -> PushJob -> PushJob) ->
  PushManager [PushJob]
applyPushJobUpdates pushIds update = do
  PushManagerEnv {pmPushJobs, pmPendingJobCount, pmStorePathIndex} <- ask
  ts <- liftIO getCurrentTime
  liftIO $ atomically $ do
    jobs <- readTVar pmPushJobs
    let (jobs', updated) = foldl' (step ts) (jobs, []) pushIds
        finished = filter PushJob.isProcessed updated
    writeTVar pmPushJobs jobs'
    modifyTVar' pmPendingJobCount (subtract (length finished))
    unless (null finished) $ do
      let finishedIds = Set.fromList (map PushJob.pushId finished)
      modifyTVar' pmStorePathIndex $ HashMap.map (Seq.filter (`Set.notMember` finishedIds))
    pure updated
  where
    step ts (!jobs, acc) pushId =
      case HashMap.lookup pushId jobs of
        Just job
          | not (PushJob.isProcessed job) ->
              let job' = transitionIfDone ts (update ts job)
               in (HashMap.insert pushId job' jobs, job' : acc)
        _ -> (jobs, acc)

    transitionIfDone ts job
      | PushJob.isProcessed job = job
      | Set.null (PushJob.pushQueue job) =
          if PushJob.hasFailedPaths job
            then PushJob.fail ts job
            else PushJob.complete ts job
      | otherwise = job

failPushJob :: Protocol.PushRequestId -> Text -> PushManager ()
failPushJob pushId reason = void $ failPushJobs [pushId] reason

-- | Mark every non-terminal job as failed and emit failure events for their
-- remaining paths.
failPendingJobs :: Text -> PushManager [PushJob]
failPendingJobs reason = do
  pmPushJobs <- asks pmPushJobs
  allIds <- HashMap.keys <$> liftIO (readTVarIO pmPushJobs)
  failPushJobs allIds reason

-- | Mark the given jobs as failed, emitting a 'PushStorePathFailed' event for
-- each path that will no longer be pushed, followed by 'PushFinished' so that
-- subscribers stop waiting. The jobs are dropped from the store path index at
-- the same time, so uploads that are still in flight do not emit any further
-- events for them. Failed jobs stay in 'pmPushJobs' so a later
-- 'getFailedPushJobs' call sees them and the daemon exits with the right
-- code; the natural path removes jobs via 'pushFinished'.
failPushJobs :: (Foldable t) => t Protocol.PushRequestId -> Text -> PushManager [PushJob]
failPushJobs pushIds reason = withJobLock $ do
  failedJobs <- applyPushJobUpdates pushIds PushJob.fail

  ts <- liftIO getCurrentTime
  sendPushEvent <- asks pmOnPushEvent
  for_ failedJobs $ \job -> do
    let pid = PushJob.pushId job
    for_ (unpushedPaths job) $ \path ->
      sendStorePathEventAt ts [pid] (PushStorePathFailed path reason)
    liftIO $ sendPushEvent pid (PushEvent ts pid PushFinished)

  pure failedJobs
  where
    -- A job that fails before closure resolution has an empty queue; fall
    -- back to the requested paths so subscribers still see the failure.
    unpushedPaths job
      | isNothing (PushJob.startedAt job) = Set.fromList $ Protocol.storePaths (pushRequest job)
      | otherwise = PushJob.pushQueue job

pendingJobCount :: PushManager Int
pendingJobCount = do
  pmPendingJobCount <- asks pmPendingJobCount
  liftIO $ readTVarIO pmPendingJobCount

-- Manage store paths

queueStorePaths :: Protocol.PushRequestId -> [FilePath] -> PushManager ()
queueStorePaths pushId storePaths = do
  PushManagerEnv {..} <- ask

  let addToQueue storePath = do
        isDuplicate <- HashMap.member storePath <$> readTVar pmStorePathIndex
        unless isDuplicate $
          writeTask pmTaskQueue (PushPath storePath)

        modifyTVar' pmStorePathIndex $ HashMap.insertWith (<>) storePath (Seq.singleton pushId)

  transactionally $ map addToQueue storePaths

removeStorePath :: FilePath -> PushManager ()
removeStorePath storePath = do
  storePathIndex <- asks pmStorePathIndex
  liftIO $ atomically $ do
    modifyTVar' storePathIndex $ HashMap.delete storePath

lookupStorePathIndex :: FilePath -> PushManager (Seq.Seq Protocol.PushRequestId)
lookupStorePathIndex storePath = do
  storePathIndex <- asks pmStorePathIndex
  references <- liftIO $ readTVarIO storePathIndex
  return $ fromMaybe Seq.empty (HashMap.lookup storePath references)

queuedStorePathCount :: PushManager Integer
queuedStorePathCount = do
  pmPushJobs <- asks pmPushJobs
  jobs <- liftIO $ readTVarIO pmPushJobs
  pure $ foldl' countQueuedPaths 0 (HashMap.elems jobs)
  where
    countQueuedPaths acc job = acc + fromIntegral (Set.size $ pushQueue job)

resolvePushJob :: Protocol.PushRequestId -> PushJob.ResolvedClosure FilePath -> PushManager ()
resolvePushJob pushId closure = withJobLock $ do
  Katip.logLocM Katip.DebugS $ Katip.ls $ showClosureStats closure

  updatedJobs <- applyPushJobUpdates [pushId] (PushJob.populateQueue closure)

  -- A job that already reached a terminal state (e.g. failed during shutdown)
  -- must not start pushing, or it would emit events after 'PushFinished'.
  when (null updatedJobs) $
    Katip.logLocM Katip.DebugS $
      Katip.ls $
        "Push job " <> (show pushId :: Text) <> " is no longer active, ignoring resolved closure"

  for_ updatedJobs $ \pushJob -> do
    pushStarted pushJob
    let skippedPaths = Set.difference (PushJob.rcAllPaths closure) (PushJob.rcMissingPaths closure)
    ts <- liftIO getCurrentTime
    forM_ skippedPaths $ \path ->
      sendStorePathEventAt ts [pushId] (PushStorePathSkipped path)
    queueStorePaths pushId $ Set.toList (PushJob.rcMissingPaths closure)
    when (PushJob.isProcessed pushJob) $
      pushFinished pushJob
  where
    showClosureStats :: PushJob.ResolvedClosure FilePath -> Text
    showClosureStats PushJob.ResolvedClosure {..} =
      let skippedPaths = Set.difference rcAllPaths rcMissingPaths
          queuedCount = length rcMissingPaths
          skippedCount = length skippedPaths
          totalCount = queuedCount + skippedCount
       in T.intercalate
            "\n"
            [ "Resolved push job " <> show pushId,
              "Total paths: " <> show totalCount,
              "Queued paths: " <> show queuedCount,
              "Skipped paths: " <> show skippedCount
            ]

handleTask :: Task -> PushManager ()
handleTask task = do
  pushParams <- asks pmPushParams
  case task of
    QueryMissingPaths pushId ->
      runQueryMissingPathsTask pushParams pushId
    HandleMissingPathsResponse pushId batchResponse ->
      runHandleMissingPathsResponseTask pushParams pushId batchResponse
    PushPath filePath ->
      runPushPathTask pushParams filePath

runQueryMissingPathsTask :: PushParams PushManager () -> Protocol.PushRequestId -> PushManager ()
runQueryMissingPathsTask pushParams pushId =
  resolveClosure `withException` failJob
  where
    failJob :: SomeException -> PushManager ()
    failJob err = do
      failPushJob pushId $ "Failed to resolve closure: " <> toS (displayException err)

      Katip.katipAddContext (Katip.sl "error" (displayException err)) $
        Katip.logLocM Katip.ErrorS $
          Katip.ls $
            "Failed to resolve closure for push job " <> (show pushId :: Text)

    resolveClosure = do
      Katip.logLocM Katip.DebugS $ Katip.ls $ "Resolving closure for push job " <> (show pushId :: Text)

      withPushJob pushId $ \pushJob -> do
        let sps = Protocol.storePaths (pushRequest pushJob)
            store = pushParamsStore pushParams

        -- Resolve paths and track which ones are invalid
        (errors, validPaths) <- liftIO $ resolveStorePaths store sps

        -- Log warnings for invalid paths
        liftIO $ for_ errors $ uncurry logStorePathWarning

        -- Emit PushStorePathInvalid events for invalid paths
        ts <- liftIO getCurrentTime
        forM_ errors $ \(path, err) ->
          sendStorePathEventForActiveJobsAt ts [pushId] (PushStorePathInvalid path (formatStorePathError err))

        paths <- computeClosure store validPaths

        -- Use async batch manager for narinfo queries (non-blocking)
        batchManager <- asks pmNarinfoQueryManager
        NarinfoQuery.submitRequest batchManager pushId paths

runHandleMissingPathsResponseTask :: PushParams PushManager () -> Protocol.PushRequestId -> NarinfoQuery.NarinfoResponse -> PushManager ()
runHandleMissingPathsResponseTask pushParams pushId batchResponse =
  processQueryResponse `withException` failJob
  where
    failJob :: SomeException -> PushManager ()
    failJob err = do
      failPushJob pushId $ "Failed to query missing paths: " <> toS (displayException err)

      Katip.katipAddContext (Katip.sl "error" (displayException err)) $
        Katip.logLocM Katip.ErrorS $
          Katip.ls $
            "Failed to process batch response for push job " <> (show pushId :: Text)

    processQueryResponse = do
      Katip.logLocM Katip.DebugS $ Katip.ls $ "Processing batch response for push job " <> (show pushId :: Text)

      let allStorePaths = Set.toList $ NarinfoQuery.nrAllPaths batchResponse
          missingStorePaths = Set.toList $ NarinfoQuery.nrMissingPaths batchResponse
          store = pushParamsStore pushParams

      storePathsToPush <- pushOnClosureAttempt pushParams allStorePaths missingStorePaths

      resolvedClosure <- do
        allPaths <- liftIO $ mapM (storeToFilePath store) allStorePaths
        pathsToPush <- liftIO $ mapM (storeToFilePath store) storePathsToPush
        return $
          PushJob.ResolvedClosure
            { rcAllPaths = Set.fromList allPaths,
              rcMissingPaths = Set.fromList pathsToPush
            }

      resolvePushJob pushId resolvedClosure

runPushPathTask :: PushParams PushManager () -> FilePath -> PushManager ()
runPushPathTask pushParams filePath = do
  pushStorePath `withException` failStorePath
  where
    failStorePath =
      pushStorePathFailed filePath . toS . displayException

    pushStorePath = do
      qs <- asks pmTaskSemaphore
      E.bracket_ (QSem.waitQSem qs) (QSem.signalQSem qs) $ do
        Katip.logLocM Katip.DebugS $ Katip.ls $ "Pushing store path " <> filePath

        let store = pushParamsStore pushParams
        storePath <- liftIO $ parseStorePath store (toS filePath)

        retryAll $ uploadStorePath pushParams storePath

newPushStrategy ::
  Store ->
  Maybe Token ->
  PushOptions ->
  Text ->
  BinaryCache.CompressionMethod ->
  (StorePath -> PushStrategy PushManager ())
newPushStrategy store authToken opts cacheName compressionMethod storePath =
  let onAlreadyPresent = do
        sp <- liftIO $ storePathToPath store storePath
        Katip.logFM Katip.InfoS $ Katip.ls $ "Skipping " <> (toS sp :: Text)
        -- TODO: needs another event type here
        pushStorePathDone (toS sp)

      onError err = do
        let errText = toS (displayException err)
        sp <- liftIO $ storePathToPath store storePath
        Katip.katipAddContext (Katip.sl "error" errText) $
          Katip.logFM Katip.InfoS (Katip.ls $ "Failed to push " <> (toS sp :: Text))
        pushStorePathFailed (toS sp) errText

      onAttempt retryStatus size = do
        sp <- liftIO $ storePathToPath store storePath
        let retryContext =
              if rsIterNumber retryStatus > 0
                then Katip.katipAddContext (Katip.sl "retry" (rsIterNumber retryStatus))
                else identity
        retryContext $
          Katip.logFM Katip.InfoS $
            Katip.ls $
              "Pushing " <> (toS sp :: Text)
        pushStorePathAttempt (toS sp) size retryStatus

      onUncompressedNARStream _ size = do
        sp <- liftIO $ storePathToPath store storePath
        progressEmitIntervalNs <- asks pmProgressEmitIntervalNs
        lastEmitNsRef <- liftIO $ newIORef =<< getMonotonicTimeNSec
        currentBytesRef <- liftIO $ newIORef (0 :: Int64)
        lastEmittedBytesRef <- liftIO $ newIORef (0 :: Int64)
        C.awaitForever $ \chunk -> do
          let newBytes = fromIntegral (BS.length chunk)
          currentBytes <- liftIO $ atomicModifyIORef' currentBytesRef (\b -> (b + newBytes, b + newBytes))
          lastEmitNs <- liftIO $ readIORef lastEmitNsRef
          nowNs <- liftIO getMonotonicTimeNSec

          when (nowNs - lastEmitNs >= progressEmitIntervalNs || currentBytes == size) $ do
            liftIO $ writeIORef lastEmitNsRef nowNs
            lastEmitted <- liftIO $ readIORef lastEmittedBytesRef
            let emitBytes = currentBytes - lastEmitted
            liftIO $ writeIORef lastEmittedBytesRef currentBytes
            when (emitBytes > 0) $
              lift $
                lift $
                  pushStorePathProgress (toS sp) currentBytes emitBytes

          C.yield chunk

      onDone = do
        sp <- liftIO $ storePathToPath store storePath
        Katip.logFM Katip.InfoS $ Katip.ls $ "Pushed " <> (toS sp :: Text)
        pushStorePathDone (toS sp)
   in PushStrategy
        { onAlreadyPresent = onAlreadyPresent,
          on401 = liftIO . handleCacheResponse cacheName authToken,
          onError = onError,
          onAttempt = onAttempt,
          onUncompressedNARStream = onUncompressedNARStream,
          onDone = onDone,
          Client.Push.compressionMethod = compressionMethod,
          Client.Push.compressionLevel = Client.OptionsParser.compressionLevel opts,
          Client.Push.chunkSize = Client.OptionsParser.chunkSize opts,
          Client.Push.numConcurrentChunks = Client.OptionsParser.numConcurrentChunks opts,
          Client.Push.omitDeriver = Client.OptionsParser.omitDeriver opts
        }

-- Push events

pushStarted :: PushJob -> PushManager ()
pushStarted pushJob@PushJob {pushId} = do
  case PushJob.startedAt pushJob of
    Nothing -> return ()
    Just timestamp -> do
      sendPushEvent <- asks pmOnPushEvent
      liftIO $ do
        sendPushEvent pushId $
          PushEvent timestamp pushId PushStarted

pushFinished :: PushJob -> PushManager ()
pushFinished pushJob@PushJob {pushId} = void $ runMaybeT $ do
  let defaultDuration = 0
  pushDuration <- MaybeT $ pure $ Just (fromMaybe defaultDuration $ PushJob.duration pushJob)

  defaultCompletedAt <- liftIO getCurrentTime
  completedAt <- MaybeT $ pure $ Just (fromMaybe defaultCompletedAt $ PushJob.completedAt pushJob)

  Katip.logLocM Katip.InfoS $
    Katip.ls $
      T.intercalate
        " "
        [ "Push job",
          show pushId :: Text,
          "finished in",
          show pushDuration
        ]

  sendPushEvent <- asks pmOnPushEvent
  liftIO $ do
    sendPushEvent pushId $
      PushEvent completedAt pushId PushFinished

  lift $ removePushJob pushId

sendStorePathEvent :: (Foldable f) => f Protocol.PushRequestId -> PushEventMessage -> PushManager ()
sendStorePathEvent pushIds msg = do
  timestamp <- liftIO getCurrentTime
  sendStorePathEventAt timestamp pushIds msg

sendStorePathEventAt :: (Foldable f) => UTCTime -> f Protocol.PushRequestId -> PushEventMessage -> PushManager ()
sendStorePathEventAt timestamp pushIds msg = do
  sendPushEvent <- asks pmOnPushEvent
  liftIO $ forM_ pushIds $ \pushId ->
    sendPushEvent pushId (PushEvent timestamp pushId msg)

pushStorePathAttempt :: FilePath -> Int64 -> RetryStatus -> PushManager ()
pushStorePathAttempt storePath size retryStatus = withJobLock $ do
  let pushRetryStatus = newPushRetryStatus retryStatus
  pushIds <- lookupStorePathIndex storePath
  sendStorePathEvent pushIds (PushStorePathAttempt storePath size pushRetryStatus)

pushStorePathProgress :: FilePath -> Int64 -> Int64 -> PushManager ()
pushStorePathProgress storePath currentBytes newBytes = withJobLock $ do
  pushIds <- lookupStorePathIndex storePath
  sendStorePathEvent pushIds (PushStorePathProgress storePath currentBytes newBytes)

pushStorePathDone :: FilePath -> PushManager ()
pushStorePathDone storePath =
  finishStorePath storePath PushJob.markStorePathPushed (PushStorePathDone storePath)

pushStorePathFailed :: FilePath -> Text -> PushManager ()
pushStorePathFailed storePath errMsg =
  finishStorePath storePath PushJob.markStorePathFailed (PushStorePathFailed storePath errMsg)

-- | Record the outcome of a store path push on every job that is still
-- waiting for it, then emit the event to those jobs only. Jobs that already
-- reached a terminal state (e.g. failed during shutdown) receive nothing, so
-- 'PushFinished' stays the last event a subscriber sees.
finishStorePath :: FilePath -> (FilePath -> PushJob -> PushJob) -> PushEventMessage -> PushManager ()
finishStorePath storePath markStorePath msg = withJobLock $ do
  pushIds <- lookupStorePathIndex storePath
  updatedJobs <- applyPushJobUpdates pushIds (\_ -> markStorePath storePath)
  sendStorePathEvent (map PushJob.pushId updatedJobs) msg
  for_ (filter PushJob.isProcessed updatedJobs) pushFinished
  removeStorePath storePath

-- Helpers

-- | Run a job operation without interleaving its state changes and events with
-- another job operation. STM keeps the internal state consistent; this lock
-- also keeps the externally visible event stream consistent with that state.
withJobLock :: PushManager a -> PushManager a
withJobLock action = do
  env@PushManagerEnv {pmJobLock} <- ask
  liftIO $ MVar.withMVar pmJobLock $ \_ -> runPushManager env action

sendStorePathEventForActiveJobsAt :: (Foldable f) => UTCTime -> f Protocol.PushRequestId -> PushEventMessage -> PushManager ()
sendStorePathEventForActiveJobsAt timestamp pushIds msg = withJobLock $ do
  pushJobs <- asks pmPushJobs >>= liftIO . readTVarIO
  let isActive pushId = maybe False (not . PushJob.isProcessed) $ HashMap.lookup pushId pushJobs
  sendStorePathEventAt timestamp (filter isActive $ toList pushIds) msg

storeToFilePath :: (MonadIO m) => Store -> StorePath -> m FilePath
storeToFilePath store storePath = do
  fp <- liftIO $ storePathToPath store storePath
  pure $ toS fp

withException :: (E.MonadCatch m) => m a -> (SomeException -> m a) -> m a
withException action handler = action `E.catchAll` (\e -> handler e >> E.throwM e)

-- STM helpers

transactionally :: (Foldable t, MonadIO m) => t (STM ()) -> m ()
transactionally = liftIO . atomically . sequence_

updateTimestampTVar :: (MonadIO m) => TVar UTCTime -> m ()
updateTimestampTVar tvar = liftIO $ do
  now <- getCurrentTime
  atomically $ writeTVar tvar now

incrementTVar :: TVar Int -> STM ()
incrementTVar tvar = modifyTVar' tvar (+ 1)

decrementTVar :: TVar Int -> STM ()
decrementTVar tvar = modifyTVar' tvar (subtract 1)

-- | Run a transaction with a timeout.
-- Returns True if the transaction completed, False if timed out.
atomicallyWithTimeout ::
  TimeoutOptions ->
  -- | A TVar timestamp to compare against
  TVar UTCTime ->
  -- | The transaction to run
  STM () ->
  IO Bool
atomicallyWithTimeout TimeoutOptions {..} timeVar transaction = do
  timeoutVar <- newTVarIO False
  Async.withAsync (updateShutdownTimeout timeoutVar) $ \_ ->
    waitForGracefulShutdown timeoutVar
  where
    waitForGracefulShutdown timeout =
      atomically $
        (transaction >> return True)
          `orElse` (checkShutdownTimeout timeout >> return False)

    updateShutdownTimeout timeoutVar =
      forever $ do
        now <- getCurrentTime
        atomically $ do
          timestamp <- readTVar timeVar
          let isTimeout =
                secondsToNominalDiffTime (realToFrac toTimeout) <= now `diffUTCTime` timestamp
          writeTVar timeoutVar isTimeout
        threadDelay $ ceiling (toPollingInterval * 1000.0 * 1000.0)

    checkShutdownTimeout timeout = check =<< readTVar timeout