{-# LANGUAGE BangPatterns #-}

module Ouroboros.Consensus.Util.TraceSize (
    -- * Generic
    traceSize
    -- * Ledger DB specific
  , LedgerDbSize(..)
  , traceLedgerDbSize
  ) where

import           Cardano.Prelude (CountFailure, computeHeapSize)
import           Control.Monad (when)
import           Control.Monad.IO.Class
import           Control.Tracer
import           Data.Word

import           Ouroboros.Consensus.Block

import           Ouroboros.Consensus.Storage.LedgerDB.InMemory (LedgerDB)
import qualified Ouroboros.Consensus.Storage.LedgerDB.InMemory as LedgerDB

{-------------------------------------------------------------------------------
  Generic
-------------------------------------------------------------------------------}

-- | Generic helper to trace a value and its size
traceSize :: MonadIO m
          => Tracer m (a, Either CountFailure Word64)
          -> Tracer m a
traceSize :: Tracer m (a, Either CountFailure Word64) -> Tracer m a
traceSize (Tracer (a, Either CountFailure Word64) -> m ()
f) = (a -> m ()) -> Tracer m a
forall (m :: * -> *) a. (a -> m ()) -> Tracer m a
Tracer ((a -> m ()) -> Tracer m a) -> (a -> m ()) -> Tracer m a
forall a b. (a -> b) -> a -> b
$ \a
a -> do
    Either CountFailure Word64
sz <- IO (Either CountFailure Word64) -> m (Either CountFailure Word64)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Either CountFailure Word64) -> m (Either CountFailure Word64))
-> IO (Either CountFailure Word64)
-> m (Either CountFailure Word64)
forall a b. (a -> b) -> a -> b
$ a -> IO (Either CountFailure Word64)
forall a. a -> IO (Either CountFailure Word64)
computeHeapSize a
a
    (a, Either CountFailure Word64) -> m ()
f (a
a, Either CountFailure Word64
sz)

{-------------------------------------------------------------------------------
  Ledger DB specific
-------------------------------------------------------------------------------}

data LedgerDbSize blk = LedgerDbSize {
      -- | The tip of the ledger DB
      LedgerDbSize blk -> Point blk
ledgerDbTip       :: Point blk

      -- | Size of the ledger at the tip of the DB
    , LedgerDbSize blk -> Either CountFailure Word64
ledgerDbSizeTip   :: Either CountFailure Word64

      -- | Size of the entire (in-memory) ledger DB
    , LedgerDbSize blk -> Either CountFailure Word64
ledgerDbSizeTotal :: Either CountFailure Word64
    }
  deriving (Int -> LedgerDbSize blk -> ShowS
[LedgerDbSize blk] -> ShowS
LedgerDbSize blk -> String
(Int -> LedgerDbSize blk -> ShowS)
-> (LedgerDbSize blk -> String)
-> ([LedgerDbSize blk] -> ShowS)
-> Show (LedgerDbSize blk)
forall blk. StandardHash blk => Int -> LedgerDbSize blk -> ShowS
forall blk. StandardHash blk => [LedgerDbSize blk] -> ShowS
forall blk. StandardHash blk => LedgerDbSize blk -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [LedgerDbSize blk] -> ShowS
$cshowList :: forall blk. StandardHash blk => [LedgerDbSize blk] -> ShowS
show :: LedgerDbSize blk -> String
$cshow :: forall blk. StandardHash blk => LedgerDbSize blk -> String
showsPrec :: Int -> LedgerDbSize blk -> ShowS
$cshowsPrec :: forall blk. StandardHash blk => Int -> LedgerDbSize blk -> ShowS
Show)

-- | Trace the size of the ledger
--
-- Only traces slots for which the predicate results true (genesis will be
-- considered to be slot 0).
traceLedgerDbSize :: MonadIO m
                  => (Word64 -> Bool)
                  -> Tracer m (LedgerDbSize blk)
                  -> Tracer m (LedgerDB l (RealPoint blk))
traceLedgerDbSize :: (Word64 -> Bool)
-> Tracer m (LedgerDbSize blk)
-> Tracer m (LedgerDB l (RealPoint blk))
traceLedgerDbSize Word64 -> Bool
p (Tracer LedgerDbSize blk -> m ()
f) = (LedgerDB l (RealPoint blk) -> m ())
-> Tracer m (LedgerDB l (RealPoint blk))
forall (m :: * -> *) a. (a -> m ()) -> Tracer m a
Tracer ((LedgerDB l (RealPoint blk) -> m ())
 -> Tracer m (LedgerDB l (RealPoint blk)))
-> (LedgerDB l (RealPoint blk) -> m ())
-> Tracer m (LedgerDB l (RealPoint blk))
forall a b. (a -> b) -> a -> b
$ \(!LedgerDB l (RealPoint blk)
db) -> do
    let !ledger :: l
ledger = LedgerDB l (RealPoint blk) -> l
forall l r. LedgerDB l r -> l
LedgerDB.ledgerDbCurrent LedgerDB l (RealPoint blk)
db
        !tip :: WithOrigin (RealPoint blk)
tip    = LedgerDB l (RealPoint blk) -> WithOrigin (RealPoint blk)
forall l r. LedgerDB l r -> WithOrigin r
LedgerDB.ledgerDbTip LedgerDB l (RealPoint blk)
db

    Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (WithOrigin (RealPoint blk) -> Bool
forall blk. WithOrigin (RealPoint blk) -> Bool
shouldTrace WithOrigin (RealPoint blk)
tip) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
      Either CountFailure Word64
sizeTip   <- IO (Either CountFailure Word64) -> m (Either CountFailure Word64)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Either CountFailure Word64) -> m (Either CountFailure Word64))
-> IO (Either CountFailure Word64)
-> m (Either CountFailure Word64)
forall a b. (a -> b) -> a -> b
$ l -> IO (Either CountFailure Word64)
forall a. a -> IO (Either CountFailure Word64)
computeHeapSize l
ledger
      Either CountFailure Word64
sizeTotal <- IO (Either CountFailure Word64) -> m (Either CountFailure Word64)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Either CountFailure Word64) -> m (Either CountFailure Word64))
-> IO (Either CountFailure Word64)
-> m (Either CountFailure Word64)
forall a b. (a -> b) -> a -> b
$ LedgerDB l (RealPoint blk) -> IO (Either CountFailure Word64)
forall a. a -> IO (Either CountFailure Word64)
computeHeapSize LedgerDB l (RealPoint blk)
db
      LedgerDbSize blk -> m ()
f (LedgerDbSize blk -> m ()) -> LedgerDbSize blk -> m ()
forall a b. (a -> b) -> a -> b
$ LedgerDbSize :: forall blk.
Point blk
-> Either CountFailure Word64
-> Either CountFailure Word64
-> LedgerDbSize blk
LedgerDbSize {
              ledgerDbTip :: Point blk
ledgerDbTip       = WithOrigin (RealPoint blk) -> Point blk
forall blk. WithOrigin (RealPoint blk) -> Point blk
withOriginRealPointToPoint WithOrigin (RealPoint blk)
tip
            , ledgerDbSizeTip :: Either CountFailure Word64
ledgerDbSizeTip   = Either CountFailure Word64
sizeTip
            , ledgerDbSizeTotal :: Either CountFailure Word64
ledgerDbSizeTotal = Either CountFailure Word64
sizeTotal
            }
  where
    shouldTrace :: WithOrigin (RealPoint blk) -> Bool
    shouldTrace :: WithOrigin (RealPoint blk) -> Bool
shouldTrace WithOrigin (RealPoint blk)
Origin         = Word64 -> Bool
p Word64
0
    shouldTrace (NotOrigin RealPoint blk
pt) = Word64 -> Bool
p (SlotNo -> Word64
unSlotNo (RealPoint blk -> SlotNo
forall blk. RealPoint blk -> SlotNo
realPointSlot RealPoint blk
pt))