diff --git a/benchmarks/DownloadCountBench.hs b/benchmarks/DownloadCountBench.hs new file mode 100644 index 000000000..c6b82a6ab --- /dev/null +++ b/benchmarks/DownloadCountBench.hs @@ -0,0 +1,204 @@ +-- | Measure the cost of recording package downloads. +-- +-- Usage: +-- +-- > DownloadCountBench MODE [DOWNLOADS] [PACKAGES] [PER-FLUSH] +-- +-- where MODE is +-- +-- * @individual@: one acid-state event per download +-- +-- * @batched@: accumulate before flushing +-- +-- 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..c1c3fb133 100644 --- a/exes/Main.hs +++ b/exes/Main.hs @@ -206,6 +206,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 +229,7 @@ defaultRunFlags = RunFlags { flagRunTemp = Flag False, flagRunCacheDelay = NoFlag, flagRunLiveTemplates = Flag False, + flagRunFlushDownloads = Flag 60, flagRunBackupOutputDir = Flag "backups", flagRunBackupLinkBlobs = Flag False, flagRunBackupScrubbed = Flag False @@ -311,6 +315,10 @@ 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}) + (noArg (Flag 60)) ] runAction :: RunFlags -> IO () diff --git a/hackage-server.cabal b/hackage-server.cabal index bd88d1455..7d5f20537 100644 --- a/hackage-server.cabal +++ b/hackage-server.cabal @@ -660,6 +660,19 @@ 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 + -- 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 @@ -675,6 +688,13 @@ 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 + build-depends: acid-state + test-suite CreateUserTest import: test-defaults diff --git a/src/Distribution/Server.hs b/src/Distribution/Server.hs index b5a5d97f5..3d03584d5 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, diff --git a/src/Distribution/Server/Features.hs b/src/Distribution/Server/Features.hs index c9a3b1dbf..f597f29df 100644 --- a/src/Distribution/Server/Features.hs +++ b/src/Distribution/Server/Features.hs @@ -451,9 +451,10 @@ featureCheckpoint = mapM_ abstractStateCheckpoint . featureState checkpointAllFeatures :: [HackageFeature] -> IO () checkpointAllFeatures = mapM_ featureCheckpoint --- | Cleanly shut down a feature's 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 527ac859d..8ce744098 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 @@ -41,8 +48,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 @@ -80,13 +90,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 inMemState - 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 inMemStateComponent :: FilePath -> IO (StateComponent AcidState InMemStats) @@ -125,24 +138,27 @@ 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} inMemState onDiskState totalDownloadsCache recentDownloadsCache - downloadStream + pendingDownloads + flushLock = DownloadFeature{..} where downloadFeatureInterface = (emptyHackageFeature "download") { featureResources = [ topDownloads downloadResource , downloadCSV ] - , featurePostInit = void $ forkIO registerDownloads + , featurePostInit = void $ forkIO (flushDownloadsLoop serverFlushDownloadsPeriod) + , featurePreShutdown = shutdownFlush , featureState = [ abstractAcidStateComponent inMemState , abstractOnDiskStateComponent onDiskState ] @@ -164,14 +180,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" $ + createCheckpoint (stateHandle inMemState) + + flushDownloads :: IO () + flushDownloads = withMVar flushLock $ \() -> do + checkDayRollover + + counts <- atomically $ swapTVar pendingDownloads Map.empty + unless (Map.null counts) $ + updateState inMemState $ RegisterDownloads (Map.toList counts) + + checkDayRollover :: IO () + checkDayRollover = do today <- getToday - today' <- query (stateHandle inMemState) RecordedToday + today' <- queryState inMemState RecordedToday - --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 <- getState inMemState putState inMemState $ initInMemStats today -- we can discard the large eventlog by writing a small checkpoint @@ -180,20 +230,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 - updateState inMemState $ RegisterDownload 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/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/Framework/Feature.hs b/src/Distribution/Server/Framework/Feature.hs index 357d9865c..d840ab062 100644 --- a/src/Distribution/Server/Framework/Feature.hs +++ b/src/Distribution/Server/Framework/Feature.hs @@ -47,6 +47,9 @@ data HackageFeature = HackageFeature { , featureErrHandlers :: [(String, ServerErrorResponse)] , featurePostInit :: 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 +72,7 @@ emptyHackageFeature name = HackageFeature { featureErrHandlers= [], featurePostInit = 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)