simulation: fix filtering of expired votes

Saizan · Saizan · commit fd8d3ce25f18 · 2025-02-12T11:32:08.000+01:00
diff --git a/simulation/src/LeiosProtocol/Relay.hs b/simulation/src/LeiosProtocol/Relay.hs
@@ -10,6 +10,7 @@
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE MultiParamTypeClasses #-}
+{-# LANGUAGE NondecreasingIndentation #-}
 {-# LANGUAGE OverloadedRecordDot #-}
 {-# LANGUAGE PolyKinds #-}
 {-# LANGUAGE RankNTypes #-}
@@ -495,6 +496,8 @@ data SubmitPolicy = SubmitInOrder | SubmitAll
 
 data RelayConsumerConfig id header body m = RelayConsumerConfig
   { relay :: !RelayConfig
+  , shouldIgnore :: m (header -> Bool)
+  -- ^ headers to ignore, e.g. already received or coming too late.
   , validateHeaders :: [header] -> m ()
   , headerId :: !(header -> id)
   , prioritize :: !(Map id header -> [header] -> [header])
@@ -661,9 +664,9 @@ relayConsumerPipelined config sst =
   -- \| Takes an STM action for the updated local state, so that
   -- requestBodies can update inFlightVar in the same STM tx.
   idleM ::
-    STM m (RelayConsumerLocalState id header body n) ->
+    m (RelayConsumerLocalState id header body n) ->
     m (RelayConsumer id header body n 'StIdle m ())
-  idleM mlst = atomically $ do
+  idleM mlst = do
     lst <- mlst
     let canRequestMoreBodies = not (Map.null lst.available)
     case lst.pendingRequests of
@@ -715,39 +718,41 @@ relayConsumerPipelined config sst =
   requestBodies ::
     forall (n :: N).
     RelayConsumerLocalState id header body n ->
-    STM m (RelayConsumer id header body n 'StIdle m ())
+    m (RelayConsumer id header body n 'StIdle m ())
   requestBodies lst = do
-    -- New headers are filtered before becoming available, but we have
-    -- to filter `lst.available` again in the same STM tx that sets them as
-    -- `inFlight`.
-    inFlight <- readTVar sst.inFlightVar
-    relayBuffer <- readTVar sst.relayBufferVar
-    let (ignored, available1) =
-          Map.partitionWithKey
-            ( \k _ ->
-                k `Set.member` inFlight
-                  || k `RB.member` relayBuffer
-            )
-            lst.available
-    -- Ignored headers are set to Nothing in the buffer so they can be acknowledged later.
-    let buffer' = lst.buffer <> Map.map (const Nothing) ignored
-
-    let hdrsToRequest =
-          take (fromIntegral config.maxBodiesToRequest) $
-            config.prioritize available1 (mapMaybe (`Map.lookup` available1) $ Foldable.toList $ lst.window)
-    let idsToRequest = map config.headerId hdrsToRequest
-    let idsToRequestSet = Set.fromList idsToRequest
-    modifyTVar' sst.inFlightVar $ Set.union idsToRequestSet
-    let available2 = Map.withoutKeys available1 idsToRequestSet
-    let !lst' = lst{pendingRequests = Succ lst.pendingRequests, available = available2, buffer = buffer'}
-    return $
-      TS.YieldPipelined
-        (MsgRequestBodies idsToRequest)
-        ( TS.ReceiverAwait $ \case
-            MsgRespondBodies bodies ->
-              TS.ReceiverDone (CollectBodies hdrsToRequest bodies)
-        )
-        (requestHeadersNonBlocking lst')
+    isIgnored <- config.shouldIgnore
+    atomically $ do
+      -- New headers are filtered before becoming available, but we have
+      -- to filter `lst.available` again in the same STM tx that sets them as
+      -- `inFlight`.
+      inFlight <- readTVar sst.inFlightVar
+      let (ignored, available1) =
+            Map.partitionWithKey
+              ( \k hd ->
+                  k `Set.member` inFlight
+                    || isIgnored hd
+              )
+              lst.available
+
+      -- Ignored headers are set to Nothing in the buffer so they can be acknowledged later.
+      let buffer' = lst.buffer <> Map.map (const Nothing) ignored
+
+      let hdrsToRequest =
+            take (fromIntegral config.maxBodiesToRequest) $
+              config.prioritize available1 (mapMaybe (`Map.lookup` available1) $ Foldable.toList $ lst.window)
+      let idsToRequest = map config.headerId hdrsToRequest
+      let idsToRequestSet = Set.fromList idsToRequest
+      modifyTVar' sst.inFlightVar $ Set.union idsToRequestSet
+      let available2 = Map.withoutKeys available1 idsToRequestSet
+      let !lst' = lst{pendingRequests = Succ lst.pendingRequests, available = available2, buffer = buffer'}
+      return $
+        TS.YieldPipelined
+          (MsgRequestBodies idsToRequest)
+          ( TS.ReceiverAwait $ \case
+              MsgRespondBodies bodies ->
+                TS.ReceiverDone (CollectBodies hdrsToRequest bodies)
+          )
+          (requestHeadersNonBlocking lst')
 
   windowAdjust ::
     forall (n :: N).
@@ -860,6 +865,7 @@ relayConsumerPipelined config sst =
           -- though not all have replies.
           buffer1 = lst.buffer <> idsRequestedWithBodiesReceived
 
+          !_ = assert (relayConsumerLocalStateInvariant $ lst{buffer = buffer1}) ()
           -- We have to update the window here eagerly and not
           -- delay it to requestBodies, otherwise we could end up blocking in
           -- idle on more pipelined results rather than being able to
@@ -888,7 +894,7 @@ relayConsumerPipelined config sst =
           extraToSubmit = case config.submitPolicy of
             SubmitInOrder -> []
             SubmitAll ->
-              mapMaybe (\id' -> join (Map.lookup id' buffer1)) $
+              mapMaybe (\id' -> join (Map.lookup id' buffer2)) $
                 Foldable.toList window'
 
           -- And set them to `Nothing` in the buffer so they can be
@@ -937,19 +943,19 @@ relayConsumerPipelined config sst =
     RelayConsumerLocalState id header body n ->
     StrictSeq id ->
     Map id header ->
-    STM m (RelayConsumerLocalState id header body n)
+    m (RelayConsumerLocalState id header body n)
   acknowledgeIds lst idsSeq _ | Seq.null idsSeq = pure lst
   acknowledgeIds lst idsSeq idsMap = do
-    relayBuffer <- readTVar sst.relayBufferVar
-    inFlight <- readTVar sst.inFlightVar
+    isIgnored <- config.shouldIgnore
+    inFlight <- readTVarIO sst.inFlightVar
 
     let
       -- Divide the new ids in two: those that are already in the
       -- relay buffer or in flight (both locally and shared in-flight)
       -- and those that are not. We'll request some bodies from the latter.
       (ignoredIds, availableIdsMp) =
         Map.partitionWithKey
-          (\id' _ -> id' `RB.member` relayBuffer || id' `Set.member` inFlight)
+          (\id' hd -> isIgnored hd || id' `Set.member` inFlight)
           idsMap
 
       availableIdsU =
diff --git a/simulation/src/LeiosProtocol/Short/Node.hs b/simulation/src/LeiosProtocol/Short/Node.hs
@@ -99,6 +99,8 @@ data LeiosNodeState m = LeiosNodeState
   , relayIBState :: !(RelayIBState m)
   , relayEBState :: !(RelayEBState m)
   , relayVoteState :: !(RelayVoteState m)
+  , prunedVoteStateToVar :: !(TVar m SlotNo)
+  -- ^ TODO: refactor into RelayState.
   , ibDeliveryTimesVar :: !(TVar m (Map InputBlockId UTCTime))
   , taskQueue :: !(TaskMultiQueue LeiosNodeTask m)
   , waitingForRBVar :: !(TVar m (Map (HeaderHash RankingBlock) [STM m ()]))
@@ -107,6 +109,8 @@ data LeiosNodeState m = LeiosNodeState
   -- ^ waiting for ledger state of RB block to be validated.
   , ledgerStateVar :: !(TVar m (Map (HeaderHash RankingBlock) LedgerState))
   , ibsNeededForEBVar :: !(TVar m (Map EndorseBlockId (Set InputBlockId)))
+  , votesForEBVar :: !(TVar m (Map EndorseBlockId (Map VoteId Word64)))
+  -- ^ TODO: prune of EBs that won't make it into chain anymore.
   }
 
 data LeiosNodeTask
@@ -220,8 +224,9 @@ relayIBConfig ::
   LeiosNodeConfig ->
   ([InputBlockHeader] -> m ()) ->
   SubmitBlocks m InputBlockHeader InputBlockBody ->
+  RelayIBState m ->
   RelayConsumerConfig InputBlockId InputBlockHeader InputBlockBody m
-relayIBConfig _tracer cfg validateHeaders submitBlocks =
+relayIBConfig _tracer cfg validateHeaders submitBlocks st =
   RelayConsumerConfig
     { relay = RelayConfig{maxWindowSize = coerce cfg.leios.ibDiffusion.maxWindowSize}
     , headerId = (.id)
@@ -231,15 +236,19 @@ relayIBConfig _tracer cfg validateHeaders submitBlocks =
     , maxHeadersToRequest = cfg.leios.ibDiffusion.maxHeadersToRequest
     , maxBodiesToRequest = cfg.leios.ibDiffusion.maxBodiesToRequest
     , submitBlocks
+    , shouldIgnore = do
+        buff <- readTVarIO st.relayBufferVar
+        return $ flip RB.member buff . (.id)
     }
 
 relayEBConfig ::
-  MonadDelay m =>
+  (MonadDelay m, MonadSTM m) =>
   Tracer m LeiosNodeEvent ->
   LeiosNodeConfig ->
   SubmitBlocks m EndorseBlockId EndorseBlock ->
+  RelayEBState m ->
   RelayConsumerConfig EndorseBlockId (RelayHeader EndorseBlockId) EndorseBlock m
-relayEBConfig _tracer cfg submitBlocks =
+relayEBConfig _tracer cfg submitBlocks st =
   RelayConsumerConfig
     { relay = RelayConfig{maxWindowSize = coerce cfg.leios.ebDiffusion.maxWindowSize}
     , headerId = (.id)
@@ -250,15 +259,20 @@ relayEBConfig _tracer cfg submitBlocks =
     , maxBodiesToRequest = cfg.leios.ebDiffusion.maxBodiesToRequest
     , submitBlocks = \hbs t k ->
         submitBlocks (map (first (.id)) hbs) t (k . map (\(i, b) -> (RelayHeader i b.slot, b)))
+    , shouldIgnore = do
+        buff <- readTVarIO st.relayBufferVar
+        return $ flip RB.member buff . (.id)
     }
 
 relayVoteConfig ::
-  MonadDelay m =>
+  (MonadDelay m, Monad (STM m), MonadSTM m, MonadTime m) =>
   Tracer m LeiosNodeEvent ->
   LeiosNodeConfig ->
   SubmitBlocks m VoteId VoteMsg ->
+  RelayVoteState m ->
+  LeiosNodeState m ->
   RelayConsumerConfig VoteId (RelayHeader VoteId) VoteMsg m
-relayVoteConfig _tracer cfg submitBlocks =
+relayVoteConfig _tracer cfg submitBlocks _ leiosState =
   RelayConsumerConfig
     { relay = RelayConfig{maxWindowSize = coerce cfg.leios.voteDiffusion.maxWindowSize}
     , headerId = (.id)
@@ -269,6 +283,12 @@ relayVoteConfig _tracer cfg submitBlocks =
     , maxBodiesToRequest = cfg.leios.voteDiffusion.maxBodiesToRequest
     , submitBlocks = \hbs t k ->
         submitBlocks (map (first (.id)) hbs) t (k . map (\(i, b) -> (RelayHeader i b.slot, b)))
+    , shouldIgnore = atomically $ do
+        buffer <- readTVar leiosState.relayVoteState.relayBufferVar
+        prunedTo <- readTVar leiosState.prunedVoteStateToVar
+        return $ \hd ->
+          hd.slot < prunedTo
+            || hd.id `RB.member` buffer
     }
 
 queueAndWait :: (MonadSTM m, MonadDelay m) => LeiosNodeState m -> LeiosNodeTask -> [CPUTask] -> m ()
@@ -298,6 +318,8 @@ newLeiosNodeState cfg = do
   waitingForRBVar <- newTVarIO Map.empty
   waitingForLedgerStateVar <- newTVarIO Map.empty
   taskQueue <- atomically $ newTaskMultiQueue cfg.processingQueueBound
+  prunedVoteStateToVar <- newTVarIO (toEnum 0)
+  votesForEBVar <- newTVarIO Map.empty
   return $ LeiosNodeState{..}
 
 leiosNode ::
@@ -352,21 +374,21 @@ leiosNode tracer cfg followers peers = do
 
   ibThreads <-
     setupRelay
-      (relayIBConfig tracer cfg valHeaderIB submitIB)
+      (relayIBConfig tracer cfg valHeaderIB submitIB relayIBState)
       relayIBState
       (map (.protocolIB) followers)
       (map (.protocolIB) peers)
 
   ebThreads <-
     setupRelay
-      (relayEBConfig tracer cfg submitEB)
+      (relayEBConfig tracer cfg submitEB relayEBState)
       relayEBState
       (map (.protocolEB) followers)
       (map (.protocolEB) peers)
 
   voteThreads <-
     setupRelay
-      (relayVoteConfig tracer cfg submitVote)
+      (relayVoteConfig tracer cfg submitVote relayVoteState leiosState)
       relayVoteState
       (map (.protocolVote) followers)
       (map (.protocolVote) peers)
@@ -419,13 +441,15 @@ pruneVoteBuffer ::
 pruneVoteBuffer _tracer cfg st = go (toEnum 0)
  where
   go p = do
-    let last_vote_recv = snd $ stageRangeOf cfg.leios p VoteRecv
     let last_vote_send = snd $ stageRangeOf cfg.leios p VoteSend
+    let last_vote_recv = snd $ stageRangeOf cfg.leios p VoteRecv
+    let pruneTo = succ last_vote_send
     _ <- waitNextSlot cfg.slotConfig (succ last_vote_recv)
     atomically $ do
       modifyTVar' st.relayVoteState.relayBufferVar $
         RB.filter $
-          \RB.EntryWithTicket{value} -> (snd value).slot <= last_vote_send
+          \RB.EntryWithTicket{value} -> (snd value).slot >= pruneTo
+      writeTVar st.prunedVoteStateToVar $! pruneTo
     go (succ p)
 
 computeLedgerStateThread ::
@@ -479,6 +503,18 @@ adoptEB leiosState eb = do
   let ibsNeeded = Map.fromList [(eb.id, Set.fromList eb.inputBlocks Set.\\ ibs)]
   modifyTVar' leiosState.ibsNeededForEBVar (`Map.union` ibsNeeded)
 
+adoptVote :: MonadSTM m => LeiosNodeState m -> VoteMsg -> STM m ()
+adoptVote leiosState v = do
+  -- We keep tally for each EB as votes arrive, so the relayVoteBuffer
+  -- can be pruned without effects on EB certification.
+  modifyTVar' leiosState.votesForEBVar $
+    Map.unionWith Map.union $
+      Map.fromListWith
+        Map.union
+        [ (eb, Map.singleton v.id v.votes)
+        | eb <- v.endorseBlocks
+        ]
+
 dispatchValidation ::
   forall m.
   (MonadMVar m, MonadFork m, MonadAsync m, MonadSTM m, MonadTime m, MonadDelay m) =>
@@ -509,7 +545,9 @@ dispatchValidation tracer cfg leiosState req =
       adoptEB leiosState eb
     traceEnterState [eb] EventEB
   valVote v completion = labelTask . (ValVote,) . (\p -> cpuTask p cfg.leios.delays.voteMsgValidation v,) $ do
-    atomically $ completion [v]
+    atomically $ do
+      completion [v]
+      adoptVote leiosState v
     traceEnterState [v] EventVote
 
   go :: ValidationRequest m -> STM m [(LeiosNodeTask, (CPUTask, m ()))]
@@ -589,7 +627,11 @@ generator tracer cfg st = do
             adoptEB st eb
           traceWith tracer (LeiosNodeEvent Generate (EventEB eb))
         SomeAction Generate.Vote v -> (GenVote,) $ do
-          atomically $ modifyTVar' st.relayVoteState.relayBufferVar (RB.snoc v.id (RelayHeader v.id v.slot, v))
+          atomically $ do
+            modifyTVar'
+              st.relayVoteState.relayBufferVar
+              (RB.snoc v.id (RelayHeader v.id v.slot, v))
+            adoptVote st v
           traceWith tracer (LeiosNodeEvent Generate (EventVote v))
   let LeiosNodeConfig{..} = cfg
   leiosBlockGenerator $ LeiosGeneratorConfig{submit = mapM_ submitOne, ..}
@@ -602,8 +644,7 @@ mkBuffersView cfg st = BuffersView{..}
     -- though it's getting more expensive as we go.
     chain <- PraosNode.preferredChain st.praosState
     bufferEB <- readTVar st.relayEBState.relayBufferVar
-    bufferVotes <- map snd . RB.values <$> readTVar st.relayVoteState.relayBufferVar
-
+    votesForEB <- readTVar st.votesForEBVar
     -- RBs in the same chain should not contain certificates for the same pipeline.
     let pipelinesInChain =
           Set.fromList $
@@ -612,14 +653,6 @@ mkBuffersView cfg st = BuffersView{..}
             , (ebId, _) <- rb.blockBody.endorseBlocks
             , Just (_, eb) <- [RB.lookup bufferEB ebId]
             ]
-    -- TODO: cache?
-    let votesForEB =
-          Map.fromListWith
-            Map.union
-            [ (eb, Map.singleton v.id v.votes)
-            | v <- bufferVotes
-            , eb <- v.endorseBlocks
-            ]
     let totalVotes = fromIntegral . sum . Map.elems
     let tryCertify eb = do
           votes <- Map.lookup eb.id votesForEB
diff --git a/simulation/src/LeiosProtocol/SimTestRelay.hs b/simulation/src/LeiosProtocol/SimTestRelay.hs
@@ -137,6 +137,9 @@ relayNode
     let relayConsumerConfig =
           RelayConsumerConfig
             { relay = relayConfig
+            , shouldIgnore = atomically $ do
+                rb <- readTVar relayBufferVar
+                return $ \hd -> RB.member (testHeaderId hd) rb
             , -- sequential validation of headers
               validateHeaders = map (const 0.1) >>> sum >>> \d -> when (d >= 0) $ threadDelay d
             , headerId = testHeaderId
