hbs2/hbs2-storage-ncq/lib/HBS2/Storage/NCQ2/Internal/Probes.hs

103 lines
3.6 KiB
Haskell

{-# Language RecordWildCards #-}
module HBS2.Storage.NCQ2.Internal.Probes where
import HBS2.Prelude
import HBS2.Hash
import HBS2.Data.Types.Refs
import HBS2.System.Logger.Simple.ANSI
import HBS2.Misc.PrettyStuff
import HBS2.Data.Log.Structured.NCQ
import HBS2.Storage.NCQ2.Internal.Types
import HBS2.Storage.NCQ.Types
import Control.Monad.Trans.Maybe
import Data.Coerce
import Data.HashMap.Strict qualified as HM
import Data.List qualified as List
import Data.Maybe
import Data.Vector ((!))
import Data.Vector qualified as V
import Lens.Micro.Platform
import System.Random.MWC qualified as MWC
import UnliftIO
randomTrackedFile :: MonadUnliftIO m => NCQStorage2 -> m (Maybe FileKey)
randomTrackedFile ncq@NCQStorage2{..} = runMaybeT do
files0 <- lift (ncqListTrackedFiles ncq)
let files = V.toList $ V.filter (isNotPending . view _2) files0
guard (not (null files))
i <- liftIO $ MWC.uniformRM (0, length files - 1) ncqRndGen
pure (view _1 (files !! i))
randomTrackedFilePair :: MonadUnliftIO m => NCQStorage2 -> m (Maybe (FileKey, FileKey))
randomTrackedFilePair ncq@NCQStorage2{..} = runMaybeT do
files0 <- lift (ncqListTrackedFiles ncq)
let files = V.toList $ V.filter (isNotPending . view _2) files0
guard (length files >= 2)
(a, b) <- liftIO $ fix \loop -> do
i <- MWC.uniformRM (0, length files - 1) ncqRndGen
j <- MWC.uniformRM (0, length files - 1) ncqRndGen
if i == j then loop else pure (min i j, max i j)
let fa = view _1 (files !! a)
let fb = view _1 (files !! b)
pure (fa, fb)
ncqTombCountProbeFor :: MonadUnliftIO m => NCQStorage2 -> FileKey -> m (Maybe Int)
ncqTombCountProbeFor ncq@NCQStorage2{..} fkey = runMaybeT do
let fIndex = ncqGetFileName ncq $ toFileName (IndexFile fkey)
(bs, nh) <- liftIO (nwayHashMMapReadOnly fIndex) >>= toMPlus
liftIO do
ref <- newTVarIO 0
nwayHashScanAll nh bs $ \_ k v -> do
let NCQIdxEntry _ s = decodeEntry v
when (k /= ncqEmptyKey && s < 64) $
atomically $ modifyTVar' ref (+1)
readTVarIO ref
ncqKeyNumIntersectionProbeFor :: MonadUnliftIO m => NCQStorage2 -> (FileKey, FileKey) -> m (Maybe Int)
ncqKeyNumIntersectionProbeFor ncq@NCQStorage2{..} (fka, fkb) = runMaybeT do
let key = FactKey $ coerce $ hashObject @HbSync $ serialise $ List.sort [fka, fkb]
known <- lift (readTVarIO ncqFacts <&> HM.member key)
guard (not known)
let fIndexA = ncqGetFileName ncq (toFileName (IndexFile fka))
let fIndexB = ncqGetFileName ncq (toFileName (IndexFile fkb))
idxPair' <- liftIO $ try @_ @IOException do
(,) <$> nwayHashMMapReadOnly fIndexA
<*> nwayHashMMapReadOnly fIndexB
((bs1,n1),(bs2,n2)) <- case idxPair' of
Right (Just x, Just y) -> pure (x,y)
_ -> warn ("can't load index pair" <+> pretty (fka, fkb)) >> mzero
liftIO do
ref <- newTVarIO 0
nwayHashScanAll n1 bs1 $ \_ k _ -> when (k /= ncqEmptyKey) do
here <- ncqLookupIndex (coerce k) (bs2,n2)
when (isJust here) $ atomically $ modifyTVar' ref (+1)
readTVarIO ref
ncqTombCountProbe :: MonadUnliftIO m => NCQStorage2 -> m ()
ncqTombCountProbe ncq = useVersion ncq $ const $ void $ runMaybeT do
fk <- MaybeT (randomTrackedFile ncq)
count <- MaybeT (ncqTombCountProbeFor ncq fk)
debug $ yellow "ncqTombCountProbe" <+> pretty fk <+> pretty count
ncqKeyNumIntersectionProbe :: MonadUnliftIO m => NCQStorage2 -> m ()
ncqKeyNumIntersectionProbe ncq = useVersion ncq $ const $ void $ runMaybeT do
(fa, fb) <- MaybeT (randomTrackedFilePair ncq)
n <- MaybeT (ncqKeyNumIntersectionProbeFor ncq (fa, fb))
debug $ yellow "ncqKeyNumIntersectionProbe" <+> pretty fa <+> pretty fb <+> pretty n