Daemon.hs

module Cachix.Daemon
  ( Types.Daemon,
    Types.runDaemon,
    new,
    start,
    run,
    stop,
    stopIO,
    subscribe,
    subscribeAll,
  )
where

import Cachix.API qualified as API
import Cachix.Client.Command.Push qualified as Command.Push
import Cachix.Client.Config qualified as Config
import Cachix.Client.Config.Orphans ()
import Cachix.Client.Env as Env
import Cachix.Client.OptionsParser (DaemonOptions, PushOptions, daemonNarinfoQueryOptions)
import Cachix.Client.OptionsParser qualified as Options
import Cachix.Client.Push
import Cachix.Client.Retry (retryClientM)
import Cachix.Client.Servant (cachixClient)
import Cachix.Daemon.EventLoop qualified as EventLoop
import Cachix.Daemon.Listen as Listen
import Cachix.Daemon.Log qualified as Log
import Cachix.Daemon.Protocol as Protocol
import Cachix.Daemon.Push as Push
import Cachix.Daemon.PushManager qualified as PushManager
import Cachix.Daemon.SocketStore qualified as SocketStore
import Cachix.Daemon.Subscription as Subscription
import Cachix.Daemon.Types as Types
import Cachix.Daemon.Types.PushManager qualified as PushManager
import Cachix.Daemon.Types.SocketStore (SocketId)
import Cachix.Daemon.Worker qualified as Worker
import Cachix.Types.BinaryCache (BinaryCacheName)
import Cachix.Types.BinaryCache qualified as BinaryCache
import Control.Concurrent.STM.TMChan
import Control.Exception.Safe (catchAny, tryAny)
import Data.IORef (IORef, atomicModifyIORef', newIORef)
import Data.Text qualified as T
import Hercules.CNix.Store (Store, withStore)
import Hercules.CNix.Util qualified as CNix.Util
import Katip qualified
import Network.Socket qualified as Socket
import Network.Socket.ByteString qualified as Socket.BS
import Protolude hiding (bracket)
import System.Environment (lookupEnv)
import System.IO.Error (isResourceVanishedError)
import System.Posix.Process (getProcessID)
import System.Posix.Signals qualified as Signal
import System.Timeout (timeout)
import UnliftIO (MonadUnliftIO, withRunInIO)
import UnliftIO.Async qualified as Async
import UnliftIO.Exception (bracket)

-- | Configure a new daemon. Use 'run' to start it.
--
-- The caller must provide a logger with scribes already registered (e.g. by
-- wrapping the call in 'Log.withLogger').
new ::
  -- | The Cachix environment.
  Env ->
  -- | A handle to the Nix store.
  Store ->
  -- | Daemon-specific options.
  DaemonOptions ->
  -- | A logger with registered scribes.
  Log.Logger ->
  -- | Push options, like compression settings and number of jobs.
  PushOptions ->
  -- | The name of the binary cache to push to.
  BinaryCacheName ->
  -- | The configured daemon environment.
  IO DaemonEnv
new daemonEnv nixStore daemonOptions daemonLogger daemonPushOptions daemonCacheName = do
  Log.logMsg daemonLogger Katip.InfoS "Starting Cachix Daemon"

  daemonEventLoop <- EventLoop.new
  daemonPid <- getProcessID

  envSocket <- fmap toS <$> lookupEnv "CACHIX_DAEMON_SOCKET"
  let cmdSocket = Options.daemonSocketPath daemonOptions
  daemonSocketPath <- maybe getSocketPath pure (envSocket <|> cmdSocket)

  daemonSocketThread <- newEmptyMVar
  daemonClients <- SocketStore.newSocketStore

  daemonPushSecret <- Command.Push.getPushSecretRequired (config daemonEnv) daemonCacheName
  let authToken = getAuthTokenFromPushSecret daemonPushSecret
  daemonBinaryCache <- Push.getBinaryCache daemonLogger daemonEnv authToken daemonCacheName

  daemonSubscriptionManagerThread <- newEmptyMVar
  daemonSubscriptionManager <- Subscription.newSubscriptionManager
  let onPushEvent pushId event = do
        Subscription.pushEvent daemonSubscriptionManager pushId event
        case Types.eventMessage event of
          Types.PushFinished -> Subscription.queueUnsubscribe daemonSubscriptionManager pushId
          _ -> return ()

  let pushParams = Push.newPushParams nixStore (clientenv daemonEnv) daemonBinaryCache daemonPushSecret daemonPushOptions
  daemonPushManager <- PushManager.newPushManagerEnv daemonPushOptions (daemonNarinfoQueryOptions daemonOptions) pushParams onPushEvent daemonLogger

  daemonWorkerThreads <- newEmptyMVar

  return $ DaemonEnv {daemonOptions = daemonOptions, ..}

-- | Configure and run the daemon as a CLI command.
-- Equivalent to running 'withStore', new', and 'run', together with some signal handling.
start :: Env -> DaemonOptions -> PushOptions -> BinaryCacheName -> IO ()
start daemonEnv daemonOptions daemonPushOptions daemonCacheName =
  withStore $ \store -> do
    let logLevel =
          if Config.verbose (Env.cachixoptions daemonEnv)
            then Debug
            else Info
    logger <- Log.new Log.namespace Nothing logLevel
    result <- Log.withLogger logger $ \registeredLogger -> do
      daemon <- new daemonEnv store daemonOptions registeredLogger daemonPushOptions daemonCacheName
      void $ runDaemon daemon installSignalHandlers
      run daemon
    exitWith (toExitCode result)

-- | Run a daemon from a given configuration.
run :: DaemonEnv -> IO (Either DaemonError ())
run daemon = fmap join <$> runDaemon daemon $ do
  DaemonEnv {..} <- ask

  printConfiguration

  Async.async (runSubscriptionManager daemonSubscriptionManager)
    >>= (liftIO . putMVar daemonSubscriptionManagerThread)

  -- Start the narinfo batch processor
  PushManager.startBatchProcessor daemonPushManager

  -- Double the number of workers to keep queues filled with jobs
  let numWorkers = Options.numJobs daemonPushOptions * 2
  Worker.startWorkers
    numWorkers
    (PushManager.pmTaskQueue daemonPushManager)
    (liftIO . PushManager.runPushManager daemonPushManager . PushManager.handleTask)
    >>= (liftIO . putMVar daemonWorkerThreads)

  Async.async (Listen.listen daemonEventLoop daemonSocketPath)
    >>= (liftIO . putMVar daemonSocketThread)

  eventLoopRes <- EventLoop.run daemonEventLoop $ \case
    AddSocketClient conn ->
      SocketStore.addSocket conn (Listen.handleClient daemonEventLoop daemonClients) daemonClients
    RemoveSocketClient socketId ->
      SocketStore.removeSocket socketId daemonClients
    ReconnectSocket ->
      -- TODO: implement reconnection logic
      EventLoop.exitLoopWith (Left DaemonSocketError) daemonEventLoop
    ReceivedMessage socketId clientMsg ->
      case clientMsg of
        ClientPushRequest pushRequest -> do
          -- Create push job upfront so we can subscribe before queuing
          pushJob <- liftIO $ PushManager.newPushJob pushRequest
          let pushId = PushManager.pushId pushJob

          Katip.katipAddContext (Katip.sl "push_id" pushId) $ do
            when (Protocol.subscribeToUpdates pushRequest) $ do
              chan <- liftIO $ subscribe daemon (Just pushId)
              publisherThread <- liftIO $ Async.async $ publishToClient socketId chan daemon
              SocketStore.addPublisherThread socketId pushId publisherThread daemonClients

            -- Now add the job
            success <- queueJob pushJob
            unless success $
              Katip.logFM Katip.ErrorS "Failed to queue push request"
        ClientStop -> do
          DaemonEnv {daemonOptions = opts, daemonClients = clients} <- ask
          if Options.daemonAllowRemoteStop opts
            then EventLoop.send daemonEventLoop ShutdownGracefully
            else do
              let errorMsg = "Remote stop is disabled on this daemon"
              SocketStore.sendAll socketId (Protocol.newMessage (Protocol.DaemonError (Protocol.UnsupportedCommand errorMsg))) clients
        ClientDiagnosticsRequest -> do
          -- Get push secret to check for signing key and auth token
          let pushParams = PushManager.pmPushParams daemonPushManager
              pushSecret = pushParamsSecret pushParams
              hasSigningKey = case pushSecret of
                PushSigningKey {} -> True
                PushToken {} -> False
              maybeToken = getAuthTokenFromPushSecret pushSecret
          -- Verify auth by making an API call
          authResult <- case maybeToken of
            Nothing -> pure $ Left "No auth token configured"
            Just token -> do
              res <- liftIO $ retryClientM (clientenv daemonEnv) $ API.getCache cachixClient token daemonCacheName
              pure $ case res of
                Right _ -> Right ()
                Left err -> Left (show err)
          let diagnostics =
                Protocol.DaemonDiagnostics
                  { Protocol.diagCacheName = daemonCacheName,
                    Protocol.diagCacheUri = BinaryCache.uri daemonBinaryCache,
                    Protocol.diagCachePublic = BinaryCache.isPublic daemonBinaryCache,
                    Protocol.diagHasSigningKey = hasSigningKey,
                    Protocol.diagAuthOk = isRight authResult,
                    Protocol.diagError = either Just (const Nothing) authResult
                  }
          SocketStore.sendAll socketId (Protocol.newMessage (Protocol.DaemonDiagnosticsResult diagnostics)) daemonClients
        _ -> return ()
    ShutdownGracefully -> do
      Katip.logFM Katip.InfoS "Shutting down daemon..."
      pushResult <- shutdownGracefully
      Katip.logFM Katip.InfoS "Daemon shut down. Exiting."
      EventLoop.exitLoopWith pushResult daemonEventLoop

  return $ case eventLoopRes of
    Left err -> Left (DaemonEventLoopError err)
    Right (Left err) -> Left err
    Right (Right ()) -> Right ()

stop :: Daemon ()
stop = do
  eventloop <- asks daemonEventLoop
  EventLoop.send eventloop ShutdownGracefully

stopIO :: DaemonEnv -> IO ()
stopIO DaemonEnv {daemonEventLoop} =
  EventLoop.sendIO daemonEventLoop ShutdownGracefully

installSignalHandlers :: Daemon ()
installSignalHandlers = do
  withRunInIO $ \runInIO -> do
    mainThreadId <- myThreadId
    -- Track Ctrl+C attempts
    interruptRef <- newIORef False

    -- Install signal handlers using runInIO to properly run Daemon actions from IO
    _ <- Signal.installHandler Signal.sigTERM (Signal.Catch (runInIO (termHandler mainThreadId))) Nothing
    _ <- Signal.installHandler Signal.sigINT (Signal.Catch (runInIO (intHandler mainThreadId interruptRef))) Nothing

    return ()
  where
    -- SIGTERM: Trigger immediate shutdown
    termHandler :: ThreadId -> Daemon ()
    termHandler mainThreadId = do
      Katip.logFM Katip.InfoS "sigTERM received. Exiting immediately..."
      startExitTimer mainThreadId
      liftIO CNix.Util.triggerInterrupt
      eventLoop <- asks daemonEventLoop
      -- Signal directly to the event loop to ensure exit even if queue is full
      EventLoop.exitLoopWithFailure EventLoopClosed eventLoop

    -- SIGINT: First try to shutdown gracefully, on second press force exit
    intHandler :: ThreadId -> IORef Bool -> Daemon ()
    intHandler mainThreadId interruptRef =
      handleInterrupt mainThreadId interruptRef

    handleInterrupt :: ThreadId -> IORef Bool -> Daemon ()
    handleInterrupt mainThreadId interruptRef = do
      liftIO CNix.Util.triggerInterrupt
      isSecondInterrupt <- liftIO $ atomicModifyIORef' interruptRef (True,)
      eventLoop <- asks daemonEventLoop

      if isSecondInterrupt
        then do
          Katip.logFM Katip.InfoS "Exiting immediately..."
          startExitTimer mainThreadId
          -- Force shutdown at the event loop level to ensure exit even if queue is full
          EventLoop.exitLoopWithFailure EventLoopClosed eventLoop
          -- Interrupt the main thread right away so we don't wait for graceful shutdown.
          liftIO $ throwTo mainThreadId ExitSuccess
        else do
          Katip.logFM Katip.InfoS "Shutting down gracefully (Ctrl+C again to force exit)..."
          EventLoop.send eventLoop ShutdownGracefully

    -- Start a timer to ensure we exit even if the event loop hangs
    startExitTimer :: ThreadId -> Daemon ()
    startExitTimer mainThreadId = do
      void $ liftIO $ forkIO $ do
        threadDelay (15 * 1000 * 1000) -- 15 seconds
        throwTo mainThreadId ExitSuccess

queueJob :: PushManager.PushJob -> Daemon Bool
queueJob pushJob = do
  DaemonEnv {daemonPushManager} <- ask
  liftIO $ PushManager.runPushManager daemonPushManager (PushManager.addPushJob pushJob)

subscribe :: DaemonEnv -> Maybe PushRequestId -> IO (TMChan PushEvent)
subscribe daemonEnv maybePushId = do
  chan <- liftIO newBroadcastTMChanIO
  liftIO $ atomically $ case maybePushId of
    Just pushId -> subscribeToSTM (daemonSubscriptionManager daemonEnv) pushId (SubChannel chan)
    Nothing -> subscribeToAllSTM (daemonSubscriptionManager daemonEnv) (SubChannel chan)
  liftIO $ atomically $ dupTMChan chan

subscribeAll :: DaemonEnv -> IO (TMChan PushEvent)
subscribeAll daemonEnv = subscribe daemonEnv Nothing

-- Publish messages from the channel to the client over the socket
publishToClient :: SocketId -> TMChan PushEvent -> DaemonEnv -> IO ()
publishToClient socketId chan daemonEnv = go
  where
    go = do
      msg <- atomically $ readTMChan chan
      case msg of
        Nothing -> return ()
        Just evt -> do
          let daemonMsg = DaemonPushEvent evt
          SocketStore.sendAll socketId (Protocol.newMessage daemonMsg) (daemonClients daemonEnv)
          go

-- | Print the daemon configuration to the log.
printConfiguration :: Daemon ()
printConfiguration = do
  config <- showConfiguration
  Katip.logFM Katip.InfoS $ Katip.ls $ "Configuration:\n" <> config

-- | Fetch debug information about the daemon configuration.
showConfiguration :: Daemon Text
showConfiguration = do
  DaemonEnv {..} <- ask
  pure $
    T.intercalate
      "\n"
      [ "PID: " <> show daemonPid,
        "Socket: " <> toS daemonSocketPath,
        "Workers: " <> show (Options.numJobs daemonPushOptions),
        "Cache name: " <> toS daemonCacheName,
        "Cache URI: " <> BinaryCache.uri daemonBinaryCache,
        "Cache public keys: " <> show (BinaryCache.publicSigningKeys daemonBinaryCache),
        "Cache is public: " <> show (BinaryCache.isPublic daemonBinaryCache),
        "Compression: " <> show (Push.getCompressionMethod daemonPushOptions daemonBinaryCache)
      ]

shutdownGracefully :: Daemon (Either DaemonError ())
shutdownGracefully = do
  DaemonEnv {..} <- ask

  -- 1. Drain: set latch (reject new jobs), wait for in-flight jobs to complete
  drained <- drainWithLogging daemonPushManager

  -- Fail any remaining jobs that didn't finish after the drain timeout
  unless drained $ do
    stuck <-
      PushManager.runPushManager daemonPushManager $
        PushManager.failPendingJobs "Daemon stopped before this path finished uploading"
    Katip.logFM Katip.WarningS $
      Katip.logStr
        ("Failed " <> show (length stuck) <> " stuck jobs after drain timeout." :: Text)

  -- 2. Compute the exit result now so clients see the real exit code, even if worker cleanup below blocks on a stuck upload.
  failedJobs <-
    PushManager.runPushManager daemonPushManager PushManager.getFailedPushJobs
  let pushResult =
        if null failedJobs
          then Right ()
          else Left DaemonPushFailure

  -- 3. Stop the batch processor and close the task queue so workers exit after their current task.
  liftIO $ PushManager.stopBatchProcessor daemonPushManager
  liftIO $ PushManager.closePushManager daemonPushManager

  -- 4. Disconnect clients with a goodbye message that includes the push result
  Async.mapConcurrently_ (sayGoodbye daemonClients pushResult) =<< SocketStore.toList daemonClients

  -- 5. Stop worker threads. If drain succeeded, workers are idle and exit
  --    immediately on the closed queue. If drain timed out, workers are stuck
  --    mid-upload; interrupt them so the daemon process can exit promptly.
  withTakeMVar daemonWorkerThreads $
    if drained then Worker.stopWorkers else Worker.abortWorkers

  -- 6. Close all event subscriptions
  withTakeMVar daemonSubscriptionManagerThread (shutdownSubscriptions daemonSubscriptionManager)

  return pushResult
  where
    drainWithLogging daemonPushManager = do
      queuedStorePathCount <- PushManager.runPushManager daemonPushManager PushManager.queuedStorePathCount
      when (queuedStorePathCount > 0) $
        Katip.logFM Katip.InfoS $
          Katip.logStr $
            "Remaining store paths: " <> (show queuedStorePathCount :: Text)

      Katip.logFM Katip.DebugS "Waiting for push manager to drain remaining jobs..."
      let timeoutOptions =
            PushManager.TimeoutOptions
              { PushManager.toTimeout = 60.0,
                PushManager.toPollingInterval = 1.0
              }
      drained <- liftIO $ PushManager.drainPushManager timeoutOptions daemonPushManager
      if drained
        then Katip.logFM Katip.DebugS "Push manager drained."
        else Katip.logFM Katip.WarningS "Push manager drain timed out. Some jobs may not have completed."
      pure drained

    shutdownSubscriptions daemonSubscriptionManager subscriptionManagerThread = do
      Katip.logFM Katip.DebugS "Shutting down event manager..."
      liftIO $ stopSubscriptionManager daemonSubscriptionManager
      _ <- Async.wait subscriptionManagerThread
      Katip.logFM Katip.DebugS "Event manager shut down."

    sayGoodbye socketStore exitResult socket = do
      let clientSock = SocketStore.socket socket
          clientThread = SocketStore.handlerThread socket
          clientSocketId = SocketStore.socketId socket

      -- Stop the handler first so we have exclusive read access to the socket
      -- below. sendAll uses an independent sendLock, so the bye message still
      -- goes out cleanly.
      Async.cancel clientThread

      sendResult <-
        liftIO $
          tryAny $
            Listen.serverBye clientSocketId socketStore exitResult
      case sendResult of
        Left err ->
          Katip.logFM Katip.WarningS $
            Katip.ls $
              "Failed to send DaemonExit: " <> displayException err
        Right () -> pure ()

      -- Half-close the write side. Pending bytes (the bye message) flush to
      -- the client; the client reads them, then sees EOF.
      liftIO $ safeShutdown clientSock Socket.ShutdownSend

      -- Bounded wait for the client to read the bye message and close its end.
      mDisconnect <-
        liftIO $
          timeout disconnectTimeoutUs $
            try $
              Socket.BS.recv clientSock 4096
      case mDisconnect of
        Nothing ->
          Katip.logFM Katip.DebugS "Client did not disconnect within timeout."
        Just (Left err)
          | isResourceVanishedError err ->
              Katip.logFM Katip.DebugS "Client did not disconnect cleanly."
        Just (Left err) ->
          Katip.logFM Katip.DebugS $
            Katip.ls $
              "Client socket threw an error: " <> displayException err
        Just (Right _) ->
          Katip.logFM Katip.DebugS "Client disconnected."

      liftIO $ safeShutdown clientSock Socket.ShutdownBoth

    safeShutdown sock mode =
      Socket.shutdown sock mode `catchAny` \_ -> pure ()

    disconnectTimeoutUs :: Int
    disconnectTimeoutUs = 5_000_000

withTakeMVar :: (MonadUnliftIO m) => MVar a -> (a -> m ()) -> m ()
withTakeMVar mvar f = do
  bracket acquire release wrapper
  where
    acquire = liftIO $ tryTakeMVar mvar

    release Nothing = pure ()
    release (Just x) = liftIO $ putMVar mvar x

    wrapper Nothing = pure ()
    wrapper (Just x) = f x