diff --git a/benchmarks/DownloadCountBench.hs b/benchmarks/DownloadCountBench.hs new file mode 100644 index 000000000..c3714ad1a --- /dev/null +++ b/benchmarks/DownloadCountBench.hs @@ -0,0 +1,218 @@ +-- | Measure the cost of recording package downloads. +-- +-- Downloads used to be written to the acid-state log one event per download. +-- They are now accumulated in memory and written one event per flush (see +-- "Distribution.Server.Features.DownloadCount"). Both paths still exist in +-- the code -- 'RegisterDownload' is retained so that old event logs can still +-- be replayed -- so we can measure the old and new behaviour from a single +-- binary, without checking out an older tree. +-- +-- Usage: +-- +-- > DownloadCountBench MODE [DOWNLOADS] [PACKAGES] [PER-FLUSH] +-- +-- where MODE is +-- +-- * @individual@: one acid-state event per download, plus the per-download +-- 'RecordedToday' query the old code did. This is the pre-batching path. +-- It does not include the 'Control.Concurrent.Chan' write and thread +-- wake-up the old code also paid for each download, so the numbers here +-- understate the old cost slightly. +-- +-- * @batched@: accumulate into a 'TVar' exactly as the download hook now +-- does, and write one 'RegisterDownloads' event per PER-FLUSH downloads. +-- At 50 downloads/second, the server's 60 second flush interval +-- corresponds to PER-FLUSH = 3000. +-- +-- Run each mode in its own process, since the memory figures are process-wide: +-- +-- > cabal run --enable-benchmarks DownloadCountBench -- individual +-- > cabal run --enable-benchmarks DownloadCountBench -- batched +module Main where + +import Control.Concurrent.STM (atomically, modifyTVar', newTVarIO, swapTVar) +import Control.Monad (forM_, unless, when) +import Data.Acid (AcidState, closeAcidState, openLocalStateFrom, query, update) +import qualified Data.Map.Strict as Map +import Data.Time.Calendar (fromGregorian) +import Distribution.Package (PackageId, PackageIdentifier (..), mkPackageName) +import Distribution.Server.Features.DownloadCount.State +import Distribution.Server.Framework.MemSize (memSize, memSizeKb) +import Distribution.Server.Util.CountingMap (cmTotal) +import Distribution.Version (mkVersion) +import GHC.Clock (getMonotonicTimeNSec) +import GHC.Stats (RTSStats (..), getRTSStats, getRTSStatsEnabled) +import System.Directory + ( doesDirectoryExist, + getFileSize, + getTemporaryDirectory, + listDirectory, + removePathForcibly, + ) +import System.Environment (getArgs, getProgName) +import System.Exit (die) +import System.FilePath (()) +import System.Mem (performMajorGC) +import Text.Printf (printf) + +data Mode = Individual | Batched + deriving (Eq) + +main :: IO () +main = do + args <- getArgs + (mode, downloads, packages, perFlush) <- case args of + (m : rest) -> do + mode <- case m of + "individual" -> pure Individual + "batched" -> pure Batched + _ -> badUsage + let arg n def = case drop n rest of + (x : _) -> read x + [] -> def + pure + ( mode, + arg 0 200000, -- total num downloads + arg 1 3000, -- number of packages + arg 2 200 -- batch size + ) + _ -> badUsage + + statsEnabled <- getRTSStatsEnabled + unless statsEnabled $ + die "Run with +RTS -T (the benchmark is built with -with-rtsopts=-T)" + + tmp <- getTemporaryDirectory + let stateDir = tmp "hackage-download-count-bench" + removePathForcibly stateDir + + printf + "%s: %d downloads over %d distinct package versions\n" + (modeName mode) + (downloads :: Int) + (packages :: Int) + when (mode == Batched) $ + printf " flushing every %d downloads\n" (perFlush :: Int) + + st <- openLocalStateFrom stateDir (initInMemStats today) + let pkgids = downloadStream downloads packages + + -- The work being measured. + elapsed <- timed $ case mode of + Individual -> recordIndividually st pkgids + Batched -> recordBatched st perFlush pkgids + + -- Check both modes actually recorded everything, so that we cannot + -- accidentally compare a fast path that does less work. + final <- query st GetInMemStats + let recorded = cmTotal (inMemCounts final) + unless (recorded == downloads) $ + die $ + "Recorded " + ++ show recorded + ++ " downloads, expected " + ++ show downloads + + performMajorGC + stats <- getRTSStats + logBytes <- dirSize stateDir + + closeAcidState st + + -- Reopening replays every event written since the last checkpoint, so + -- this is what the event log costs the server at start-up. + replay <- timed $ do + st' <- openLocalStateFrom stateDir (initInMemStats today) + _ <- query st' GetInMemStats + closeAcidState st' + + removePathForcibly stateDir + + printf " time to record %10.3f s\n" (seconds elapsed) + printf " event log on disk %10d KB\n" (logBytes `div` 1024) + printf " start-up replay %10.3f s\n" (seconds replay) + printf + " allocated %10d MB\n" + (allocated_bytes stats `div` (1024 * 1024)) + printf + " peak live heap %10d KB\n" + (max_live_bytes stats `div` 1024) + printf " major GCs %10d\n" (major_gcs stats) + printf " InMemStats memSize %10d KB\n" (memSizeKb (memSize final)) + where + today = fromGregorian 2026 8 24 + + badUsage = do + pname <- getProgName + die $ + "usage: " + ++ pname + ++ " (individual|batched) [DOWNLOADS] [PACKAGES] [PER-FLUSH]" + + modeName Individual = "individual (one event per download)" + modeName Batched = "batched (one event per flush)" + +-- | The stream of downloads to record. +-- +-- Spread uniformly over the given number of distinct package versions, which +-- is the least favourable case for batching: real download traffic is heavily +-- skewed towards a few packages, which collapses further within a flush. +downloadStream :: Int -> Int -> [PackageId] +downloadStream downloads packages = + [pkgid ((i * 7919) `mod` packages) | i <- [1 .. downloads]] + where + pkgid n = + PackageIdentifier + (mkPackageName ("package-" ++ show n)) + (mkVersion [1, n `mod` 20]) + +-- | The pre-batching path: one 'RecordedToday' query and one +-- 'RegisterDownload' event per download. +recordIndividually :: AcidState InMemStats -> [PackageId] -> IO () +recordIndividually st pkgids = + forM_ pkgids $ \pkgid -> do + _ <- query st RecordedToday + update st (RegisterDownload pkgid) + +-- | The current path: accumulate in a 'TVar', write one event per flush. +recordBatched :: AcidState InMemStats -> Int -> [PackageId] -> IO () +recordBatched st perFlush pkgids = do + acc <- newTVarIO Map.empty + let go _ [] = flush acc + go n (pkgid : rest) = do + atomically $ modifyTVar' acc (Map.insertWith (+) pkgid 1) + if n >= perFlush + then flush acc >> go 1 rest + else go (n + 1 :: Int) rest + go 1 pkgids + where + flush acc = do + counts <- atomically $ swapTVar acc Map.empty + unless (Map.null counts) $ + update st (RegisterDownloads (Map.toList counts)) + +-- | Time an action, in nanoseconds. +timed :: IO () -> IO Word +timed action = do + before <- getMonotonicTimeNSec + action + after <- getMonotonicTimeNSec + pure (fromIntegral (after - before)) + +seconds :: Word -> Double +seconds ns = fromIntegral ns / 1e9 + +-- | Total size of every file under a directory. +dirSize :: FilePath -> IO Integer +dirSize dir = do + isDir <- doesDirectoryExist dir + if not isDir + then pure 0 + else do + entries <- listDirectory dir + sum <$> mapM entrySize entries + where + entrySize entry = do + let path = dir entry + isDir <- doesDirectoryExist path + if isDir then dirSize path else getFileSize path diff --git a/exes/Main.hs b/exes/Main.hs index 2e4723f11..7782b708b 100644 --- a/exes/Main.hs +++ b/exes/Main.hs @@ -5,6 +5,7 @@ module Main where +import Distribution.ReadE (ReadE (ReadE)) import qualified Distribution.Server as Server import Distribution.Server (ListenOn(..), ServerConfig(..), Server) import Distribution.Server.Framework.Feature @@ -58,6 +59,7 @@ import Control.Arrow import qualified Data.ByteString.Lazy as BS import qualified Distribution.Server.Util.GZip as GZip import qualified Text.Parsec as Parse +import Text.Read (readMaybe) import Paths_hackage_server as Paths (version) @@ -206,6 +208,9 @@ data RunFlags = RunFlags { flagRunTemp :: Flag Bool, flagRunCacheDelay :: Flag String, flagRunLiveTemplates :: Flag Bool, + -- | The period before which to flush download counts to + -- disk, in seconds + flagRunFlushDownloads :: Flag Int, -- Online backup flags flagRunBackupOutputDir :: Flag FilePath, flagRunBackupLinkBlobs :: Flag Bool, @@ -226,6 +231,7 @@ defaultRunFlags = RunFlags { flagRunTemp = Flag False, flagRunCacheDelay = NoFlag, flagRunLiveTemplates = Flag False, + flagRunFlushDownloads = Flag 60, flagRunBackupOutputDir = Flag "backups", flagRunBackupLinkBlobs = Flag False, flagRunBackupScrubbed = Flag False @@ -311,7 +317,16 @@ runCommand = "Do not cache templates, for quicker feedback during development." flagRunLiveTemplates (\v flags -> flags { flagRunLiveTemplates = v }) (noArg (Flag True)) + , option [] ["flush-downloads-period"] + "Period on which download counts are flushed to disk, in seconds." + flagRunFlushDownloads (\v flags -> flags {flagRunFlushDownloads = v}) + (optArg "SECONDS" parseFlushDownloadPeriod ("60", Flag 60) (\(Flag n) -> [Just (show n)])) ] + where + parseFlushDownloadPeriod :: ReadE (Flag Int) + parseFlushDownloadPeriod = ReadE $ \s -> case readMaybe s of + Nothing -> Left "Could not parse flush download period" + Just seconds -> Right (Flag seconds) runAction :: RunFlags -> IO () runAction opts = do @@ -330,6 +345,7 @@ runAction opts = do loPortNum = port, loIP = ip } + flushDownloadPeriod = fromFlagOrDefault (confDownloadFlush defaults) $ flagRunFlushDownloads opts config = defaults { confHostUri = hosturi, confUserContentUri = usercontenturi, @@ -340,7 +356,8 @@ runAction opts = do confTmpDir = tmpDir, confCacheDelay = cacheDelay, confLiveTemplates = liveTemplates, - confVerbosity = verbosity + confVerbosity = verbosity, + confDownloadFlush = flushDownloadPeriod } outputDir = fromFlag (flagRunBackupOutputDir opts) linkBlobs = fromFlag (flagRunBackupLinkBlobs opts) diff --git a/hackage-server.cabal b/hackage-server.cabal index ac0972b60..f96b18ad4 100644 --- a/hackage-server.cabal +++ b/hackage-server.cabal @@ -702,6 +702,20 @@ benchmark RevDeps ghc-options: -with-rtsopts=-s other-modules: RevDepCommon +benchmark DownloadCountBench + import: test-defaults + type: exitcode-stdio-1.0 + hs-source-dirs: benchmarks + main-is: DownloadCountBench.hs + build-tool-depends: hackage-server:hackage-server + -- NOTE: version constraints inherited from hackage-server + -- We are not using 'gauge' as it is not supported for + -- GHC 9.10+ + build-depends: + , acid-state + , stm + ghc-options: -with-rtsopts=-T + test-suite PaginationTest import: test-defaults type: exitcode-stdio-1.0 @@ -717,6 +731,14 @@ test-suite BrowseQueryParserTest other-modules: Util build-depends: attoparsec +test-suite DownloadCountTest + import: test-defaults + type: exitcode-stdio-1.0 + main-is: DownloadCountTest.hs + build-tool-depends: hackage-server:hackage-server + -- NOTE: version constraints inherited from hackage-server + build-depends: acid-state + test-suite CreateUserTest import: test-defaults diff --git a/src/Distribution/Server.hs b/src/Distribution/Server.hs index b5a5d97f5..37d36d428 100644 --- a/src/Distribution/Server.hs +++ b/src/Distribution/Server.hs @@ -1,4 +1,5 @@ {-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE NumericUnderscores #-} module Distribution.Server ( -- * Server control Server(..), @@ -75,6 +76,8 @@ data ServerConfig = ServerConfig { confStaticDir :: FilePath, confTmpDir :: FilePath, confCacheDelay:: Int, + -- | Period before which downloads are flushed to disk + confDownloadFlush :: Int, confLiveTemplates :: Bool } deriving (Show) @@ -108,6 +111,7 @@ defaultServerConfig = do confStaticDir = dataDir, confTmpDir = "state" "tmp", confCacheDelay= 0, + confDownloadFlush = 60, confLiveTemplates = False } @@ -128,7 +132,7 @@ hasSavedState = doesDirectoryExist . confDbStateDir mkServerEnv :: ServerConfig -> IO ServerEnv mkServerEnv config@(ServerConfig verbosity hostURI userContentURI requiredBaseHostHeader _ stateDir _ tmpDir - cacheDelay liveTemplates) = do + cacheDelay flushDownloadPeriod liveTemplates) = do createDirectoryIfMissing False stateDir let blobStoreDir = confBlobStoreDir config staticDir = confStaticFilesDir config @@ -149,7 +153,8 @@ mkServerEnv config@(ServerConfig verbosity hostURI userContentURI requiredBaseHo serverBlobStore = store, serverCron = cron, serverTmpDir = tmpDir, - serverCacheDelay = cacheDelay * 1000000, --microseconds + serverCacheDelay = cacheDelay * 1_000_000, --microseconds + serverFlushDownloadsPeriod = flushDownloadPeriod * 1_000_000,--microseconds serverBaseURI = hostURI, serverUserContentBaseURI = userContentURI, serverRequiredBaseHostHeader = requiredBaseHostHeader, @@ -197,6 +202,8 @@ run server@Server{ serverEnv = env } = do cronJobAction = checkpoint server } + mapM_ Feature.featureRunHook (serverFeatures server) + runServer listenOn $ do handlePutPostQuotas diff --git a/src/Distribution/Server/Features.hs b/src/Distribution/Server/Features.hs index c9a3b1dbf..c8aa49a16 100644 --- a/src/Distribution/Server/Features.hs +++ b/src/Distribution/Server/Features.hs @@ -451,9 +451,12 @@ featureCheckpoint = mapM_ abstractStateCheckpoint . featureState checkpointAllFeatures :: [HackageFeature] -> IO () checkpointAllFeatures = mapM_ featureCheckpoint --- | Cleanly shut down a feature's state components. +-- | Cleanly shut down a feature: run its pre-shutdown hook, then close its +-- state components. featureShutdown :: HackageFeature -> IO () -featureShutdown = mapM_ abstractStateClose . featureState +featureShutdown feature = do + featurePreShutdown feature + mapM_ abstractStateClose (featureState feature) -- | Cleanly shut down all features' state components. shutdownAllFeatures :: [HackageFeature] -> IO () diff --git a/src/Distribution/Server/Features/DownloadCount.hs b/src/Distribution/Server/Features/DownloadCount.hs index 93dd607cb..a861e4c91 100644 --- a/src/Distribution/Server/Features/DownloadCount.hs +++ b/src/Distribution/Server/Features/DownloadCount.hs @@ -2,11 +2,18 @@ {-# LANGUAGE DerivingStrategies #-} {-# LANGUAGE DeriveGeneric #-} {-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE NumericUnderscores #-} {-# OPTIONS_GHC -Wno-orphans #-} -- | Download counts -- -- We maintain -- +-- 0. In-memory (cache): downloads that have arrived since the last flush. +-- Downloads are accumulated here and flushed into (1) as a single +-- transaction once per flush interval, so that a busy server does not write +-- one acid-state event per download. A hard crash loses at most one flush +-- interval worth of counts. +-- -- 1. In-memory (ACID): today's download counts per package version -- -- 2. In-memory (cache): total download count over the last 30 days per package @@ -43,8 +50,11 @@ import Distribution.Server.Util.CountingMap (cmFromCSV, cmToList) import Data.Time.Calendar (Day, addDays) import Data.Time.Clock (getCurrentTime, utctDay) -import Control.Concurrent.Chan -import Control.Concurrent (forkIO) +import Control.Concurrent (forkIO, threadDelay) +import Control.Concurrent.MVar (MVar, newMVar, withMVar) +import Control.Concurrent.STM (TVar, atomically, modifyTVar', newTVarIO, swapTVar) +import Control.Exception (SomeException, try) +import qualified Data.Map.Strict as Map import GHC.Generics (Generic) import Data.Aeson (ToJSON) import qualified Data.Aeson as Aeson @@ -82,13 +92,16 @@ initDownloadFeature serverEnv@ServerEnv{serverStateDir} = do totalDownloads) <- computeRecentAndTotalDownloads =<< getState onDiskState recentCache <- newMemStateWHNF recentDownloads totalsCache <- newMemStateWHNF totalDownloads - downChan <- newChan + pendingDownloads <- newTVarIO Map.empty + flushLock <- newMVar () return $ \core users -> do let feature = downloadFeature core users serverEnv inMemBackend - onDiskState totalsCache recentCache downChan + onDiskState totalsCache recentCache + pendingDownloads flushLock - registerHook (packageDownloadHook core) (writeChan downChan) + registerHook (packageDownloadHook core) $ \pkgid -> + atomically $ modifyTVar' pendingDownloads (Map.insertWith (+) pkgid 1) return feature onDiskStateComponent :: FilePath -> StateComponent OnDiskState OnDiskStats @@ -113,17 +126,19 @@ downloadFeature :: CoreFeature -> StateComponent OnDiskState OnDiskStats -> MemState TotalDownloads -> MemState RecentDownloads - -> Chan PackageId + -> TVar (Map.Map PackageId Int) + -> MVar () -> DownloadFeature downloadFeature CoreFeature{} UserFeature{..} - ServerEnv{serverStateDir} + ServerEnv{serverStateDir, serverVerbosity, serverFlushDownloadsPeriod} inMemBackend onDiskState totalDownloadsCache recentDownloadsCache - downloadStream + pendingDownloads + flushLock = DownloadFeature{..} where inMemStore = Store.backendStore inMemBackend @@ -132,9 +147,10 @@ downloadFeature CoreFeature{} featureResources = [ topDownloads downloadResource , downloadCSV ] - , featurePostInit = void $ forkIO registerDownloads + , featureRunHook = void $ forkIO (flushDownloadsLoop serverFlushDownloadsPeriod) , featureState = Store.backendState inMemBackend ++ [abstractOnDiskStateComponent onDiskState] + , featurePreShutdown = shutdownFlush , featureCaches = [ CacheComponent { cacheDesc = "recent package downloads cache", @@ -153,14 +169,48 @@ downloadFeature CoreFeature{} totalPackageDownloads :: MonadIO m => m TotalDownloads totalPackageDownloads = readMemState totalDownloadsCache - registerDownloads = forever $ do - pkg <- readChan downloadStream + flushDownloadsLoop :: Int -> IO () + flushDownloadsLoop flushInterval = forever $ do + threadDelay flushInterval + flushDownloadsSafe + + logErrors :: String -> IO () -> IO () + logErrors what action = do + outcome <- try action + case outcome of + Right () -> return () + Left err -> lognotice serverVerbosity $ + what ++ ": " ++ show (err :: SomeException) + + flushDownloadsSafe :: IO () + flushDownloadsSafe = + logErrors "Error recording download counts" flushDownloads + + shutdownFlush :: IO () + shutdownFlush = do + flushDownloadsSafe + logErrors "Error checkpointing download counts" $ + Store.checkpointInMemStats inMemStore + + flushDownloads :: IO () + flushDownloads = withMVar flushLock $ \() -> do + checkDayRollover + + counts <- atomically $ swapTVar pendingDownloads Map.empty + unless (Map.null counts) $ + Store.registerDownloads inMemStore (Map.toList counts) + + checkDayRollover :: IO () + checkDayRollover = do today <- getToday today' <- Store.recordedToday inMemStore - --TODO: do this asyncronously rather than blocking this request + --TODO: this should be a daily cron job rather than being polled by the + -- flush loop: the rollover does a lot of I/O (it rewrites the whole + -- on-disk history) and it holds up the flush of the counts while it + -- runs. when (today /= today') $ do - -- For the first download each day we reset the in-memory stats and.. + -- For the first flush each day we reset the in-memory stats and.. inMemStats <- Store.getInMemStats inMemStore Store.replaceInMemStats inMemStore $ initInMemStats today -- we can discard the large eventlog by writing a small checkpoint @@ -169,21 +219,20 @@ downloadFeature CoreFeature{} -- Write yesterday's downloads to the log appendToLog (dcPath serverStateDir) inMemStats - -- Update the on-disk statistics and recompute recent downloads - onDiskStats' <- updateHistory inMemStats <$> getState onDiskState - writeOnDiskStats (dcPath serverStateDir "ondisk") onDiskStats' - --TODO: this is still stupid, writing it out only to read it back - -- we should be able to update the in memory ones incrementally + -- Update the on-disk statistics and recompute recent downloads. + -- Only the packages downloaded yesterday need writing out. + (onDiskStats', + changedPkgs) <- updateHistory inMemStats <$> getState onDiskState + writeOnDiskStatsFor (dcPath serverStateDir "ondisk") + changedPkgs onDiskStats' + --TODO: we still recompute these from the whole history rather than + -- updating them with yesterday's downloads (recentDownloads, - totalDownloads) <- computeRecentAndTotalDownloads =<< getState onDiskState + totalDownloads) <- computeRecentAndTotalDownloads onDiskStats' writeMemState recentDownloadsCache recentDownloads writeMemState totalDownloadsCache totalDownloads - Store.registerDownload inMemStore pkg - - - downloadResource = DownloadResource { topDownloads = (resourceAt "/packages/top.:format") { resourceDesc = [ (GET, "Get top downloaded packages for the last 30 days")] diff --git a/src/Distribution/Server/Features/DownloadCount/Acid.hs b/src/Distribution/Server/Features/DownloadCount/Acid.hs index b300b86db..78c5b648e 100644 --- a/src/Distribution/Server/Features/DownloadCount/Acid.hs +++ b/src/Distribution/Server/Features/DownloadCount/Acid.hs @@ -17,7 +17,7 @@ acidStore stateDir = do Store.recordedToday = queryState inMemState State.RecordedToday , Store.getInMemStats = queryState inMemState State.GetInMemStats , Store.replaceInMemStats = \stats -> updateState inMemState (State.ReplaceInMemStats stats) - , Store.registerDownload = \pkgid -> updateState inMemState (State.RegisterDownload pkgid) + , Store.registerDownloads = \pkgs -> updateState inMemState (State.RegisterDownloads pkgs) , Store.checkpointInMemStats = liftIO (createCheckpoint (stateHandle inMemState)) } , Store.backendState = [abstractAcidStateComponent inMemState] diff --git a/src/Distribution/Server/Features/DownloadCount/State.hs b/src/Distribution/Server/Features/DownloadCount/State.hs index cdd591885..c128fb5d4 100644 --- a/src/Distribution/Server/Features/DownloadCount/State.hs +++ b/src/Distribution/Server/Features/DownloadCount/State.hs @@ -6,7 +6,7 @@ module Distribution.Server.Features.DownloadCount.State where import Data.Time.Calendar (Day(..)) import Data.Foldable (forM_) import Control.Arrow (first) -import Control.Monad (liftM) +import Control.Monad (liftM, unless) import Data.List (foldl', groupBy) import Data.Function (on) import Control.Monad.Reader (ask, asks) @@ -168,9 +168,15 @@ accumTotalDownloads pkgName (OnDiskPerPkg perPkg) = Pure updates/queries ------------------------------------------------------------------------------} -updateHistory :: InMemStats -> OnDiskStats -> OnDiskStats +-- | Fold a day's downloads into the historical statistics. +-- +-- The name of packages whose statistics changed is also returned, +-- so that only those need writing out. +updateHistory :: InMemStats -> OnDiskStats -> (OnDiskStats, [PackageName]) updateHistory (InMemStats day perPkg) (OnDiskStats (NCM _ m)) = - OnDiskStats (NCM 0 (Map.unionWith cmUnion m updatesMap)) + ( OnDiskStats (NCM 0 (Map.unionWith cmUnion m updatesMap)) + , Map.keys updatesMap + ) where updatesMap :: Map.Map PackageName OnDiskPerPkg updatesMap = Map.fromList @@ -220,11 +226,16 @@ readOnDiskPerPkg pkgFile = evaluate =<< (runGetLazy safeGet <$> BSL.hGetContents h) writeOnDiskStats :: FilePath -> OnDiskStats -> IO () -writeOnDiskStats stateDir (OnDiskStats (NCM _ onDisk)) = do +writeOnDiskStats stateDir stats@(OnDiskStats (NCM _ onDisk)) = + writeOnDiskStatsFor stateDir (Map.keys onDisk) stats + +-- | Write out the statistics for the given packages only. +writeOnDiskStatsFor :: FilePath -> [PackageName] -> OnDiskStats -> IO () +writeOnDiskStatsFor stateDir pkgNames (OnDiskStats (NCM _ onDisk)) = do createDirectoryIfMissing True stateDir - forM_ (Map.toList onDisk) $ \(pkgName, perPkg) -> do - let pkgFile = stateDir display pkgName - writeFileAtomic pkgFile $ runPutLazy (safePut perPkg) + forM_ pkgNames $ \pkgName -> + forM_ (Map.lookup pkgName onDisk) $ \perPkg -> + writeFileAtomic (stateDir display pkgName) $ runPutLazy (safePut perPkg) {------------------------------------------------------------------------------ The append-only all-time log @@ -251,13 +262,27 @@ replaceInMemStats = put recordedToday :: Query InMemStats Day recordedToday = asks inMemToday +-- | Record a single download. +-- +-- Superseded by 'registerDownloads', which coalesces many downloads into a +-- single event. This is retained only so that existing acid-state event logs +-- (which may contain 'RegisterDownload' events) can still be replayed. +-- +-- This should be dropped at some point registerDownload :: PackageId -> Update InMemStats () -registerDownload pkgId = do +registerDownload pkgId = registerDownloads [(pkgId, 1)] + +-- | Record a batch of downloads: @(package, number of downloads)@ pairs. +registerDownloads :: [(PackageId, Int)] -> Update InMemStats () +registerDownloads pkgs = unless (null pkgs) $ do InMemStats day counts <- get - put $ InMemStats day (cmInsert pkgId 1 counts) + put $! InMemStats day (foldl' insert counts pkgs) + where + insert counts (pkgId, n) = cmInsert pkgId n counts makeAcidic ''InMemStats [ 'getInMemStats , 'replaceInMemStats , 'recordedToday , 'registerDownload + , 'registerDownloads ] diff --git a/src/Distribution/Server/Features/DownloadCount/Store.hs b/src/Distribution/Server/Features/DownloadCount/Store.hs index e1cd90952..94b655b2e 100644 --- a/src/Distribution/Server/Features/DownloadCount/Store.hs +++ b/src/Distribution/Server/Features/DownloadCount/Store.hs @@ -22,6 +22,7 @@ data Store = Store { recordedToday :: forall m. MonadIO m => m Day , getInMemStats :: forall m. MonadIO m => m State.InMemStats , replaceInMemStats :: forall m. MonadIO m => State.InMemStats -> m () - , registerDownload :: forall m. MonadIO m => PackageId -> m () + -- | Record a batch of downloads: @(package, number of downloads)@ pairs. + , registerDownloads :: forall m. MonadIO m => [(PackageId, Int)] -> m () , checkpointInMemStats :: forall m. MonadIO m => m () } diff --git a/src/Distribution/Server/Framework/Feature.hs b/src/Distribution/Server/Framework/Feature.hs index 357d9865c..0e36b26c5 100644 --- a/src/Distribution/Server/Framework/Feature.hs +++ b/src/Distribution/Server/Framework/Feature.hs @@ -47,6 +47,12 @@ data HackageFeature = HackageFeature { , featureErrHandlers :: [(String, ServerErrorResponse)] , featurePostInit :: IO () + -- | Run just before the server starts accepting HTTP connections, which is + -- useful to start background worker threads. + , featureRunHook :: IO () + -- | Run when the server is shutting down, before the feature's state + -- components are closed. Must not throw. + , featurePreShutdown :: IO () , featureReloadFiles :: IO () , featureState :: [AbstractStateComponent] @@ -69,6 +75,8 @@ emptyHackageFeature name = HackageFeature { featureErrHandlers= [], featurePostInit = return (), + featureRunHook = return (), + featurePreShutdown = return (), featureReloadFiles = return (), featureState = error $ "'featureState' not defined for feature '" ++ name ++ "'", diff --git a/src/Distribution/Server/Framework/ServerEnv.hs b/src/Distribution/Server/Framework/ServerEnv.hs index 9e4cfdeef..5beeb8b42 100644 --- a/src/Distribution/Server/Framework/ServerEnv.hs +++ b/src/Distribution/Server/Framework/ServerEnv.hs @@ -73,6 +73,11 @@ data ServerEnv = ServerEnv { -- increasing the time taken to update the cache we can push this further. serverCacheDelay :: Int, + -- | A tunable parameter for the download counts in-memory accumulation. + -- Download counts are accumulated in-memory and flushed to disk every 'serverFlushDownloadPeriods' + -- microseconds. + serverFlushDownloadsPeriod :: Int, + serverVerbosity :: Verbosity } diff --git a/tests/DownloadCountTest.hs b/tests/DownloadCountTest.hs new file mode 100644 index 000000000..f0c978dbc --- /dev/null +++ b/tests/DownloadCountTest.hs @@ -0,0 +1,144 @@ +module Main where + +import Control.Monad (forM_, unless) +import Data.Acid (AcidState, query, update) +import Data.Acid.Memory (openMemoryState) +import qualified Data.Map.Strict as Map +import Data.Time.Calendar (fromGregorian) +import Distribution.Package (PackageIdentifier (..), mkPackageName) +import Distribution.Server.Features.DownloadCount.State +import Distribution.Server.Util.CountingMap (cmEmpty, cmFind, cmInsert, cmToList, cmTotal) +import Distribution.Version (mkVersion) +import System.Directory (getTemporaryDirectory, removePathForcibly) +import System.Exit (die) +import System.FilePath (()) + +main :: IO () +main = do + -- Recording downloads one at a time and recording them as a batch must + -- produce exactly the same statistics. + oneAtATime <- withStats $ \st -> + forM_ downloadStream $ \pkgid -> update st (RegisterDownload pkgid) + batched <- withStats $ \st -> + update st (RegisterDownloads (tally downloadStream)) + unless (oneAtATime == batched) $ + die $ + "Batched downloads do not match individual downloads:\n" + ++ show (cmToList (inMemCounts oneAtATime)) + ++ "\nversus\n" + ++ show (cmToList (inMemCounts batched)) + + -- Sanity check the counts themselves, so that the test above cannot be + -- satisfied by both paths being equally wrong. + let counts = inMemCounts batched + unless (cmTotal counts == length downloadStream) $ + die $ + "Wrong total: " + ++ show (cmTotal counts) + ++ " expected " + ++ show (length downloadStream) + forM_ (tally downloadStream) $ \(pkgid, n) -> + unless (cmFind pkgid counts == n) $ + die $ + "Wrong count for " + ++ show pkgid + ++ ": " + ++ show (cmFind pkgid counts) + ++ " expected " + ++ show n + + -- Splitting a batch across several updates must accumulate, not replace. + split <- withStats $ \st -> do + update st (RegisterDownloads (tally (take 3 downloadStream))) + update st (RegisterDownloads (tally (drop 3 downloadStream))) + unless (split == batched) $ + die "Successive batches do not accumulate" + + -- An empty batch must be a no-op, and in particular must not disturb the + -- day the statistics are recorded against. + empty <- withStats $ \st -> update st (RegisterDownloads []) + unless (inMemCounts empty == inMemCounts (initInMemStats today)) $ + die "Empty batch changed the counts" + unless (inMemToday batched == today) $ + die "Recording downloads changed the recorded day" + + -- Rolling a day over writes out only the packages that were downloaded + -- that day. The result on disk must be the same as writing every package. + checkIncrementalWrite + + putStrLn "OK" + where + withStats :: (AcidState InMemStats -> IO ()) -> IO InMemStats + withStats action = do + st <- openMemoryState (initInMemStats today) + action st + query st GetInMemStats + + today = fromGregorian 2026 8 24 + + -- Deliberately interleaved and repeated, so that ordering differences + -- between the two paths would show up. + downloadStream = + [ pkg "foo" [1, 0], + pkg "bar" [2], + pkg "foo" [1, 0], + pkg "foo" [2, 1], + pkg "bar" [2], + pkg "foo" [1, 0] + ] + + pkg name version = + PackageIdentifier (mkPackageName name) (mkVersion version) + +tally :: [PackageIdentifier] -> [(PackageIdentifier, Int)] +tally = Map.toList . Map.fromListWith (+) . map (,1) + +checkIncrementalWrite :: IO () +checkIncrementalWrite = do + tmp <- getTemporaryDirectory + let dir = tmp "hackage-download-count-test" + removePathForcibly dir + + let yesterday = fromGregorian 2026 8 23 + history = + fst $ + updateHistory + ( statsFor + yesterday + [ pkg "foo" [1, 0], + pkg "bar" [2] + ] + ) + cmEmpty + writeOnDiskStats dir history + + let today = fromGregorian 2026 8 24 + (history', changed) = + updateHistory + ( statsFor + today + [ pkg "foo" [1, 0], + pkg "baz" [3] + ] + ) + history + expectedChanged = [mkPackageName "baz", mkPackageName "foo"] + unless (changed == expectedChanged) $ + die $ + "Expected " ++ show expectedChanged ++ " to change, got " ++ show changed + writeOnDiskStatsFor dir changed history' + + reread <- readOnDiskStats dir + unless (cmToList reread == cmToList history') $ + die $ + "On-disk statistics do not match after an incremental write:\n" + ++ show (cmToList reread) + ++ "\nversus\n" + ++ show (cmToList history') + removePathForcibly dir + where + statsFor day pkgids = + InMemStats day (foldr (`cmInsert` 1) cmEmpty pkgids) + + pkg name version = + PackageIdentifier (mkPackageName name) (mkVersion version)