PerasCertDB: implement garbage collection

Author: amesgen

ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/Storage/PerasCertDB/API.hs

@@ -23,6 +23,8 @@ import qualified Ouroboros.Network.AnchoredFragment as AF
 data PerasCertDB m blk = PerasCertDB
   { addCert :: PerasCert blk -> m ()
   , getWeightSnapshot :: STM m (PerasWeightSnapshot blk)
+  , garbageCollect :: SlotNo -> m ()
+  -- ^ Garbage-collect state older than the given slot number.
   , closeDB :: m ()
   }
   deriving NoThunks via OnlyCheckWhnfNamed "PerasCertDB" (PerasCertDB m blk)

ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/Storage/PerasCertDB/Impl.hs

@@ -23,10 +23,9 @@ module Ouroboros.Consensus.Storage.PerasCertDB.Impl
 import Control.Monad (join)
 import Control.Tracer (Tracer, nullTracer, traceWith)
 import Data.Kind (Type)
+import qualified Data.Map.Merge.Strict as Map
 import Data.Map.Strict (Map)
 import qualified Data.Map.Strict as Map
-import Data.Set (Set)
-import qualified Data.Set as Set
 import GHC.Generics (Generic)
 import NoThunks.Class
 import Ouroboros.Consensus.Block
@@ -58,20 +57,19 @@ openDB ::
   Complete PerasCertDbArgs m blk ->
   m (PerasCertDB m blk)
 openDB args = do
-  pcdbRoundNos <- newTVarIO Set.empty
-  pcdbWeightByPoint <- newTVarIO Map.empty
+  pcdbVolatileState <- newTVarIO initialPerasVolatileCertState
   let env =
         PerasCertDbEnv
           { pcdbTracer
-          , pcdbRoundNos
-          , pcdbWeightByPoint
+          , pcdbVolatileState
           }
   h <- PerasCertDbHandle <$> newTVarIO (PerasCertDbOpen env)
   traceWith pcdbTracer OpenedPerasCertDB
   pure
     PerasCertDB
       { addCert = getEnv1 h implAddCert
       , getWeightSnapshot = getEnvSTM h implGetWeightSnapshot
+      , garbageCollect = getEnv1 h implGarbageCollect
       , closeDB = implCloseDB h
       }
  where
@@ -93,13 +91,8 @@ data PerasCertDbState m blk
 
 data PerasCertDbEnv m blk = PerasCertDbEnv
   { pcdbTracer :: !(Tracer m (TraceEvent blk))
-  , pcdbRoundNos :: !(StrictTVar m (Set PerasRoundNo))
+  , pcdbVolatileState :: !(StrictTVar m (PerasVolatileCertState blk))
   -- ^ The 'RoundNo's of all certificates currently in the db.
-  , pcdbWeightByPoint :: !(StrictTVar m (Map (Point blk) PerasWeight))
-  -- ^ The weight of boosted blocks w.r.t. the certificates currently in the
-  -- db.
-  --
-  -- INVARIANT: In sync with 'pcdbRoundNos'.
   }
   deriving NoThunks via OnlyCheckWhnfNamed "PerasCertDbEnv" (PerasCertDbEnv m blk)
 
@@ -154,20 +147,25 @@ implAddCert ::
 implAddCert env cert = do
   traceWith pcdbTracer $ AddingPerasCert roundNo boostedPt
   join $ atomically $ do
-    roundNos <- readTVar pcdbRoundNos
-    if Set.member roundNo roundNos
+    PerasVolatileCertState{pvcsCerts, pvcsWeightByPoint} <- readTVar pcdbVolatileState
+    if Map.member roundNo pvcsCerts
       then do
         pure $ traceWith pcdbTracer $ IgnoredCertAlreadyInDB roundNo boostedPt
       else do
-        writeTVar pcdbRoundNos $ Set.insert roundNo roundNos
-        -- Note that the same block might be boosted by multiple points.
-        modifyTVar pcdbWeightByPoint $ Map.insertWith (<>) boostedPt boostPerCert
+        writeTVar
+          pcdbVolatileState
+          PerasVolatileCertState
+            { pvcsCerts =
+                Map.insert roundNo cert pvcsCerts
+            , -- Note that the same block might be boosted by multiple points.
+              pvcsWeightByPoint =
+                Map.insertWith (<>) boostedPt boostPerCert pvcsWeightByPoint
+            }
         pure $ traceWith pcdbTracer $ AddedPerasCert roundNo boostedPt
  where
   PerasCertDbEnv
     { pcdbTracer
-    , pcdbRoundNos
-    , pcdbWeightByPoint
+    , pcdbVolatileState
     } = env
 
   roundNo = perasCertRound cert
@@ -176,8 +174,72 @@ implAddCert env cert = do
 implGetWeightSnapshot ::
   IOLike m =>
   PerasCertDbEnv m blk -> STM m (PerasWeightSnapshot blk)
-implGetWeightSnapshot PerasCertDbEnv{pcdbWeightByPoint} =
-  PerasWeightSnapshot <$> readTVar pcdbWeightByPoint
+implGetWeightSnapshot PerasCertDbEnv{pcdbVolatileState} =
+  PerasWeightSnapshot . pvcsWeightByPoint <$> readTVar pcdbVolatileState
+
+implGarbageCollect ::
+  forall m blk.
+  (IOLike m, StandardHash blk) =>
+  PerasCertDbEnv m blk -> SlotNo -> m ()
+implGarbageCollect PerasCertDbEnv{pcdbVolatileState} slot =
+  atomically $ modifyTVar pcdbVolatileState gc
+ where
+  gc :: PerasVolatileCertState blk -> PerasVolatileCertState blk
+  gc PerasVolatileCertState{pvcsCerts, pvcsWeightByPoint} =
+    PerasVolatileCertState
+      { pvcsCerts = certsToKeep
+      , pvcsWeightByPoint =
+          Map.merge
+            -- Do not touch weight of boosted blocks that we do not subtract any
+            -- weight from.
+            Map.preserveMissing
+            -- Irrelevant, the key set of @weightToRemove@ is a subset of the
+            -- key set of @pvcsWeightByPoint@.
+            Map.dropMissing
+            (Map.zipWithMaybeMatched $ \_pt -> subtractWeight)
+            pvcsWeightByPoint
+            weightToRemove
+      }
+   where
+    (certsToRemove, certsToKeep) =
+      Map.partition isTooOld pvcsCerts
+    isTooOld cert =
+      pointSlot (perasCertBoostedBlock cert) < NotOrigin slot
+    weightToRemove =
+      Map.fromListWith
+        (<>)
+        [ (perasCertBoostedBlock cert, boostPerCert)
+        | cert <- Map.elems certsToRemove
+        ]
+
+  subtractWeight :: PerasWeight -> PerasWeight -> Maybe PerasWeight
+  subtractWeight (PerasWeight w1) (PerasWeight w2)
+    | w1 > w2 = Just $ PerasWeight (w1 - w2)
+    | otherwise = Nothing
+
+{-------------------------------------------------------------------------------
+  Implementation-internal types
+-------------------------------------------------------------------------------}
+
+-- | Volatile Peras certificate state, i.e. certificates that could influence
+-- chain selection by boosting a volatile block.
+data PerasVolatileCertState blk = PerasVolatileCertState
+  { pvcsCerts :: !(Map PerasRoundNo (PerasCert blk))
+  -- ^ The boosted blocks by 'RoundNo' of all certificates currently in the db.
+  , pvcsWeightByPoint :: !(Map (Point blk) PerasWeight)
+  -- ^ The weight of boosted blocks w.r.t. the certificates currently in the db.
+  --
+  -- INVARIANT: In sync with 'pvcsCerts'.
+  }
+  deriving stock (Show, Generic)
+  deriving anyclass NoThunks
+
+initialPerasVolatileCertState :: PerasVolatileCertState blk
+initialPerasVolatileCertState =
+  PerasVolatileCertState
+    { pvcsCerts = Map.empty
+    , pvcsWeightByPoint = Map.empty
+    }
 
 {-------------------------------------------------------------------------------
   Trace types

ouroboros-consensus/test/storage-test/Test/Ouroboros/Storage/PerasCertDB/Model.hs

@@ -11,13 +11,14 @@ module Test.Ouroboros.Storage.PerasCertDB.Model
   , closeDB
   , addCert
   , getWeightSnapshot
+  , garbageCollect
   ) where
 
 import qualified Data.Map as Map
 import Data.Set (Set)
 import qualified Data.Set as Set
 import GHC.Generics (Generic)
-import Ouroboros.Consensus.Block (PerasCert, StandardHash, boostPerCert, perasCertBoostedBlock)
+import Ouroboros.Consensus.Block
 import Ouroboros.Consensus.Storage.PerasCertDB.API (PerasWeightSnapshot (..))
 
 data Model blk = Model
@@ -56,3 +57,9 @@ getWeightSnapshot Model{certs} = snap
             Map.empty
             certs
       }
+
+garbageCollect :: StandardHash blk => SlotNo -> Model blk -> Model blk
+garbageCollect slot model@Model{certs} =
+  model{certs = Set.filter keepCert certs}
+ where
+  keepCert cert = pointSlot (perasCertBoostedBlock cert) >= NotOrigin slot

ouroboros-consensus/test/storage-test/Test/Ouroboros/Storage/PerasCertDB/StateMachine.hs

@@ -51,13 +51,15 @@ instance StateModel Model where
     CloseDB :: Action Model ()
     AddCert :: PerasCert TestBlock -> Action Model ()
     GetWeightSnapshot :: Action Model (PerasWeightSnapshot TestBlock)
+    GarbageCollect :: SlotNo -> Action Model ()
 
   arbitraryAction _ (Model model)
     | model.open =
         frequency
           [ (1, pure $ Some CloseDB)
           , (20, Some <$> genAddCert)
           , (20, pure $ Some GetWeightSnapshot)
+          , (5, Some . GarbageCollect . SlotNo <$> arbitrary)
           ]
     | otherwise = pure $ Some OpenDB
    where
@@ -82,6 +84,7 @@ instance StateModel Model where
     CloseDB -> Model.closeDB model
     AddCert cert -> Model.addCert model cert
     GetWeightSnapshot -> model
+    GarbageCollect slot -> Model.garbageCollect slot model
 
   precondition (Model model) = \case
     OpenDB -> not model.open
@@ -93,6 +96,7 @@ instance StateModel Model where
          where
           p cert' = perasCertRound cert /= perasCertRound cert' || cert == cert'
         GetWeightSnapshot -> True
+        GarbageCollect _slot -> True
 
 deriving stock instance Show (Action Model a)
 deriving stock instance Eq (Action Model a)
@@ -114,6 +118,9 @@ instance RunModel Model (StateT (PerasCertDB IO TestBlock) IO) where
     GetWeightSnapshot -> do
       perasCertDB <- get
       lift $ atomically $ PerasCertDB.getWeightSnapshot perasCertDB
+    GarbageCollect slot -> do
+      perasCertDB <- get
+      lift $ PerasCertDB.garbageCollect perasCertDB slot
 
   postcondition (Model model, _) GetWeightSnapshot _ actual = do
     let expected = Model.getWeightSnapshot model