diff --git a/cachix.cabal b/cachix.cabal --- a/cachix.cabal +++ b/cachix.cabal @@ -1,6 +1,6 @@ cabal-version: 2.2 name: cachix -version: 1.8.0 +version: 1.9.0 synopsis: Command-line client for Nix binary cache hosting https://cachix.org @@ -88,6 +88,7 @@ Cachix.Daemon.EventLoop Cachix.Daemon.Listen Cachix.Daemon.Log + Cachix.Daemon.NarinfoQuery Cachix.Daemon.PostBuildHook Cachix.Daemon.Progress Cachix.Daemon.Protocol @@ -97,6 +98,7 @@ Cachix.Daemon.ShutdownLatch Cachix.Daemon.SocketStore Cachix.Daemon.Subscription + Cachix.Daemon.TTLCache Cachix.Daemon.Types Cachix.Daemon.Types.Daemon Cachix.Daemon.Types.Error @@ -169,6 +171,7 @@ , network , nix-narinfo , optparse-applicative + , pqueue , pretty-terminal , prettyprinter , process @@ -242,9 +245,11 @@ main-is: Main.hs hs-source-dirs: test other-modules: + Daemon.NarinfoQuerySpec Daemon.PostBuildHookSpec Daemon.ProtocolSpec Daemon.PushManagerSpec + Daemon.TTLCacheSpec DeploySpec InstallationModeSpec NetRcSpec @@ -267,6 +272,7 @@ , hercules-ci-cnix-store , here , hspec + , katip , protolude , retry , servant-auth-client @@ -274,6 +280,7 @@ , stm , temporary , time + , unliftio , versions build-tool-depends: hspec-discover:hspec-discover diff --git a/src/Cachix/Client.hs b/src/Cachix/Client.hs --- a/src/Cachix/Client.hs +++ b/src/Cachix/Client.hs @@ -11,6 +11,7 @@ ( CachixCommand (..), DaemonCommand (..), PushArguments (..), + daemonNarinfoQueryOptions, getOpts, ) import Cachix.Client.Version (cachixVersion) @@ -39,8 +40,8 @@ Config configCommand -> Config.run cachixOptions configCommand Daemon (DaemonRun daemonOptions pushOptions mcacheName) -> Daemon.start env daemonOptions pushOptions mcacheName Daemon (DaemonStop daemonOptions) -> Daemon.Client.stop env daemonOptions - Daemon (DaemonPushPaths daemonOptions storePaths) -> Daemon.Client.push env daemonOptions storePaths - Daemon (DaemonWatchExec pushOptions cacheName cmd args) -> Command.watchExecDaemon env pushOptions cacheName cmd args + Daemon (DaemonPushPaths daemonOptions daemonPushOptions storePaths) -> Daemon.Client.push env daemonOptions daemonPushOptions storePaths + Daemon (DaemonWatchExec daemonOptions pushOptions cacheName cmd args) -> Command.watchExecDaemon env pushOptions (daemonNarinfoQueryOptions daemonOptions) cacheName cmd args DeployCommand (DeployOptions.Agent opts) -> AgentCommand.run cachixOptions opts DeployCommand (DeployOptions.Activate opts) -> ActivateCommand.run env opts GenerateKeypair name -> Command.generateKeypair env name @@ -52,8 +53,8 @@ Remove name -> Command.remove env name Use name useOptions -> Command.use env name useOptions Version -> putText cachixVersion - WatchExec watchExecMode pushArgs name cmd args -> - Command.watchExec env watchExecMode pushArgs name cmd args + WatchExec watchExecMode pushArgs batchOptions name cmd args -> + Command.watchExec env watchExecMode pushArgs batchOptions name cmd args WatchStore watchArgs name -> Command.watchStore env watchArgs name -- | Install client-wide signal handlers. diff --git a/src/Cachix/Client/Command/Push.hs b/src/Cachix/Client/Command/Push.hs --- a/src/Cachix/Client/Command/Push.hs +++ b/src/Cachix/Client/Command/Push.hs @@ -17,7 +17,6 @@ import Cachix.Client.Env (Env (..)) import Cachix.Client.Exception (CachixException (..)) import Cachix.Client.HumanSize (humanSize) -import Cachix.Client.OptionsParser (PushOptions (..)) import Cachix.Client.OptionsParser as Options (PushOptions (..)) import Cachix.Client.Push as Push import Cachix.Client.Retry (retryHttp) diff --git a/src/Cachix/Client/Command/Watch.hs b/src/Cachix/Client/Command/Watch.hs --- a/src/Cachix/Client/Command/Watch.hs +++ b/src/Cachix/Client/Command/Watch.hs @@ -17,6 +17,7 @@ import Cachix.Client.Push import Cachix.Client.WatchStore qualified as WatchStore import Cachix.Daemon qualified as Daemon +import Cachix.Daemon.NarinfoQuery (NarinfoQueryOptions) import Cachix.Daemon.PostBuildHook qualified as Daemon.PostBuildHook import Cachix.Daemon.Progress qualified as Daemon.Progress import Cachix.Daemon.Types @@ -26,8 +27,6 @@ import Data.Conduit.Combinators qualified as C import Data.Conduit.TMChan qualified as C import Data.Generics.Labels () -import Data.HashMap.Strict as HashMap -import Data.IORef import Data.Text.IO (hGetLine) import GHC.IO.Handle (hDuplicate, hDuplicateTo) import Hercules.CNix.Store (withStore) @@ -50,18 +49,18 @@ -- -- In auto mode, registers a post-build hook if the user is trusted. -- Otherwise, falls back to watching the entire Nix store. -watchExec :: Env -> WatchExecMode -> PushOptions -> BinaryCacheName -> Text -> [Text] -> IO () -watchExec env watchExecMode pushOptions cacheName cmd args = do +watchExec :: Env -> WatchExecMode -> PushOptions -> NarinfoQueryOptions -> BinaryCacheName -> Text -> [Text] -> IO () +watchExec env watchExecMode pushOptions batchOptions cacheName cmd args = do case watchExecMode of PostBuildHook -> - watchExecDaemon env pushOptions cacheName cmd args + watchExecDaemon env pushOptions batchOptions cacheName cmd args Store -> watchExecStore env pushOptions cacheName cmd args Auto -> do nixEnv <- InstallationMode.getNixEnv if InstallationMode.isTrusted nixEnv - then watchExecDaemon env pushOptions cacheName cmd args + then watchExecDaemon env pushOptions batchOptions cacheName cmd args else do putErrText fallbackWarning watchExecStore env pushOptions cacheName cmd args @@ -72,12 +71,17 @@ -- | Run a command and push any new paths to the binary cache. -- -- Requires the user to be a trusted user in a multi-user installation. -watchExecDaemon :: Env -> PushOptions -> BinaryCacheName -> Text -> [Text] -> IO () -watchExecDaemon env pushOpts cacheName cmd args = +watchExecDaemon :: Env -> PushOptions -> NarinfoQueryOptions -> BinaryCacheName -> Text -> [Text] -> IO () +watchExecDaemon env pushOpts batchOptions cacheName cmd args = Daemon.PostBuildHook.withSetup Nothing $ \hookEnv -> withTempFile (Daemon.PostBuildHook.tempDir hookEnv) "daemon-log-capture" $ \_ logHandle -> withStore $ \store -> do - let daemonOptions = DaemonOptions {daemonSocketPath = Just (Daemon.PostBuildHook.daemonSock hookEnv)} + let daemonOptions = + DaemonOptions + { daemonAllowRemoteStop = False, + daemonNarinfoQueryOptions = batchOptions, + daemonSocketPath = Just (Daemon.PostBuildHook.daemonSock hookEnv) + } daemon <- Daemon.new env store daemonOptions (Just logHandle) pushOpts cacheName exitCode <- @@ -98,7 +102,7 @@ -- Launch the daemon in the background and subscribe to all push events startDaemonThread daemon = do daemonThread <- Async.async $ Daemon.run daemon - daemonChan <- Daemon.subscribe daemon + daemonChan <- Daemon.subscribeAll daemon return (daemonThread, daemonChan) shutdownDaemonThread daemon logHandle (daemonThread, daemonChan) = do @@ -113,34 +117,10 @@ ExitSuccess -> return () postWatchExec chan = do - statsRef <- newIORef HashMap.empty + progressState <- Daemon.Progress.newProgressState runConduit $ C.sourceTMChan chan - .| C.mapM_ (displayPushEvent statsRef) - - -- Deduplicate events by path and retry status - displayPushEvent statsRef PushEvent {eventMessage} = liftIO $ do - stats <- readIORef statsRef - case eventMessage of - PushStorePathAttempt path pathSize retryStatus -> do - case HashMap.lookup path stats of - Just (progress, prevRetryStatus) - | prevRetryStatus == retryStatus -> return () - | otherwise -> do - newProgress <- Daemon.Progress.update progress retryStatus - writeIORef statsRef $ HashMap.insert path (newProgress, retryStatus) stats - Nothing -> do - progress <- Daemon.Progress.new stderr (toS path) pathSize retryStatus - writeIORef statsRef $ HashMap.insert path (progress, retryStatus) stats - PushStorePathProgress path _ newBytes -> do - case HashMap.lookup path stats of - Nothing -> return () - Just (progress, _) -> Daemon.Progress.tick progress newBytes - PushStorePathDone path -> do - mapM_ (Daemon.Progress.complete . fst) (HashMap.lookup path stats) - PushStorePathFailed path _ -> do - mapM_ (Daemon.Progress.fail . fst) (HashMap.lookup path stats) - _ -> return () + .| C.mapM_ (liftIO . Daemon.Progress.displayPushEvent progressState) printLog h = getLineLoop where diff --git a/src/Cachix/Client/NixVersion.hs b/src/Cachix/Client/NixVersion.hs --- a/src/Cachix/Client/NixVersion.hs +++ b/src/Cachix/Client/NixVersion.hs @@ -33,7 +33,7 @@ (exitcode, out, err) <- readProcessWithExitCode "nix-env" ["--version"] mempty unless (err == "") $ putStrLn $ "nix-env stderr: " <> err return $ case exitcode of - ExitFailure i -> Left $ "'nix-env --version' exited with " <> show i + ExitFailure i -> Left $ "'nix-env --version' exited with " <> Protolude.show i ExitSuccess -> Right (toS out) parseNixVersion :: Text -> Either Text Versioning diff --git a/src/Cachix/Client/OptionsParser.hs b/src/Cachix/Client/OptionsParser.hs --- a/src/Cachix/Client/OptionsParser.hs +++ b/src/Cachix/Client/OptionsParser.hs @@ -2,6 +2,7 @@ ( CachixCommand (..), DaemonCommand (..), DaemonOptions (..), + DaemonPushOptions (..), PushArguments (..), -- * Push options @@ -23,6 +24,8 @@ -- * Global options Flags (..), + -- * Daemon options + -- * Misc BinaryCacheName, getOpts, @@ -33,6 +36,8 @@ import Cachix.Client.InstallationMode qualified as InstallationMode import Cachix.Client.URI (URI) import Cachix.Client.URI qualified as URI +import Cachix.Daemon.NarinfoQuery (NarinfoQueryOptions (..)) +import Cachix.Daemon.NarinfoQuery qualified as NarinfoQuery import Cachix.Deploy.OptionsParser qualified as DeployOptions import Cachix.Types.BinaryCache (BinaryCacheName) import Cachix.Types.BinaryCache qualified as BinaryCache @@ -44,6 +49,26 @@ import Protolude.Conv import Prelude qualified +-- | Create enable/disable flags with --[no-] syntax +enableDisableFlag :: + -- | Default value when no flag is provided + Bool -> + -- | Flag name (without -- prefix) + Prelude.String -> + -- | Help text describing what the flag does + Prelude.String -> + Parser Bool +enableDisableFlag defaultValue flagName helpText = + Prelude.last + <$> some + ( flag' True (hidden <> internal <> long flagName <> help ("Enable " <> helpText)) + <|> flag' False (hidden <> internal <> long ("no-" <> flagName) <> help ("Disable " <> helpText)) + <|> flag' True (long ("[no-]" <> flagName) <> help ("Enable/disable " <> helpText <> " (default: " <> defaultText <> ")")) + ) + <|> pure defaultValue + where + defaultText = if defaultValue then "enabled" else "disabled" + data Flags = Flags { configPath :: Config.ConfigPath, hostname :: Maybe URI, @@ -59,7 +84,7 @@ | Import PushOptions Text URI | Pin PinOptions | WatchStore PushOptions Text - | WatchExec WatchExecMode PushOptions Text Text [Text] + | WatchExec WatchExecMode PushOptions NarinfoQueryOptions Text Text [Text] | Use BinaryCacheName InstallationMode.UseOptions | Remove BinaryCacheName | DeployCommand DeployOptions.DeployCommand @@ -139,17 +164,24 @@ } data DaemonCommand - = DaemonPushPaths DaemonOptions [FilePath] + = DaemonPushPaths DaemonOptions DaemonPushOptions [FilePath] | DaemonRun DaemonOptions PushOptions BinaryCacheName | DaemonStop DaemonOptions - | DaemonWatchExec PushOptions BinaryCacheName Text [Text] + | DaemonWatchExec DaemonOptions PushOptions BinaryCacheName Text [Text] deriving (Show) data DaemonOptions = DaemonOptions - { daemonSocketPath :: Maybe FilePath + { daemonAllowRemoteStop :: Bool, + daemonNarinfoQueryOptions :: NarinfoQueryOptions, + daemonSocketPath :: Maybe FilePath } deriving (Show) +data DaemonPushOptions = DaemonPushOptions + { shouldWait :: Bool + } + deriving (Show) + -- | CLI parser entry point getOpts :: IO (Flags, CachixCommand) getOpts = do @@ -474,6 +506,7 @@ daemonPush = DaemonPushPaths <$> daemonOptionsParser + <*> daemonPushOptionsParser <*> many (strArgument (metavar "PATHS...")) daemonRun :: Parser DaemonCommand @@ -489,21 +522,76 @@ daemonWatchExec :: Parser DaemonCommand daemonWatchExec = DaemonWatchExec - <$> pushOptionsParser + <$> daemonOptionsParser + <*> pushOptionsParser <*> cacheNameParser <*> strArgument (metavar "CMD") <*> many (strArgument (metavar "-- ARGS")) daemonOptionsParser :: Parser DaemonOptions daemonOptionsParser = - DaemonOptions <$> socketOption + DaemonOptions + <$> remoteStopOption + <*> batchConfigParser + <*> socketOption where socketOption = optional . strOption $ long "socket" <> short 's' <> metavar "SOCKET" + <> help "Path to the daemon socket" + remoteStopOption = + enableDisableFlag True "remote-stop" "the remote stop command which allows clients to remotely shut down the daemon. Remote stop should be disabled in environments where the lifecycle of the daemon is handled by a service manager, like systemd." + +batchConfigParser :: Parser NarinfoQueryOptions +batchConfigParser = + NarinfoQueryOptions + <$> narinfoBatchSizeOption + <*> narinfoBatchTimeoutOption + <*> narinfoCacheTTLOption + <*> narinfoMaxCacheSizeOption + where + narinfoBatchSizeOption = + option auto $ + long "narinfo-batch-size" + <> metavar "INT" + <> help "Maximum number of paths to batch together (default: 100)" + <> value (NarinfoQuery.nqoMaxBatchSize NarinfoQuery.defaultNarinfoQueryOptions) + <> showDefault + + narinfoBatchTimeoutOption = + option auto $ + long "narinfo-batch-timeout" + <> metavar "SECONDS" + <> help "Maximum time to wait before processing a batch in seconds. Use 0 for immediate processing (no batching). (default: 0.5)" + <> value (realToFrac (NarinfoQuery.nqoMaxWaitTime NarinfoQuery.defaultNarinfoQueryOptions)) + + narinfoCacheTTLOption = + option auto $ + long "narinfo-cache-ttl" + <> metavar "SECONDS" + <> help "Time-to-live for cached narinfo results in seconds. Use 0 to disable caching (default: 300.0)" + <> value (realToFrac (NarinfoQuery.nqoCacheTTL NarinfoQuery.defaultNarinfoQueryOptions)) + <> showDefault + + narinfoMaxCacheSizeOption = + option auto $ + long "narinfo-max-cache-size" + <> metavar "INT" + <> help "Maximum number of entries in the narinfo cache. Use 0 for unlimited (default: 0)" + <> value (NarinfoQuery.nqoMaxCacheSize NarinfoQuery.defaultNarinfoQueryOptions) + <> showDefault + +daemonPushOptionsParser :: Parser DaemonPushOptions +daemonPushOptionsParser = + DaemonPushOptions + <$> switch + ( long "wait" + <> help "Wait for the push operation to complete" + ) + deployCommand :: Parser CachixCommand deployCommand = DeployCommand <$> DeployOptions.parser @@ -512,6 +600,7 @@ WatchExec <$> watchExecModeParser <*> pushOptionsParser + <*> batchConfigParser <*> cacheNameParser <*> strArgument (metavar "CMD") <*> many (strArgument (metavar "-- ARGS")) diff --git a/src/Cachix/Daemon.hs b/src/Cachix/Daemon.hs --- a/src/Cachix/Daemon.hs +++ b/src/Cachix/Daemon.hs @@ -7,6 +7,7 @@ stop, stopIO, subscribe, + subscribeAll, ) where @@ -14,7 +15,7 @@ import Cachix.Client.Config qualified as Config import Cachix.Client.Config.Orphans () import Cachix.Client.Env as Env -import Cachix.Client.OptionsParser (DaemonOptions, PushOptions) +import Cachix.Client.OptionsParser (DaemonOptions, PushOptions, daemonNarinfoQueryOptions) import Cachix.Client.OptionsParser qualified as Options import Cachix.Client.Push import Cachix.Daemon.EventLoop qualified as EventLoop @@ -27,6 +28,7 @@ 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 @@ -90,11 +92,11 @@ let onPushEvent = Subscription.pushEvent daemonSubscriptionManager let pushParams = Push.newPushParams nixStore (clientenv daemonEnv) daemonBinaryCache daemonPushSecret daemonPushOptions - daemonPushManager <- PushManager.newPushManagerEnv daemonPushOptions pushParams onPushEvent daemonLogger + daemonPushManager <- PushManager.newPushManagerEnv daemonPushOptions (daemonNarinfoQueryOptions daemonOptions) pushParams onPushEvent daemonLogger daemonWorkerThreads <- newEmptyMVar - return $ DaemonEnv {..} + 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. @@ -117,6 +119,9 @@ Async.async (runSubscriptionManager daemonSubscriptionManager) >>= (liftIO . putMVar daemonSubscriptionManagerThread) + -- Start the narinfo batch processor + PushManager.startBatchProcessor daemonPushManager + Worker.startWorkers (Options.numJobs daemonPushOptions) (PushManager.pmTaskQueue daemonPushManager) @@ -134,10 +139,22 @@ ReconnectSocket -> -- TODO: implement reconnection logic EventLoop.exitLoopWith (Left DaemonSocketError) daemonEventLoop - ReceivedMessage clientMsg -> + ReceivedMessage socketId clientMsg -> case clientMsg of - ClientPushRequest pushRequest -> queueJob pushRequest - ClientStop -> EventLoop.send daemonEventLoop ShutdownGracefully + ClientPushRequest pushRequest -> do + maybePushId <- queueJob pushRequest + case maybePushId of + Just pushId -> when (Protocol.subscribeToUpdates pushRequest) $ do + chan <- liftIO $ subscribe daemon (Just pushId) + void $ liftIO $ Async.async $ publishToClient socketId chan daemon + Nothing -> 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 _ -> return () ShutdownGracefully -> do Katip.logFM Katip.InfoS "Shutting down daemon..." @@ -206,22 +223,33 @@ threadDelay (15 * 1000 * 1000) -- 15 seconds throwTo mainThreadId ExitSuccess -queueJob :: Protocol.PushRequest -> Daemon () +queueJob :: Protocol.PushRequest -> Daemon (Maybe PushRequestId) queueJob pushRequest = do - daemonPushManager <- asks daemonPushManager - -- TODO: subscribe the socket to updates if available - - -- Queue the job - void $ - PushManager.runPushManager daemonPushManager (PushManager.addPushJob pushRequest) + DaemonEnv {daemonPushManager} <- ask + liftIO $ PushManager.runPushManager daemonPushManager (PushManager.addPushJob pushRequest) -subscribe :: DaemonEnv -> IO (TMChan PushEvent) -subscribe DaemonEnv {..} = do +subscribe :: DaemonEnv -> Maybe PushRequestId -> IO (TMChan PushEvent) +subscribe daemonEnv maybePushId = do chan <- liftIO newBroadcastTMChanIO - liftIO $ atomically $ do - subscribeToAllSTM daemonSubscriptionManager (SubChannel chan) - dupTMChan chan + 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 = do + forever $ 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) + -- | Print the daemon configuration to the log. printConfiguration :: Daemon () printConfiguration = do @@ -251,6 +279,9 @@ -- Stop the push manager and wait for any remaining paths to be uploaded shutdownPushManager daemonPushManager + + -- Stop the narinfo batch processor + liftIO $ PushManager.stopBatchProcessor daemonPushManager -- Stop worker threads withTakeMVar daemonWorkerThreads Worker.stopWorkers diff --git a/src/Cachix/Daemon/Client.hs b/src/Cachix/Daemon/Client.hs --- a/src/Cachix/Daemon/Client.hs +++ b/src/Cachix/Daemon/Client.hs @@ -1,21 +1,30 @@ module Cachix.Daemon.Client (push, stop) where import Cachix.Client.Env as Env -import Cachix.Client.OptionsParser (DaemonOptions (..)) +import Cachix.Client.Exception (CachixException (..)) +import Cachix.Client.OptionsParser (DaemonOptions (..), DaemonPushOptions (..)) +import Cachix.Client.OptionsParser qualified as Options import Cachix.Client.Retry qualified as Retry import Cachix.Daemon.Listen (getSocketPath) +import Cachix.Daemon.Progress qualified as Daemon.Progress import Cachix.Daemon.Protocol as Protocol +import Cachix.Daemon.Types.Error (DaemonError (..), HasExitCode (..)) +import Cachix.Daemon.Types.PushEvent (PushEvent (..), PushEventMessage (..)) import Control.Concurrent.Async qualified as Async import Control.Concurrent.STM.TBMQueue import Data.Aeson qualified as Aeson import Data.ByteString qualified as BS +import Data.ByteString.Char8 qualified as BS8 import Data.IORef +import Data.Text qualified as T import Data.Time.Clock +import Hercules.CNix.Store qualified as Store import Network.Socket qualified as Socket import Network.Socket.ByteString qualified as Socket.BS import Network.Socket.ByteString.Lazy qualified as Socket.LBS import Protolude import System.Environment (lookupEnv) +import System.IO (hIsTerminalDevice) import System.IO.Error (isResourceVanishedError) data SocketError @@ -33,52 +42,21 @@ SocketStalled -> "The socket has stopped responding to pings" SocketDecodingError err -> "Failed to decode the message from socket: " <> toS err --- | Queue up push requests with the daemon --- --- TODO: wait for the daemon to respond that it has received the request -push :: Env -> DaemonOptions -> [FilePath] -> IO () -push _env daemonOptions storePaths = - withDaemonConn (daemonSocketPath daemonOptions) $ \sock -> do - Socket.LBS.sendAll sock $ Protocol.newMessage pushRequest - where - pushRequest = - Protocol.ClientPushRequest $ - PushRequest {storePaths = storePaths} - --- | Tell the daemon to stop and wait for it to gracefully exit -stop :: Env -> DaemonOptions -> IO () -stop _env daemonOptions = - withDaemonConn (daemonSocketPath daemonOptions) $ \sock -> do - let size = 100 - (rx, tx) <- atomically $ (,) <$> newTBMQueue size <*> newTBMQueue size - - rxThread <- Async.async (handleIncoming rx sock) - txThread <- Async.async (handleOutgoing tx sock) - - lastPongRef <- newIORef =<< getCurrentTime - pingThread <- Async.async (runPingThread lastPongRef rx tx) +-- | Run socket communication with ping/pong handling +withSocketComm :: Socket.Socket -> (STM (Maybe (Either SocketError Protocol.DaemonMessage)) -> (Protocol.ClientMessage -> STM ()) -> IO a) -> IO a +withSocketComm sock action = do + let size = 1000 + (rx, tx) <- atomically $ (,) <$> newTBMQueue size <*> newTBMQueue size - mapM_ Async.link [rxThread, txThread, pingThread] + lastPongRef <- newIORef =<< getCurrentTime + rxThread <- Async.async (handleIncoming lastPongRef rx sock) + txThread <- Async.async (handleOutgoing tx sock) + pingThread <- Async.async (runPingThread lastPongRef rx tx) - -- Request the daemon to stop - atomically $ writeTBMQueue tx Protocol.ClientStop + let threads = [rxThread, txThread, pingThread] + mapM_ Async.link threads - fix $ \loop -> do - mmsg <- atomically (readTBMQueue rx) - case mmsg of - Nothing -> return () - Just (Left err) -> do - putErrText $ toS $ displayException err - exitFailure - Just (Right msg) -> - case msg of - Protocol.DaemonPong -> do - writeIORef lastPongRef =<< getCurrentTime - loop - Protocol.DaemonExit exitStatus -> - case exitCode exitStatus of - 0 -> exitSuccess - code -> exitWith (ExitFailure code) + finally (action (readTBMQueue rx) (writeTBMQueue tx)) (mapM_ Async.cancel threads) where runPingThread lastPongRef rx tx = go where @@ -93,27 +71,26 @@ threadDelay (2 * 1000 * 1000) go - handleOutgoing tx sock = go + handleOutgoing tx sock' = go where go = do mmsg <- atomically $ readTBMQueue tx case mmsg of Nothing -> return () Just msg -> do - Retry.retryAll $ const $ Socket.LBS.sendAll sock $ Protocol.newMessage msg + Retry.retryAll $ const $ Socket.LBS.sendAll sock' $ Protocol.newMessage msg go - handleIncoming rx sock = go BS.empty + handleIncoming lastPongRef rx sock' = go BS.empty where socketClosed = atomically $ writeTBMQueue rx (Left SocketClosed) go leftovers = do - ebs <- liftIO $ try $ Socket.BS.recv sock 4096 + ebs <- liftIO $ try $ Socket.BS.recv sock' 4096 case ebs of Left err | isResourceVanishedError err -> socketClosed Left _err -> socketClosed - -- If the socket returns 0 bytes, then it is closed Right bs | BS.null bs -> socketClosed Right bs -> do let (rawMsgs, newLeftovers) = Protocol.splitMessages (BS.append leftovers bs) @@ -125,9 +102,92 @@ putErrText terr atomically $ writeTBMQueue rx (Left (SocketDecodingError terr)) Right msg -> do - atomically $ writeTBMQueue rx (Right msg) + case msg of + Protocol.DaemonPong -> do + writeIORef lastPongRef =<< getCurrentTime + atomically $ writeTBMQueue rx (Right msg) + _ -> atomically $ writeTBMQueue rx (Right msg) go newLeftovers + +-- | Queue up push requests with the daemon and wait for completion +push :: Env -> DaemonOptions -> DaemonPushOptions -> [FilePath] -> IO () +push _env daemonOptions daemonPushOptions cliPaths = do + hasStdin <- not <$> hIsTerminalDevice stdin + inputStorePaths <- + case (hasStdin, cliPaths) of + (False, []) -> throwIO $ NoInput "You need to specify store paths either as stdin or as a command argument" + (True, []) -> map T.unpack . T.words <$> getContents + -- If we get both stdin and cli args, prefer cli args. + -- This avoids hangs in cases where stdin is non-interactive but unused by caller. + (_, paths) -> return paths + + storePaths <- Store.withStore $ \store -> do + inputStorePaths' <- mapM (Store.followLinksToStorePath store . BS8.pack) inputStorePaths + mapM (fmap (toS . BS8.unpack) . Store.storePathToPath store) inputStorePaths' + + withDaemonConn (daemonSocketPath daemonOptions) $ \sock -> do + let shouldWait = Options.shouldWait daemonPushOptions + let pushRequest = Protocol.ClientPushRequest (PushRequest {storePaths = storePaths, subscribeToUpdates = shouldWait}) + + Socket.LBS.sendAll sock $ Protocol.newMessage pushRequest + unless shouldWait exitSuccess + + withSocketComm sock $ \receive _send -> do + progressState <- Daemon.Progress.newProgressState + failureRef <- newIORef False + + fix $ \loop -> do + mmsg <- atomically receive + case mmsg of + Nothing -> return () + Just (Left err) -> do + putErrText $ toS $ displayException err + exitFailure + Just (Right msg) -> do + case msg of + Protocol.DaemonPong -> loop + Protocol.DaemonPushEvent pushEvent -> do + Daemon.Progress.displayPushEvent progressState pushEvent + case eventMessage pushEvent of + PushStorePathFailed _ _ -> do + writeIORef failureRef True + loop + PushFinished -> do + hasFailures <- readIORef failureRef + if hasFailures + then exitWith (toExitCode DaemonPushFailure) + else exitSuccess + _ -> loop + _ -> loop + +-- | Tell the daemon to stop and wait for it to gracefully exit +stop :: Env -> DaemonOptions -> IO () +stop _env daemonOptions = + withDaemonConn (daemonSocketPath daemonOptions) $ \sock -> do + withSocketComm sock $ \receive send -> do + atomically $ send Protocol.ClientStop + + fix $ \loop -> do + mmsg <- atomically receive + case mmsg of + Nothing -> return () + Just (Left err) -> do + putErrText $ toS $ displayException err + exitFailure + Just (Right msg) -> + case msg of + Protocol.DaemonPong -> loop + Protocol.DaemonExit exitStatus -> + case exitCode exitStatus of + 0 -> exitSuccess + code -> exitWith (ExitFailure code) + Protocol.DaemonError errorMsg -> do + case errorMsg of + Protocol.UnsupportedCommand errMsg -> do + putErrText $ toS errMsg + exitWith (ExitFailure 2) + _ -> loop withDaemonConn :: Maybe FilePath -> (Socket.Socket -> IO a) -> IO a withDaemonConn optionalSocketPath f = do diff --git a/src/Cachix/Daemon/Listen.hs b/src/Cachix/Daemon/Listen.hs --- a/src/Cachix/Daemon/Listen.hs +++ b/src/Cachix/Daemon/Listen.hs @@ -95,7 +95,7 @@ msgs <- catMaybes <$> mapM decodeMessage rawMsgs forM_ msgs $ \msg -> do - EventLoop.send eventloop (ReceivedMessage msg) + EventLoop.send eventloop (ReceivedMessage socketId msg) case msg of Protocol.ClientPing -> liftIO $ Socket.LBS.sendAll conn $ Protocol.newMessage DaemonPong diff --git a/src/Cachix/Daemon/NarinfoQuery.hs b/src/Cachix/Daemon/NarinfoQuery.hs new file mode 100644 --- /dev/null +++ b/src/Cachix/Daemon/NarinfoQuery.hs @@ -0,0 +1,419 @@ +{-# LANGUAGE DerivingStrategies #-} +{-# LANGUAGE NamedFieldPuns #-} + +module Cachix.Daemon.NarinfoQuery + ( -- * Types + NarinfoQueryManager, + NarinfoRequest (..), + NarinfoResponse (..), + NarinfoQueryOptions (..), + defaultNarinfoQueryOptions, + + -- * Operations + new, + submitRequest, + start, + stop, + cleanupStaleEntries, + + -- * Cache operations (for testing) + lookupCache, + nqmCache, + ) +where + +import Cachix.Daemon.TTLCache (TTLCache) +import Cachix.Daemon.TTLCache qualified as TTLCache +import Control.Concurrent.STM +import Control.Monad.Extra (partitionM) +import Control.Monad.IO.Unlift (MonadUnliftIO) +import Data.Sequence qualified as Seq +import Data.Set qualified as Set +import Data.Time (NominalDiffTime, UTCTime, addUTCTime, diffUTCTime, getCurrentTime) +import Hercules.CNix.Store (StorePath, getStorePathHash) +import Katip qualified +import Protolude +import UnliftIO.Async qualified as Async + +-- | Type alias for the positive narinfo cache. +-- Keys are store path hashes. +type NarinfoCache = TTLCache ByteString + +-- | Configuration for the narinfo query manager +data NarinfoQueryOptions = NarinfoQueryOptions + { -- | Maximum number of paths to accumulate before processing a query + nqoMaxBatchSize :: !Int, + -- | Maximum time to wait before processing a query (in seconds) + -- Use 0 for immediate processing + nqoMaxWaitTime :: !NominalDiffTime, + -- | TTL for cache entries (in seconds) + -- Use 0 to disable caching + nqoCacheTTL :: !NominalDiffTime, + -- | Maximum number of entries in the cache + -- Use 0 for unlimited + nqoMaxCacheSize :: !Int + } + deriving stock (Eq, Show) + +-- | Default configuration with reasonable values +defaultNarinfoQueryOptions :: NarinfoQueryOptions +defaultNarinfoQueryOptions = + NarinfoQueryOptions + { nqoMaxBatchSize = 100, + nqoMaxWaitTime = 0.5, -- 500ms + nqoCacheTTL = 300.0, -- 5 minutes + nqoMaxCacheSize = 0 -- Unlimited + } + +-- | A request to check narinfo for store paths +data NarinfoRequest requestId = NarinfoRequest + { -- | Unique identifier for this request + qrRequestId :: !requestId, + -- | Store paths to check + qrStorePaths :: ![StorePath], + -- | Cached existing paths + qrCachedPaths :: ![StorePath] + } + +-- | Response to a narinfo query request +data NarinfoResponse = NarinfoResponse + { -- | All paths in the dependency closure + nrAllPaths :: !(Set StorePath), + -- | Paths missing from the cache + nrMissingPaths :: !(Set StorePath) + } + deriving stock (Eq, Show) + +-- | Internal state of the query manager +data QueryState requestId = QueryState + { -- | Accumulated requests waiting to be processed + qsPendingRequests :: !(Seq (NarinfoRequest requestId)), + -- | All unique paths from pending requests + qsAccumulatedPaths :: !(Set StorePath), + -- | Time when the first request in this query group was added + qsQueryStartTime :: !(Maybe UTCTime), + -- | Current timeout TVar (Nothing if no timeout active) + qsTimeoutVar :: !(Maybe (TVar Bool)), + -- | Whether the processor should continue running + qsRunning :: !Bool + } + +-- | Manager for narinfo queries +data NarinfoQueryManager requestId = NarinfoQueryManager + { -- | Configuration + nqmConfig :: !NarinfoQueryOptions, + -- | Internal state + nqmState :: !(TVar (QueryState requestId)), + -- | Handle to the processor thread + nqmProcessorThread :: !(MVar (Async ())), + -- | Callback to handle query responses + nqmCallback :: !(requestId -> NarinfoResponse -> IO ()), + -- | Cache for path existence information + nqmCache :: !(TVar NarinfoCache), + -- | Cached current time for TTL calculations (updated periodically) + nqmCachedTime :: !(TVar UTCTime), + -- | Handle to the time refresh thread + nqmTimeRefreshThread :: !(MVar (Async ())) + } + +-- | Create a new narinfo query manager +new :: (MonadIO m) => NarinfoQueryOptions -> (requestId -> NarinfoResponse -> IO ()) -> m (NarinfoQueryManager requestId) +new nqmConfig nqmCallback = liftIO $ do + nqmState <- newTVarIO initialState + nqmProcessorThread <- newEmptyMVar + nqmCache <- newTVarIO TTLCache.empty + nqmCachedTime <- newTVarIO =<< getCurrentTime + nqmTimeRefreshThread <- newEmptyMVar + return NarinfoQueryManager {..} + where + initialState = + QueryState + { qsPendingRequests = Seq.empty, + qsAccumulatedPaths = Set.empty, + qsQueryStartTime = Nothing, + qsTimeoutVar = Nothing, + qsRunning = True + } + +-- Cache operations + +-- | Check if a path exists in the cache and is not stale +lookupCache :: (MonadIO m) => NarinfoQueryManager requestId -> StorePath -> m Bool +lookupCache NarinfoQueryManager {nqmCache, nqmCachedTime} path = liftIO $ do + pathHash <- getStorePathHash path + now <- readTVarIO nqmCachedTime + cache <- readTVarIO nqmCache + return $ TTLCache.lookup now pathHash cache + +-- | Update cache with new entries (only existing paths) +updateCache :: (MonadIO m) => NarinfoQueryManager requestId -> Set StorePath -> m () +updateCache NarinfoQueryManager {nqmCache, nqmConfig, nqmCachedTime} paths = liftIO $ do + -- Convert paths to hashes + pathHashes <- mapM getStorePathHash (Set.toList paths) + now <- readTVarIO nqmCachedTime + let ttl = nqoCacheTTL nqmConfig + when (ttl > 0) $ do + let expiresAt = addUTCTime ttl now + atomically $ modifyTVar' nqmCache $ \cache -> + let cacheWithNewEntries = foldr (`TTLCache.insert` expiresAt) cache pathHashes + maxSize = nqoMaxCacheSize nqmConfig + currentSize = TTLCache.size cacheWithNewEntries + -- Lazy cleanup: only clean up if cache is getting large (20% over limit) + cleanupThreshold = if maxSize > 0 then maxSize + (maxSize `div` 5) else maxBound + in if currentSize > cleanupThreshold + then + let cleanedCache = TTLCache.cleanupExpired now cacheWithNewEntries + finalCache = + if maxSize > 0 && TTLCache.size cleanedCache > maxSize + then TTLCache.pruneToSize maxSize cleanedCache + else cleanedCache + in finalCache + else cacheWithNewEntries + +-- | Remove stale entries from cache +cleanupStaleEntries :: (MonadIO m) => NarinfoQueryManager requestId -> m () +cleanupStaleEntries NarinfoQueryManager {nqmCache, nqmCachedTime} = liftIO $ do + now <- readTVarIO nqmCachedTime + atomically $ modifyTVar' nqmCache $ TTLCache.cleanupExpired now + +-- | Submit a request to the query manager +-- +-- The store paths will be queued for processing. +-- The closure will _not_ be computed. +submitRequest :: + (MonadIO m) => + NarinfoQueryManager requestId -> + requestId -> + [StorePath] -> + m () +submitRequest manager@NarinfoQueryManager {nqmConfig, nqmState, nqmCachedTime} requestId storePaths = liftIO $ do + -- Check cache for each path in a single pass + (cachedExistingPaths, pathsToQuery') <- partitionM (lookupCache manager) storePaths + + -- Perform the state update atomically and determine if timeout is needed + needsTimeout <- atomically $ do + now <- readTVar nqmCachedTime + batchState <- readTVar nqmState + let isFirstRequest = Seq.null (qsPendingRequests batchState) + needsNewTimeout = isFirstRequest && isNothing (qsTimeoutVar batchState) && nqoMaxWaitTime nqmConfig > 0 + updatedPaths = qsAccumulatedPaths batchState <> Set.fromList pathsToQuery' + newRequest = NarinfoRequest requestId storePaths cachedExistingPaths + newStartTime = case qsQueryStartTime batchState of + Nothing -> Just now + justTime -> justTime + + writeTVar nqmState $ + batchState + { qsPendingRequests = qsPendingRequests batchState Seq.|> newRequest, + qsAccumulatedPaths = updatedPaths, + qsQueryStartTime = newStartTime + } + + return needsNewTimeout + + -- Set up timeout outside of STM if needed + when needsTimeout $ do + timeoutVar <- startBatchTimeout (nqoMaxWaitTime nqmConfig) + atomically $ do + modifyTVar' nqmState $ \bs -> bs {qsTimeoutVar = Just timeoutVar} + +-- | Start the query processor thread +start :: + (MonadUnliftIO m, Katip.KatipContext m) => + NarinfoQueryManager requestId -> + -- | Function to process a batch of paths + ([StorePath] -> m ([StorePath], [StorePath])) -> + m () +start manager@NarinfoQueryManager {nqmProcessorThread, nqmTimeRefreshThread} processBatch = do + -- Start time refresh thread + timeThread <- Async.async $ runTimeRefreshThread manager + liftIO $ putMVar nqmTimeRefreshThread timeThread + + -- Start processor thread + thread <- Async.async $ runQueryProcessorThread manager processBatch + liftIO $ putMVar nqmProcessorThread thread + +-- | Stop the query processor +stop :: (MonadIO m) => NarinfoQueryManager requestId -> m () +stop NarinfoQueryManager {nqmState, nqmProcessorThread, nqmTimeRefreshThread} = liftIO $ do + -- Signal shutdown + atomically $ modifyTVar' nqmState $ \batchState -> batchState {qsRunning = False} + + -- Wait for processor thread to finish + thread <- tryTakeMVar nqmProcessorThread + for_ thread Async.wait + + -- Cancel the time refresh thread + timeThread <- tryTakeMVar nqmTimeRefreshThread + for_ timeThread Async.cancel + +-- | Start a timeout for the current batch +startBatchTimeout :: NominalDiffTime -> IO (TVar Bool) +startBatchTimeout delay = do + let delayMicros = ceiling (delay * 1000000) -- Convert to microseconds + registerDelay delayMicros + +-- | Check if the current query has timed out +isQueryTimedOut :: QueryState requestId -> STM Bool +isQueryTimedOut queryState = + case qsTimeoutVar queryState of + Nothing -> return False + Just timeoutVar -> readTVar timeoutVar + +-- | Data representing a batch ready for processing +data ReadyBatch requestId = ReadyBatch + { rbRequests :: !(Seq (NarinfoRequest requestId)), + rbAllPaths :: ![StorePath], + rbBatchStartTime :: !(Maybe UTCTime) + } + +-- | Wait for a batch to be ready for processing or shutdown +-- Returns Nothing if shutdown requested, Just batch if ready to process +waitForBatchOrShutdown :: + NarinfoQueryOptions -> + TVar (QueryState requestId) -> + STM (Maybe (ReadyBatch requestId)) +waitForBatchOrShutdown config stateVar = do + queryState <- readTVar stateVar + + -- Check for shutdown first + if not (qsRunning queryState) + then return Nothing + else do + -- Check if we have any pending requests + if Seq.null (qsPendingRequests queryState) + then retry -- No work, wait for requests + else do + -- We have requests, check if batch is ready + let pathCount = Set.size (qsAccumulatedPaths queryState) + sizeReady = pathCount >= nqoMaxBatchSize config + -- If timeout is 0, process immediately + immediateMode = nqoMaxWaitTime config <= 0 + + -- Check timeout condition + timeoutReady <- isQueryTimedOut queryState + + if sizeReady || timeoutReady || immediateMode + then do + -- Batch is ready, extract data and clear state + let requests = qsPendingRequests queryState + allPaths = Set.toList (qsAccumulatedPaths queryState) + startTime = qsQueryStartTime queryState + + -- Clear the batch state + writeTVar stateVar $ + queryState + { qsPendingRequests = Seq.empty, + qsAccumulatedPaths = Set.empty, + qsQueryStartTime = Nothing, + qsTimeoutVar = Nothing + } + + return $ + Just + ReadyBatch + { rbRequests = requests, + rbAllPaths = allPaths, + rbBatchStartTime = startTime + } + else retry -- Not ready yet, wait for timeout or more requests + +-- | Time refresh thread that updates cached time every few seconds +runTimeRefreshThread :: (MonadUnliftIO m) => NarinfoQueryManager requestId -> m () +runTimeRefreshThread NarinfoQueryManager {nqmCachedTime} = do + loop + where + loop = do + -- Update cached time every second + liftIO $ threadDelay 1_000_000 + + now <- liftIO getCurrentTime + liftIO $ atomically $ writeTVar nqmCachedTime now + loop + +-- | Main query processor loop +runQueryProcessorThread :: + (MonadUnliftIO m, Katip.KatipContext m) => + NarinfoQueryManager requestId -> + ([StorePath] -> m ([StorePath], [StorePath])) -> + m () +runQueryProcessorThread manager@NarinfoQueryManager {nqmConfig, nqmState} processBatch = do + loop + where + loop = do + -- Wait for a batch to be ready or shutdown + maybeReady <- liftIO $ atomically $ waitForBatchOrShutdown nqmConfig nqmState + + case maybeReady of + Nothing -> return () -- Shutdown requested + Just readyBatch -> do + -- Process the ready batch + processReadyBatch manager processBatch readyBatch + loop + +-- | Process a ready batch +processReadyBatch :: + (MonadUnliftIO m, Katip.KatipContext m) => + NarinfoQueryManager requestId -> + ([StorePath] -> m ([StorePath], [StorePath])) -> + ReadyBatch requestId -> + m () +processReadyBatch manager@NarinfoQueryManager {nqmCallback} processBatch ReadyBatch {rbRequests, rbAllPaths, rbBatchStartTime} = do + -- Process the batch if we have paths to query + (allPathsInClosure, missingPaths) <- + if null rbAllPaths + then return ([], []) + else do + processingStartTime <- liftIO getCurrentTime + + -- Log batch statistics + let requestCount = Seq.length rbRequests + pathCount = length rbAllPaths + waitTime = case rbBatchStartTime of + Nothing -> 0 + Just startTime -> processingStartTime `diffUTCTime` startTime + + Katip.logFM Katip.DebugS "Processing narinfo batch" + & Katip.katipAddContext + ( Katip.sl "requests" requestCount + <> Katip.sl "paths" pathCount + <> Katip.sl "wait_time_s" (show waitTime :: Text) + ) + + -- Query narinfo + (allPathsInClosure, missingPaths) <- processBatch rbAllPaths + + processingEndTime <- liftIO getCurrentTime + let processingTime = processingEndTime `diffUTCTime` processingStartTime + closureSize = length allPathsInClosure + missingCount = length missingPaths + + Katip.logFM Katip.DebugS "Batch completed" + & Katip.katipAddContext + ( Katip.sl "processing_time_s" (show processingTime :: Text) + <> Katip.sl "total_paths" closureSize + <> Katip.sl "missing_paths" missingCount + ) + + return (allPathsInClosure, missingPaths) + + -- Build lookup tables including cached results + let allPathsSet = Set.fromList allPathsInClosure + missingPathsSet = Set.fromList missingPaths + + -- Update cache with results (only cache positive lookups) + let existingPathsSet = allPathsSet `Set.difference` missingPathsSet + unless (Set.null existingPathsSet) $ do + Katip.logFM Katip.DebugS "Updating narinfo cache with existing paths" + & Katip.katipAddContext (Katip.sl "count" (Set.size existingPathsSet)) + updateCache manager existingPathsSet + + -- Respond to each request using the manager's callback + liftIO $ forM_ rbRequests $ \NarinfoRequest {qrRequestId, qrStorePaths} -> do + -- Find the missing paths for this requests from the batch + let qrStorePathsSet = Set.fromList qrStorePaths + missingPathsForRequest = qrStorePathsSet `Set.intersection` missingPathsSet + response = NarinfoResponse qrStorePathsSet missingPathsForRequest + + -- Call the manager's callback with request ID and response + nqmCallback qrRequestId response diff --git a/src/Cachix/Daemon/Progress.hs b/src/Cachix/Daemon/Progress.hs --- a/src/Cachix/Daemon/Progress.hs +++ b/src/Cachix/Daemon/Progress.hs @@ -1,17 +1,23 @@ module Cachix.Daemon.Progress ( UploadProgress, + ProgressState, new, complete, fail, tick, update, + newProgressState, + displayPushEvent, ) where import Cachix.Client.HumanSize (humanSize) import Cachix.Daemon.Types (PushRetryStatus (..)) +import Cachix.Daemon.Types.PushEvent (PushEvent (..), PushEventMessage (..)) import Control.Concurrent.Async qualified as Async import Control.Concurrent.MVar +import Data.HashMap.Strict qualified as HashMap +import Data.IORef import Data.String (String) import Data.Text qualified as T import Protolude @@ -32,6 +38,13 @@ size :: Int64 } +-- | State for tracking multiple progress bars +type ProgressState = IORef (HashMap.HashMap FilePath (UploadProgress, PushRetryStatus)) + +-- | Create a new progress state for tracking multiple progress bars +newProgressState :: IO ProgressState +newProgressState = newIORef HashMap.empty + new :: Handle -> String -> Int64 -> PushRetryStatus -> IO UploadProgress new hdl path size retryStatus = do isCI <- liftIO $ (== Just "true") <$> lookupEnv "CI" @@ -166,3 +179,28 @@ clearAsciiWith pg@(Ascii.ProgressBar _ _ region) str = do cancelUpdates pg Ascii.setConsoleRegion region str + +-- | Display push events as progress bars, deduplicating by path and retry status +displayPushEvent :: ProgressState -> PushEvent -> IO () +displayPushEvent statsRef PushEvent {eventMessage} = do + stats <- readIORef statsRef + case eventMessage of + PushStorePathAttempt path pathSize retryStatus -> do + case HashMap.lookup path stats of + Just (progress, prevRetryStatus) + | prevRetryStatus == retryStatus -> return () + | otherwise -> do + newProgress <- update progress retryStatus + writeIORef statsRef $ HashMap.insert path (newProgress, retryStatus) stats + Nothing -> do + progress <- new stderr (toS path) pathSize retryStatus + writeIORef statsRef $ HashMap.insert path (progress, retryStatus) stats + PushStorePathProgress path _ newBytes -> do + case HashMap.lookup path stats of + Nothing -> return () + Just (progress, _) -> tick progress newBytes + PushStorePathDone path -> do + mapM_ (complete . fst) (HashMap.lookup path stats) + PushStorePathFailed path _ -> do + mapM_ (fail . fst) (HashMap.lookup path stats) + _ -> return () diff --git a/src/Cachix/Daemon/Protocol.hs b/src/Cachix/Daemon/Protocol.hs --- a/src/Cachix/Daemon/Protocol.hs +++ b/src/Cachix/Daemon/Protocol.hs @@ -1,8 +1,7 @@ -{-# LANGUAGE GeneralizedNewtypeDeriving #-} - module Cachix.Daemon.Protocol ( ClientMessage (..), DaemonMessage (..), + DaemonErrorMessage (..), DaemonExitStatus (..), PushRequestId, newPushRequestId, @@ -12,11 +11,10 @@ ) where +import Cachix.Daemon.Types.PushEvent (PushEvent, PushRequestId, newPushRequestId) import Data.Aeson qualified as Aeson import Data.ByteString qualified as BS import Data.ByteString.Lazy qualified as LBS -import Data.UUID (UUID) -import Data.UUID.V4 qualified as UUID import Protolude -- | JSON messages that the client can send to the daemon @@ -31,9 +29,17 @@ data DaemonMessage = DaemonPong | DaemonExit !DaemonExitStatus + | DaemonPushEvent PushEvent + | DaemonError !DaemonErrorMessage deriving stock (Eq, Generic, Show) deriving anyclass (Aeson.FromJSON, Aeson.ToJSON) +-- | Error messages that the daemon can send to the client +data DaemonErrorMessage + = UnsupportedCommand !Text + deriving stock (Eq, Generic, Show) + deriving anyclass (Aeson.FromJSON, Aeson.ToJSON) + data DaemonExitStatus = DaemonExitStatus { exitCode :: !Int, exitMessage :: !(Maybe Text) @@ -41,16 +47,10 @@ deriving stock (Eq, Generic, Show) deriving anyclass (Aeson.FromJSON, Aeson.ToJSON) -newtype PushRequestId = PushRequestId UUID - deriving stock (Generic) - deriving newtype (Eq, Ord, Show, Aeson.FromJSON, Aeson.ToJSON, Hashable) - -newPushRequestId :: (MonadIO m) => m PushRequestId -newPushRequestId = liftIO $ PushRequestId <$> UUID.nextRandom - -- | A request for the daemon to push store paths to a binary cache data PushRequest = PushRequest - { storePaths :: [FilePath] + { storePaths :: [FilePath], + subscribeToUpdates :: Bool } deriving stock (Eq, Generic, Show) deriving anyclass (Aeson.FromJSON, Aeson.ToJSON) diff --git a/src/Cachix/Daemon/PushManager.hs b/src/Cachix/Daemon/PushManager.hs --- a/src/Cachix/Daemon/PushManager.hs +++ b/src/Cachix/Daemon/PushManager.hs @@ -34,6 +34,10 @@ pushStorePathDone, pushStorePathFailed, + -- * Batch processor + startBatchProcessor, + stopBatchProcessor, + -- * Helpers atomicallyWithTimeout, ) @@ -46,10 +50,11 @@ ) 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.Types.Log (Logger) -import Cachix.Daemon.Types.PushEvent +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 @@ -57,6 +62,7 @@ import Control.Concurrent.STM.TBMQueue 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 @@ -75,8 +81,8 @@ import Servant.Conduit () import UnliftIO.QSem qualified as QSem -newPushManagerEnv :: (MonadIO m) => PushOptions -> PushParams PushManager () -> OnPushEvent -> Logger -> m PushManagerEnv -newPushManagerEnv pushOptions pmPushParams onPushEvent pmLogger = liftIO $ do +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 @@ -85,6 +91,15 @@ pmLastEventTimestamp <- newTVarIO =<< getCurrentTime let pmOnPushEvent id pushEvent = updateTimestampTVar pmLastEventTimestamp >> onPushEvent id pushEvent + -- Create query manager with callback that queues ProcessQueryResponse tasks + let batchCallback requestId response = do + atomically $ do + result <- tryWriteTBMQueue pmTaskQueue $ HandleMissingPathsResponse requestId response + case result of + Just True -> return () + _ -> retry -- Queue is full, keep trying + pmNarinfoQueryManager <- NarinfoQuery.new batchOptions batchCallback + return $ PushManagerEnv {..} runPushManager :: (MonadIO m) => PushManagerEnv -> PushManager a -> m a @@ -97,6 +112,23 @@ check (pendingJobs <= 0) atomically $ closeTBMQueue 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 addPushJob :: Protocol.PushRequest -> PushManager (Maybe Protocol.PushRequestId) @@ -109,7 +141,7 @@ let queueJob = do modifyTVar' pmPushJobs $ HashMap.insert (pushId pushJob) pushJob incrementTVar pmPendingJobCount - res <- tryWriteTBMQueue pmTaskQueue $ ResolveClosure (pushId pushJob) + res <- tryWriteTBMQueue pmTaskQueue $ QueryMissingPaths (pushId pushJob) case res of Just True -> return True _ -> retry @@ -191,7 +223,7 @@ let addToQueue storePath = do isDuplicate <- HashMap.member storePath <$> readTVar pmStorePathIndex unless isDuplicate $ - writeTBMQueue pmTaskQueue (PushStorePath storePath) + writeTBMQueue pmTaskQueue (PushPath storePath) modifyTVar' pmStorePathIndex $ HashMap.insertWith (<>) storePath [pushId] @@ -264,13 +296,15 @@ handleTask task = do pushParams <- asks pmPushParams case task of - ResolveClosure pushId -> - runResolveClosureTask pushParams pushId - PushStorePath filePath -> - runPushStorePathTask pushParams filePath + QueryMissingPaths pushId -> + runQueryMissingPathsTask pushParams pushId + HandleMissingPathsResponse pushId batchResponse -> + runHandleMissingPathsResponseTask pushParams pushId batchResponse + PushPath filePath -> + runPushPathTask pushParams filePath -runResolveClosureTask :: PushParams PushManager () -> Protocol.PushRequestId -> PushManager () -runResolveClosureTask pushParams pushId = +runQueryMissingPathsTask :: PushParams PushManager () -> Protocol.PushRequestId -> PushManager () +runQueryMissingPathsTask pushParams pushId = resolveClosure `withException` failJob where failJob :: SomeException -> PushManager () @@ -288,23 +322,49 @@ withPushJob pushId $ \pushJob -> do let sps = Protocol.storePaths (pushRequest pushJob) store = pushParamsStore pushParams - normalized <- mapM (normalizeStorePath store) sps - (allStorePaths, missingStorePaths) <- getMissingPathsForClosure pushParams (catMaybes normalized) - storePathsToPush <- pushOnClosureAttempt pushParams allStorePaths missingStorePaths + validPaths <- catMaybes <$> mapM (normalizeStorePath store) sps - 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 - } + paths <- computeClosure store validPaths - resolvePushJob pushId resolvedClosure + -- Use async batch manager for narinfo queries (non-blocking) + batchManager <- asks pmNarinfoQueryManager + NarinfoQuery.submitRequest batchManager pushId paths -runPushStorePathTask :: PushParams PushManager () -> FilePath -> PushManager () -runPushStorePathTask pushParams filePath = do +runHandleMissingPathsResponseTask :: PushParams PushManager () -> Protocol.PushRequestId -> NarinfoQuery.NarinfoResponse -> PushManager () +runHandleMissingPathsResponseTask pushParams pushId batchResponse = + processQueryResponse `withException` failJob + where + failJob :: SomeException -> PushManager () + failJob err = do + failPushJob pushId + + 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 = @@ -343,7 +403,11 @@ onAttempt retryStatus size = do sp <- liftIO $ storePathToPath store storePath - Katip.katipAddContext (Katip.sl "retry" (rsIterNumber retryStatus)) $ + 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) @@ -396,8 +460,11 @@ pushFinished :: PushJob -> PushManager () pushFinished pushJob@PushJob {pushId} = void $ runMaybeT $ do - pushDuration <- MaybeT $ pure $ PushJob.duration pushJob - completedAt <- MaybeT $ pure $ PushJob.completedAt pushJob + 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 $ diff --git a/src/Cachix/Daemon/SocketStore.hs b/src/Cachix/Daemon/SocketStore.hs --- a/src/Cachix/Daemon/SocketStore.hs +++ b/src/Cachix/Daemon/SocketStore.hs @@ -4,15 +4,18 @@ removeSocket, toList, Socket (..), + sendAll, ) where import Cachix.Daemon.Types.SocketStore (Socket (..), SocketId, SocketStore (..)) import Control.Concurrent.STM.TVar import Control.Monad.IO.Unlift (MonadUnliftIO) +import Data.ByteString.Lazy qualified as LBS import Data.HashMap.Strict qualified as HashMap import Data.UUID.V4 qualified as UUID import Network.Socket qualified +import Network.Socket.ByteString.Lazy qualified as Socket.LBS import Protolude hiding (toList) import UnliftIO.Async qualified as Async @@ -41,3 +44,10 @@ toList (SocketStore st) = do hm <- liftIO $ readTVarIO st return $ HashMap.elems hm + +sendAll :: (MonadIO m) => SocketId -> LBS.ByteString -> SocketStore -> m () +sendAll socketId msg (SocketStore stvar) = do + mSocket <- liftIO $ readTVarIO stvar + case HashMap.lookup socketId mSocket of + Just (Socket {socket}) -> liftIO $ Socket.LBS.sendAll socket msg + Nothing -> return () diff --git a/src/Cachix/Daemon/TTLCache.hs b/src/Cachix/Daemon/TTLCache.hs new file mode 100644 --- /dev/null +++ b/src/Cachix/Daemon/TTLCache.hs @@ -0,0 +1,72 @@ +{-# LANGUAGE DerivingStrategies #-} +{-# LANGUAGE NamedFieldPuns #-} + +module Cachix.Daemon.TTLCache + ( -- * Types + TTLCache (..), + + -- * Operations + empty, + insert, + lookup, + cleanupExpired, + pruneToSize, + size, + ) +where + +import Data.HashMap.Strict qualified as HashMap +import Data.PQueue.Min qualified as PQ +import Data.Time (UTCTime) +import Protolude hiding (empty) + +-- | TTL cache with efficient expiration using priority queue +data TTLCache k = TTLCache + { -- | Fast lookups by key + tcLookupMap :: !(HashMap.HashMap k UTCTime), + -- | Min-heap ordered by expiration time for efficient cleanup + tcExpirationQueue :: !(PQ.MinQueue (UTCTime, k)) + } + deriving stock (Eq, Show) + +-- | Create an empty TTL cache O(1) +empty :: TTLCache k +empty = TTLCache HashMap.empty PQ.empty + +-- | Lookup a value in the TTL cache, checking expiration O(log n) +lookup :: (Ord k, Hashable k) => UTCTime -> k -> TTLCache k -> Bool +lookup now key TTLCache {tcLookupMap} = + case HashMap.lookup key tcLookupMap of + Nothing -> False + Just expiresAt -> now < expiresAt + +-- | Get the current size of the cache O(1) +size :: TTLCache k -> Int +size = PQ.size . tcExpirationQueue + +-- | Insert a value into the TTL cache with expiration time O(log n) +insert :: (Ord k, Hashable k) => k -> UTCTime -> TTLCache k -> TTLCache k +insert key expiresAt TTLCache {tcLookupMap, tcExpirationQueue} = + TTLCache + { tcLookupMap = HashMap.insert key expiresAt tcLookupMap, + tcExpirationQueue = PQ.insert (expiresAt, key) tcExpirationQueue + } + +-- | Remove expired entries from the cache +cleanupExpired :: (Ord k, Hashable k) => UTCTime -> TTLCache k -> TTLCache k +cleanupExpired now cache@TTLCache {tcLookupMap, tcExpirationQueue} = + let (expired, remaining) = PQ.span ((<= now) . fst) tcExpirationQueue + expiredKeys = map snd expired + newLookupMap = foldr HashMap.delete tcLookupMap expiredKeys + in cache {tcLookupMap = newLookupMap, tcExpirationQueue = remaining} + +-- | Prune cache to target size by removing oldest entries +pruneToSize :: (Ord k, Hashable k) => Int -> TTLCache k -> TTLCache k +pruneToSize targetSize cache@TTLCache {tcLookupMap, tcExpirationQueue} + | HashMap.size tcLookupMap <= targetSize = cache + | otherwise = + let excessCount = HashMap.size tcLookupMap - targetSize + (toRemove, remaining) = PQ.splitAt excessCount tcExpirationQueue + keysToRemove = map snd toRemove + newLookupMap = foldr HashMap.delete tcLookupMap keysToRemove + in cache {tcLookupMap = newLookupMap, tcExpirationQueue = remaining} diff --git a/src/Cachix/Daemon/Types/Daemon.hs b/src/Cachix/Daemon/Types/Daemon.hs --- a/src/Cachix/Daemon/Types/Daemon.hs +++ b/src/Cachix/Daemon/Types/Daemon.hs @@ -12,7 +12,7 @@ import Cachix.Client.Config.Orphans () import Cachix.Client.Env as Env -import Cachix.Client.OptionsParser (PushOptions) +import Cachix.Client.OptionsParser (DaemonOptions, PushOptions) import Cachix.Daemon.Log qualified as Log import Cachix.Daemon.Protocol qualified as Protocol import Cachix.Daemon.Subscription (SubscriptionManager) @@ -43,7 +43,7 @@ | -- | Remove an existing client socket connection. For example, after it is closed. RemoveSocketClient SocketId | -- | Handle a new message from a client. - ReceivedMessage Protocol.ClientMessage + ReceivedMessage SocketId Protocol.ClientMessage data DaemonEnv = DaemonEnv { -- | Cachix client env @@ -73,7 +73,9 @@ -- | Logging env daemonLogger :: Logger, -- | The PID of the daemon process - daemonPid :: ProcessID + daemonPid :: ProcessID, + -- | Daemon configuration options + daemonOptions :: DaemonOptions } newtype Daemon a = Daemon diff --git a/src/Cachix/Daemon/Types/Error.hs b/src/Cachix/Daemon/Types/Error.hs --- a/src/Cachix/Daemon/Types/Error.hs +++ b/src/Cachix/Daemon/Types/Error.hs @@ -20,6 +20,8 @@ DaemonSocketError | -- | Failed to push some store paths DaemonPushFailure + | -- | A command is not supported by the daemon configuration + DaemonUnsupportedCommand Text deriving stock (Show, Eq) instance Exception DaemonError where @@ -31,6 +33,7 @@ unlines ["The daemon encountered an internal event loop error:", "", toS (displayException eventLoopError)] DaemonSocketError -> "There was an error with the socket" DaemonPushFailure -> "Failed to push some store paths" + DaemonUnsupportedCommand msg -> toS msg -- | A wrapper for SomeException that implements equality. newtype UnhandledException @@ -54,6 +57,7 @@ instance HasExitCode DaemonError where toExitCode err = ExitFailure $ case err of DaemonPushFailure -> 3 + DaemonUnsupportedCommand _ -> 2 _remainingErrors -> 1 instance (HasExitCode a) => HasExitCode (Maybe a) where diff --git a/src/Cachix/Daemon/Types/PushEvent.hs b/src/Cachix/Daemon/Types/PushEvent.hs --- a/src/Cachix/Daemon/Types/PushEvent.hs +++ b/src/Cachix/Daemon/Types/PushEvent.hs @@ -1,21 +1,24 @@ module Cachix.Daemon.Types.PushEvent ( PushEvent (..), PushEventMessage (..), + PushRequestId, PushRetryStatus (..), + newPushRequestId, newPushRetryStatus, ) where -import Cachix.Daemon.Protocol qualified as Protocol import Control.Retry (RetryStatus (..)) import Data.Aeson (FromJSON, ToJSON) import Data.Time (UTCTime) +import Data.UUID (UUID) +import Data.UUID.V4 qualified as UUID import Protolude data PushEvent = PushEvent { -- TODO: newtype a monotonic clock eventTimestamp :: UTCTime, - eventPushId :: Protocol.PushRequestId, + eventPushId :: PushRequestId, eventMessage :: PushEventMessage } deriving stock (Eq, Generic, Show) @@ -33,6 +36,13 @@ | PushFinished deriving stock (Eq, Generic, Show) deriving anyclass (FromJSON, ToJSON) + +newtype PushRequestId = PushRequestId UUID + deriving stock (Generic) + deriving newtype (Eq, Ord, Show, FromJSON, ToJSON, Hashable) + +newPushRequestId :: (MonadIO m) => m PushRequestId +newPushRequestId = liftIO $ PushRequestId <$> UUID.nextRandom data PushRetryStatus = PushRetryStatus {retryCount :: Int} deriving stock (Eq, Generic, Show) diff --git a/src/Cachix/Daemon/Types/PushManager.hs b/src/Cachix/Daemon/Types/PushManager.hs --- a/src/Cachix/Daemon/Types/PushManager.hs +++ b/src/Cachix/Daemon/Types/PushManager.hs @@ -18,6 +18,8 @@ import Cachix.Client.Push (PushParams) import Cachix.Daemon.Log qualified as Log +import Cachix.Daemon.NarinfoQuery (NarinfoQueryManager) +import Cachix.Daemon.NarinfoQuery qualified as NarinfoQuery import Cachix.Daemon.Protocol qualified as Protocol import Cachix.Daemon.Types.Log (Logger) import Cachix.Daemon.Types.PushEvent (PushEvent (..)) @@ -31,8 +33,9 @@ import Protolude data Task - = ResolveClosure Protocol.PushRequestId - | PushStorePath FilePath + = QueryMissingPaths Protocol.PushRequestId + | HandleMissingPathsResponse Protocol.PushRequestId NarinfoQuery.NarinfoResponse + | PushPath FilePath type PushJobStore = TVar (HashMap Protocol.PushRequestId PushJob) @@ -58,6 +61,8 @@ pmLastEventTimestamp :: TVar UTCTime, -- | The number of pending (uncompleted) jobs. pmPendingJobCount :: TVar Int, + -- | Manager for batching narinfo queries + pmNarinfoQueryManager :: NarinfoQueryManager Protocol.PushRequestId, pmLogger :: Logger } diff --git a/src/Cachix/Deploy/ActivateCommand.hs b/src/Cachix/Deploy/ActivateCommand.hs --- a/src/Cachix/Deploy/ActivateCommand.hs +++ b/src/Cachix/Deploy/ActivateCommand.hs @@ -153,7 +153,7 @@ Text.intercalate "\n" $ "Deploying agents:" : [ inBrackets agentName <> " " <> DeployResponse.url details - | (agentName, details) <- agents + | (agentName, details) <- agents ] renderSummary :: [(Text, Deployment.Deployment)] -> Text @@ -161,7 +161,7 @@ Text.intercalate "\n" $ "Deployment summary:" : [ inBrackets agentName <> " " <> renderStatus (Deployment.status deployment) - | (agentName, deployment) <- results + | (agentName, deployment) <- results ] where renderStatus = \case diff --git a/test/Daemon/NarinfoQuerySpec.hs b/test/Daemon/NarinfoQuerySpec.hs new file mode 100644 --- /dev/null +++ b/test/Daemon/NarinfoQuerySpec.hs @@ -0,0 +1,295 @@ +{-# 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 + +-- 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 = do + 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) + +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 first request (1 path) - should not trigger + NarinfoQuery.submitRequest tcManager (1 :: Int) [path1] + threadDelay 10000 + + calls1 <- readTVarIO tcBatchCalls + length calls1 `shouldBe` 0 -- No batch yet + + -- Submit second request (1 more unique path) - should trigger batch (2 paths total) + NarinfoQuery.submitRequest tcManager (2 :: Int) [path2] + threadDelay 20000 -- Wait for batch processing + calls2 <- readTVarIO tcBatchCalls + length calls2 `shouldBe` 1 -- One batch triggered + let batchPaths = case head calls2 of + Just (BatchCall paths _) -> paths + Nothing -> panic "Expected batch call" + Set.fromList batchPaths `shouldBe` Set.fromList [path1, path2] + + -- Verify both requests got responses + callbacks <- readTVarIO tcCallbackCalls + length callbacks `shouldBe` 2 + + 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 to trigger + threadDelay 60000 -- 60ms > 50ms timeout + calls <- readTVarIO tcBatchCalls + length calls `shouldBe` 1 -- Batch triggered by timeout + callbacks <- readTVarIO tcCallbackCalls + length callbacks `shouldBe` 1 -- Request got response + 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] + threadDelay 20000 -- Short wait + calls <- readTVarIO tcBatchCalls + length calls `shouldBe` 1 -- Processed immediately + callbacks <- readTVarIO tcCallbackCalls + length callbacks `shouldBe` 1 + + 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] + threadDelay 20000 + + -- 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] + threadDelay 20000 + + -- Second request - path1 from cache, path2 goes to batch + NarinfoQuery.submitRequest tcManager (2 :: Int) [path1, path2] + threadDelay 20000 + + -- 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] + + -- Both requests should have gotten responses + callbacks <- readTVarIO tcCallbackCalls + length callbacks `shouldBe` 2 + + -- Second response should contain both paths (path1 from cache + path2 from batch) + 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 + let config = defaultNarinfoQueryOptions {nqoMaxWaitTime = 0} + 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 + NarinfoQuery.submitRequest tcManager (1 :: Int) [path1, path2, path3] + -- Request 2: paths 3,4,5 (path 3 overlaps) + NarinfoQuery.submitRequest tcManager (2 :: Int) [path3, path4, path5] + + threadDelay 20_000 + + callbacks <- readTVarIO tcCallbackCalls + length callbacks `shouldBe` 2 + + -- 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) + threadDelay 120000 -- Wait 120ms, longer than 100ms timeout + 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 + + -- Both requests should still get responses + callbacks <- readTVarIO tcCallbackCalls + length callbacks `shouldBe` 2 diff --git a/test/Daemon/PushManagerSpec.hs b/test/Daemon/PushManagerSpec.hs --- a/test/Daemon/PushManagerSpec.hs +++ b/test/Daemon/PushManagerSpec.hs @@ -4,6 +4,7 @@ import Cachix.Client.OptionsParser (defaultPushOptions) import Cachix.Client.Push (PushSecret (PushToken)) import Cachix.Daemon.Log qualified as Log +import Cachix.Daemon.NarinfoQuery (defaultNarinfoQueryOptions) import Cachix.Daemon.Protocol qualified as Protocol import Cachix.Daemon.Push qualified as Daemon.Push import Cachix.Daemon.PushManager @@ -30,12 +31,12 @@ spec = do describe "push job" $ do it "starts in the queued state" $ do - let request = Protocol.PushRequest {Protocol.storePaths = ["foo", "bar"]} + let request = Protocol.PushRequest {Protocol.storePaths = ["foo", "bar"], Protocol.subscribeToUpdates = False} pushJob <- PushJob.new request PushJob.status pushJob `shouldBe` Queued it "can be resolved" $ do - let request = Protocol.PushRequest {Protocol.storePaths = ["foo", "bar"]} + let request = Protocol.PushRequest {Protocol.storePaths = ["foo", "bar"], Protocol.subscribeToUpdates = False} pathSet = Set.fromList ["foo", "bar"] closure = PushJob.ResolvedClosure pathSet pathSet initPushJob <- PushJob.new request @@ -47,7 +48,7 @@ it "marks paths as pushed" $ do - let request = Protocol.PushRequest {Protocol.storePaths = ["foo", "bar"]} + let request = Protocol.PushRequest {Protocol.storePaths = ["foo", "bar"], Protocol.subscribeToUpdates = False} pathSet = Set.fromList ["foo", "bar"] closure = PushJob.ResolvedClosure pathSet pathSet timestamp <- getCurrentTime @@ -69,7 +70,7 @@ it "marks paths as failed" $ do - let request = Protocol.PushRequest {Protocol.storePaths = ["foo", "bar"]} + let request = Protocol.PushRequest {Protocol.storePaths = ["foo", "bar"], Protocol.subscribeToUpdates = False} pathSet = Set.fromList ["foo", "bar"] closure = PushJob.ResolvedClosure pathSet pathSet @@ -90,7 +91,7 @@ it "unmark paths as failed after successful retry" $ do - let request = Protocol.PushRequest {Protocol.storePaths = ["foo", "bar"]} + let request = Protocol.PushRequest {Protocol.storePaths = ["foo", "bar"], Protocol.subscribeToUpdates = False} pathSet = Set.fromList ["foo", "bar"] closure = PushJob.ResolvedClosure pathSet pathSet @@ -111,7 +112,7 @@ describe "push manager" $ do it "queues push jobs " $ inPushManager $ do - let request = Protocol.PushRequest {Protocol.storePaths = ["foo", "bar"]} + let request = Protocol.PushRequest {Protocol.storePaths = ["foo", "bar"], Protocol.subscribeToUpdates = False} Just pushId <- addPushJob request Just pushJob <- lookupPushJob pushId liftIO $ do @@ -121,7 +122,7 @@ it "manages the lifecycle of a push job" $ inPushManager $ do let paths = ["bar", "foo"] - let pushRequest = Protocol.PushRequest {Protocol.storePaths = paths} + let pushRequest = Protocol.PushRequest {Protocol.storePaths = paths, Protocol.subscribeToUpdates = False} Just pushId <- addPushJob pushRequest let pathSet = Set.fromList paths @@ -156,7 +157,7 @@ let longTimeoutOptions = TimeoutOptions {toTimeout = 10.0, toPollingInterval = 0.1} _ <- runPushManager pm $ do - let request = Protocol.PushRequest {Protocol.storePaths = paths} + let request = Protocol.PushRequest {Protocol.storePaths = paths, Protocol.subscribeToUpdates = False} addPushJob request concurrently_ (stopPushManager longTimeoutOptions pm) $ @@ -166,7 +167,7 @@ it "shuts down on job stall" $ withPushManager $ \pm -> do _ <- runPushManager pm $ do - let request = Protocol.PushRequest {Protocol.storePaths = ["foo"]} + let request = Protocol.PushRequest {Protocol.storePaths = ["foo"], Protocol.subscribeToUpdates = False} addPushJob request stopPushManager timeoutOptions pm @@ -187,8 +188,9 @@ let binaryCache = newBinaryCache "test" pushSecret = PushToken (Token "test") pushOptions = defaultPushOptions + batchOptions = defaultNarinfoQueryOptions pushParams = Daemon.Push.newPushParams store clientEnv binaryCache pushSecret pushOptions - newPushManagerEnv pushOptions pushParams mempty logger >>= f + newPushManagerEnv pushOptions batchOptions pushParams mempty logger >>= f inPushManager :: PushManager a -> IO a inPushManager f = withPushManager (`runPushManager` f) diff --git a/test/Daemon/TTLCacheSpec.hs b/test/Daemon/TTLCacheSpec.hs new file mode 100644 --- /dev/null +++ b/test/Daemon/TTLCacheSpec.hs @@ -0,0 +1,37 @@ +module Daemon.TTLCacheSpec where + +import Cachix.Daemon.TTLCache (TTLCache) +import Cachix.Daemon.TTLCache qualified as TTLCache +import Data.Time (addUTCTime, getCurrentTime) +import Protolude +import Test.Hspec + +spec :: Spec +spec = do + describe "TTL cache" $ do + it "correctly handles expiration" $ do + now <- getCurrentTime + let past = addUTCTime (-10) now -- 10 seconds ago + future = addUTCTime 10 now -- 10 seconds from now + cache :: TTLCache Text + cache = + TTLCache.insert "expired" past $ + TTLCache.insert "valid" future TTLCache.empty + + -- Expired entry should not be found + TTLCache.lookup now "expired" cache `shouldBe` False + -- Valid entry should be found + TTLCache.lookup now "valid" cache `shouldBe` True + + it "respects size limits when pruning" $ do + now <- getCurrentTime + let future = addUTCTime 60 now + -- Create cache with 5 entries + cache :: TTLCache Text + cache = foldl' (\c i -> TTLCache.insert (show i) future c) TTLCache.empty [1 .. 5 :: Integer] + + TTLCache.size cache `shouldBe` 5 + + -- Prune to 3 entries + let pruned = TTLCache.pruneToSize 3 cache + TTLCache.size pruned `shouldBe` 3