cachix 1.8.0 → 1.9.0
raw patch · 23 files changed
+1305/−177 lines, 23 filesdep +pqueue
Dependencies added: pqueue
Files
- cachix.cabal +8/−1
- src/Cachix/Client.hs +5/−4
- src/Cachix/Client/Command/Push.hs +0/−1
- src/Cachix/Client/Command/Watch.hs +16/−36
- src/Cachix/Client/NixVersion.hs +1/−1
- src/Cachix/Client/OptionsParser.hs +95/−6
- src/Cachix/Daemon.hs +49/−18
- src/Cachix/Daemon/Client.hs +110/−50
- src/Cachix/Daemon/Listen.hs +1/−1
- src/Cachix/Daemon/NarinfoQuery.hs +419/−0
- src/Cachix/Daemon/Progress.hs +38/−0
- src/Cachix/Daemon/Protocol.hs +12/−12
- src/Cachix/Daemon/PushManager.hs +95/−28
- src/Cachix/Daemon/SocketStore.hs +10/−0
- src/Cachix/Daemon/TTLCache.hs +72/−0
- src/Cachix/Daemon/Types/Daemon.hs +5/−3
- src/Cachix/Daemon/Types/Error.hs +4/−0
- src/Cachix/Daemon/Types/PushEvent.hs +12/−2
- src/Cachix/Daemon/Types/PushManager.hs +7/−2
- src/Cachix/Deploy/ActivateCommand.hs +2/−2
- test/Daemon/NarinfoQuerySpec.hs +295/−0
- test/Daemon/PushManagerSpec.hs +12/−10
- test/Daemon/TTLCacheSpec.hs +37/−0
cachix.cabal view
@@ -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
src/Cachix/Client.hs view
@@ -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.
src/Cachix/Client/Command/Push.hs view
@@ -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)
src/Cachix/Client/Command/Watch.hs view
@@ -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
src/Cachix/Client/NixVersion.hs view
@@ -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
src/Cachix/Client/OptionsParser.hs view
@@ -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"))
src/Cachix/Daemon.hs view
@@ -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
src/Cachix/Daemon/Client.hs view
@@ -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
src/Cachix/Daemon/Listen.hs view
@@ -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
+ src/Cachix/Daemon/NarinfoQuery.hs view
@@ -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
src/Cachix/Daemon/Progress.hs view
@@ -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 ()
src/Cachix/Daemon/Protocol.hs view
@@ -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)
src/Cachix/Daemon/PushManager.hs view
@@ -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 $
src/Cachix/Daemon/SocketStore.hs view
@@ -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 ()
+ src/Cachix/Daemon/TTLCache.hs view
@@ -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}
src/Cachix/Daemon/Types/Daemon.hs view
@@ -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
src/Cachix/Daemon/Types/Error.hs view
@@ -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
src/Cachix/Daemon/Types/PushEvent.hs view
@@ -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)
src/Cachix/Daemon/Types/PushManager.hs view
@@ -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 }
src/Cachix/Deploy/ActivateCommand.hs view
@@ -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
+ test/Daemon/NarinfoQuerySpec.hs view
@@ -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
test/Daemon/PushManagerSpec.hs view
@@ -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)
+ test/Daemon/TTLCacheSpec.hs view
@@ -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