{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE ScopedTypeVariables #-}

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.Ledger.Basics

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 l = LedgerDbSize {
      -- | The tip of the ledger DB
      LedgerDbSize l -> Point l
ledgerDbTip       :: Point l

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

      -- | Size of the entire (in-memory) ledger DB
    , LedgerDbSize l -> Either CountFailure Word64
ledgerDbSizeTotal :: Either CountFailure Word64
    }
  deriving (Int -> LedgerDbSize l -> ShowS
[LedgerDbSize l] -> ShowS
LedgerDbSize l -> String
(Int -> LedgerDbSize l -> ShowS)
-> (LedgerDbSize l -> String)
-> ([LedgerDbSize l] -> ShowS)
-> Show (LedgerDbSize l)
forall l. StandardHash l => Int -> LedgerDbSize l -> ShowS
forall l. StandardHash l => [LedgerDbSize l] -> ShowS
forall l. StandardHash l => LedgerDbSize l -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [LedgerDbSize l] -> ShowS
$cshowList :: forall l. StandardHash l => [LedgerDbSize l] -> ShowS
show :: LedgerDbSize l -> String
$cshow :: forall l. StandardHash l => LedgerDbSize l -> String
showsPrec :: Int -> LedgerDbSize l -> ShowS
$cshowsPrec :: forall l. StandardHash l => Int -> LedgerDbSize l -> 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 :: forall m l. (MonadIO m, GetTip l)
                  => (Word64 -> Bool)
                  -> Tracer m (LedgerDbSize l)
                  -> Tracer m (LedgerDB l)
traceLedgerDbSize :: (Word64 -> Bool)
-> Tracer m (LedgerDbSize l) -> Tracer m (LedgerDB l)
traceLedgerDbSize Word64 -> Bool
p (Tracer LedgerDbSize l -> m ()
f) = (LedgerDB l -> m ()) -> Tracer m (LedgerDB l)
forall (m :: * -> *) a. (a -> m ()) -> Tracer m a
Tracer ((LedgerDB l -> m ()) -> Tracer m (LedgerDB l))
-> (LedgerDB l -> m ()) -> Tracer m (LedgerDB l)
forall a b. (a -> b) -> a -> b
$ \(!LedgerDB l
db) -> do
    let !ledger :: l
ledger = LedgerDB l -> l
forall l. GetTip l => LedgerDB l -> l
LedgerDB.ledgerDbCurrent LedgerDB l
db
        !tip :: Point l
tip    = l -> Point l
forall l. GetTip l => l -> Point l
getTip l
ledger

    Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Point l -> Bool
shouldTrace Point l
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 -> IO (Either CountFailure Word64)
forall a. a -> IO (Either CountFailure Word64)
computeHeapSize LedgerDB l
db
      LedgerDbSize l -> m ()
f (LedgerDbSize l -> m ()) -> LedgerDbSize l -> m ()
forall a b. (a -> b) -> a -> b
$ LedgerDbSize :: forall l.
Point l
-> Either CountFailure Word64
-> Either CountFailure Word64
-> LedgerDbSize l
LedgerDbSize {
              ledgerDbTip :: Point l
ledgerDbTip       = Point l
tip
            , ledgerDbSizeTip :: Either CountFailure Word64
ledgerDbSizeTip   = Either CountFailure Word64
sizeTip
            , ledgerDbSizeTotal :: Either CountFailure Word64
ledgerDbSizeTotal = Either CountFailure Word64
sizeTotal
            }
  where
    shouldTrace :: Point l -> Bool
    shouldTrace :: Point l -> Bool
shouldTrace Point l
GenesisPoint     = Word64 -> Bool
p Word64
0
    shouldTrace (BlockPoint SlotNo
s HeaderHash l
_) = Word64 -> Bool
p (SlotNo -> Word64
unSlotNo SlotNo
s)