State.hs

{-# LANGUAGE TemplateHaskell, StandaloneDeriving, GeneralizedNewtypeDeriving,
             DeriveDataTypeable, TypeFamilies, FlexibleInstances,
             MultiParamTypeClasses, BangPatterns #-}
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 Data.List (foldl', groupBy)
import Data.Function (on)
import Control.Monad.Reader (ask, asks)
import Control.Monad.State (get, put)
import qualified Data.Map.Lazy as Map
import System.FilePath ((</>))
import System.Directory (
    getDirectoryContents
  , createDirectoryIfMissing
  )
import qualified Data.ByteString.Lazy as BSL
import System.IO (withFile, IOMode (..))
import System.IO.Unsafe (unsafeInterleaveIO)
import Text.CSV (printCSV)
import Control.Exception (evaluate)

import Data.Acid (Update, Query, makeAcidic)
import Data.SafeCopy (base, deriveSafeCopy, safeGet, safePut)
import Data.Serialize.Get (runGetLazy)
import Data.Serialize.Put (runPutLazy)

import Distribution.Version (Version)
import Distribution.Package (
    PackageId
  , PackageName
  , packageName
  , packageVersion
  )
import Distribution.Text (simpleParse, display)
import Distribution.Simple.Utils (writeFileAtomic)

import Distribution.Server.Framework.Instances ()
import Distribution.Server.Framework.MemSize
import Distribution.Server.Util.CountingMap

{------------------------------------------------------------------------------
  Data types
------------------------------------------------------------------------------}

data InMemStats = InMemStats {
    inMemToday  :: !Day
  , inMemCounts :: !(SimpleCountingMap PackageId)
  }
  deriving (Show, Eq)

newtype OnDiskStats = OnDiskStats {
    onDiskStats :: NestedCountingMap PackageName OnDiskPerPkg
  }
  deriving (Show, Eq, MemSize)

instance CountingMap (PackageName, (Day, Version)) OnDiskStats where
  cmEmpty                            = OnDiskStats cmEmpty
  cmTotal  (OnDiskStats ncm)         = cmTotal ncm
  cmInsert kl n (OnDiskStats ncm)    = OnDiskStats $ cmInsert kl n ncm
  cmFind   k (OnDiskStats ncm)       = cmFind k ncm
  cmUnion    (OnDiskStats a)
             (OnDiskStats b)         = OnDiskStats (cmUnion a b)
  cmToList   (OnDiskStats ncm)       = cmToList ncm
  cmToCSV    (OnDiskStats ncm)       = cmToCSV ncm
  cmInsertRecord r (OnDiskStats ncm) = first OnDiskStats `liftM` cmInsertRecord r ncm

newtype OnDiskPerPkg = OnDiskPerPkg {
    onDiskPerPkgCounts :: NestedCountingMap Day (SimpleCountingMap Version)
  }
  deriving (Show, Eq, Ord, MemSize)

instance CountingMap (Day, Version) OnDiskPerPkg where
  cmEmpty  = OnDiskPerPkg cmEmpty
  cmTotal  (OnDiskPerPkg ncm) = cmTotal ncm
  cmInsert kl n (OnDiskPerPkg ncm) = OnDiskPerPkg $ cmInsert kl n ncm
  cmFind   k (OnDiskPerPkg ncm) = cmFind k ncm
  cmUnion  (OnDiskPerPkg a) (OnDiskPerPkg b) = OnDiskPerPkg (cmUnion a b)
  cmToList (OnDiskPerPkg ncm) = cmToList ncm
  cmToCSV  (OnDiskPerPkg ncm) = cmToCSV ncm
  cmInsertRecord r (OnDiskPerPkg ncm) = first OnDiskPerPkg `liftM` cmInsertRecord r ncm

newtype RecentDownloads = RecentDownloads {
    recentDownloads :: SimpleCountingMap PackageName
  }
  deriving (Show, Eq, MemSize)

instance CountingMap PackageName RecentDownloads where
  cmEmpty  = RecentDownloads cmEmpty
  cmTotal  (RecentDownloads ncm) = cmTotal ncm
  cmInsert kl n (RecentDownloads ncm) = RecentDownloads $ cmInsert kl n ncm
  cmFind   k (RecentDownloads ncm) = cmFind k ncm
  cmUnion  (RecentDownloads a) (RecentDownloads b) = RecentDownloads (cmUnion a b)
  cmToList (RecentDownloads ncm) = cmToList ncm
  cmToCSV  (RecentDownloads ncm) = cmToCSV ncm
  cmInsertRecord r (RecentDownloads ncm) = first RecentDownloads `liftM` cmInsertRecord r ncm

newtype TotalDownloads = TotalDownloads {
    totalDownloads :: SimpleCountingMap PackageName
  }
  deriving (Show, Eq, MemSize)

instance CountingMap PackageName TotalDownloads where
  cmEmpty  = TotalDownloads cmEmpty
  cmTotal  (TotalDownloads ncm) = cmTotal ncm
  cmInsert kl n (TotalDownloads ncm) = TotalDownloads $ cmInsert kl n ncm
  cmFind   k (TotalDownloads ncm) = cmFind k ncm
  cmUnion  (TotalDownloads a) (TotalDownloads b) = TotalDownloads (cmUnion a b)
  cmToList (TotalDownloads ncm) = cmToList ncm
  cmToCSV  (TotalDownloads ncm) = cmToCSV ncm
  cmInsertRecord r (TotalDownloads ncm) = first TotalDownloads `liftM` cmInsertRecord r ncm

{------------------------------------------------------------------------------
  Initial instances
------------------------------------------------------------------------------}

initInMemStats :: Day -> InMemStats
initInMemStats day = InMemStats {
    inMemToday  = day
  , inMemCounts = cmEmpty
  }

type DayRange = (Day, Day)

initRecentAndTotalDownloads :: DayRange -> OnDiskStats
                            -> (RecentDownloads, TotalDownloads)
initRecentAndTotalDownloads dayRange (OnDiskStats (NCM _ m)) =
    foldl' (\(recent, total) (pname, pstats) ->
              let !recent' = accumRecentDownloads dayRange pname pstats recent
                  !total'  = accumTotalDownloads  pname pstats total
               in (recent', total'))
           (emptyRecentDownloads, emptyTotalDownloads)
           (Map.toList m)

emptyRecentDownloads :: RecentDownloads
emptyRecentDownloads = RecentDownloads cmEmpty

accumRecentDownloads :: DayRange
                     -> PackageName -> OnDiskPerPkg
                     -> RecentDownloads -> RecentDownloads
accumRecentDownloads dayRange pkgName (OnDiskPerPkg (NCM _ perDay))
  | let rangeTotal = sum (map cmTotal (lookupRange dayRange perDay))
  , rangeTotal > 0
  = cmInsert pkgName rangeTotal

  | otherwise = id

lookupRange :: Ord k => (k,k) -> Map.Map k a -> [a]
lookupRange (l,u) m =
  let (_,ml,above)  = Map.splitLookup l m
      (middle,mu,_) = Map.splitLookup u above
   in maybe [] (\x->[x]) ml
   ++ Map.elems middle
   ++ maybe [] (\x->[x]) mu

emptyTotalDownloads :: TotalDownloads
emptyTotalDownloads = TotalDownloads cmEmpty

accumTotalDownloads :: PackageName -> OnDiskPerPkg
                    -> TotalDownloads -> TotalDownloads
accumTotalDownloads pkgName (OnDiskPerPkg perPkg) =
    cmInsert pkgName (cmTotal perPkg)

{------------------------------------------------------------------------------
  Pure updates/queries
------------------------------------------------------------------------------}

updateHistory :: InMemStats -> OnDiskStats -> OnDiskStats
updateHistory (InMemStats day perPkg) (OnDiskStats (NCM _ m)) =
    OnDiskStats (NCM 0 (Map.unionWith cmUnion m updatesMap))
  where
    updatesMap :: Map.Map PackageName OnDiskPerPkg
    updatesMap = Map.fromList
      [ (pkgname, applyUpdates pkgs)
      | pkgs <- groupBy ((==) `on` (packageName . fst))
                        (cmToList perPkg :: [(PackageId, Int)])
      , let pkgname = packageName (fst (head pkgs))
      ]

    applyUpdates :: [(PackageId, Int)] -> OnDiskPerPkg
    applyUpdates pkgs = foldr (.) id
                          [ cmInsert (day, packageVersion pkgId) count
                          | (pkgId, count) <- pkgs ]
                          cmEmpty

{------------------------------------------------------------------------------
  MemSize
------------------------------------------------------------------------------}

instance MemSize InMemStats where
  memSize (InMemStats a b) = memSize2 a b

{------------------------------------------------------------------------------
  Serializing on-disk stats
------------------------------------------------------------------------------}

deriveSafeCopy 0 'base ''InMemStats
deriveSafeCopy 0 'base ''OnDiskPerPkg

-- | This processes all of the files. If you're only interested in a
-- couple of them, use 'readOnDiskStatsLazily' instead, to avoid most
-- of the IO.
readOnDiskStatsEagerly :: FilePath -> IO OnDiskStats
readOnDiskStatsEagerly = readOnDiskStatsHelper False

-- | Compared to 'readOnDiskStatsEagerly', this defers opening and
-- processing the files; each file will be opened and processed when
-- the corresponding entry in the result map is forced.
--
-- It is therefore perhaps unwise to call this function and then
-- quickly force every entry of the map. It's not easy to predict
-- whether it will open many files before closing any of them; might
-- depend on the behaviors of the threaded RTS\/the platform\/etc.
--
-- So if you're going to quickly force many of the packages' entries
-- in the map, you could instead call 'readOnDiskStatsEagerly' to
-- avoid any risk of exhausing file descriptions: every file will be
-- closed before the next is opened.
--
-- But we haven't /actually demonstrated/ 'readOnDiskStatsLazily'
-- risks exhausting file descriptors; we merely just suspect it
-- might. See the FdExhaustion-DownloadCount-Hypothesis
-- benchmark. Thus, for now, the rest of the codebase continues to use
-- only 'readOnDiskStateLazily', since that's what it used before we
-- started investigating.
readOnDiskStatsLazily :: FilePath -> IO OnDiskStats
readOnDiskStatsLazily = readOnDiskStatsHelper True

-- | See 'readOnDiskStatsLazily' and 'readOnDiskStatsEagerly'.
readOnDiskStatsHelper :: Bool -> FilePath -> IO OnDiskStats
readOnDiskStatsHelper whetherToInterleaveIO stateDir = do
    createDirectoryIfMissing True stateDir
    pkgStrs <- getDirectoryContents stateDir
    OnDiskStats . NCM 0 . Map.fromList <$> sequence
      [ do onDiskPerPkg <- (if whetherToInterleaveIO then unsafeInterleaveIO else id) $
                             either (const cmEmpty) id
                               <$> readOnDiskPerPkg pkgFile
           return (pkgName, onDiskPerPkg)
      | Just pkgName <- map simpleParse pkgStrs
      , let pkgFile = stateDir </> display pkgName ]

readOnDiskPerPkg :: FilePath -> IO (Either String OnDiskPerPkg)
readOnDiskPerPkg pkgFile =
    withFile pkgFile ReadMode $ \h ->
      -- By evaluating the Either result from the parser we force
      -- all contents to be read
      evaluate =<< (runGetLazy safeGet <$> BSL.hGetContents h)

writeOnDiskStats :: FilePath -> OnDiskStats -> IO ()
writeOnDiskStats stateDir (OnDiskStats (NCM _ onDisk)) = do
   createDirectoryIfMissing True stateDir
   forM_ (Map.toList onDisk) $ \(pkgName, perPkg) -> do
     let pkgFile = stateDir </> display pkgName
     writeFileAtomic pkgFile $ runPutLazy (safePut perPkg)

{------------------------------------------------------------------------------
  The append-only all-time log
------------------------------------------------------------------------------}

appendToLog :: FilePath -> InMemStats -> IO ()
appendToLog stateDir (InMemStats _ inMemStats) =
  appendFile (stateDir </> "log") $ printCSV (cmToCSV inMemStats)

reconstructLog :: FilePath -> OnDiskStats -> IO ()
reconstructLog stateDir onDisk =
  writeFile (stateDir </> "log") $ printCSV (cmToCSV onDisk)

{------------------------------------------------------------------------------
  ACID stuff
------------------------------------------------------------------------------}

getInMemStats :: Query InMemStats InMemStats
getInMemStats = ask

replaceInMemStats :: InMemStats -> Update InMemStats ()
replaceInMemStats = put

recordedToday :: Query InMemStats Day
recordedToday = asks inMemToday

registerDownload :: PackageId -> Update InMemStats ()
registerDownload pkgId = do
  InMemStats day counts <- get
  put $ InMemStats day (cmInsert pkgId 1 counts)

makeAcidic ''InMemStats [ 'getInMemStats
                        , 'replaceInMemStats
                        , 'recordedToday
                        , 'registerDownload
                        ]