Pass maybe Peer Advertise value to EstablishedPeers tasks

IntersectMBO · Oct 26, 2023 · 519a137 · 519a137
1 parent 739c034
commit 519a137
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 8 deletions.
diff --git a/ouroboros-network/CHANGELOG.md b/ouroboros-network/CHANGELOG.md
@@ -13,6 +13,7 @@
 * Adds ledger peers to diffusion simulation
 * Fix diffusion tests.
 * `demo-chain-sync`: added option parser, added new options.
+* Preserve `PeerAdvertise` information when connecting to peers.
 
 ## 0.9.1.0 -- 2023-08-22
 

diff --git a/ouroboros-network/src/Ouroboros/Network/PeerSelection/Governor/EstablishedPeers.hs b/ouroboros-network/src/Ouroboros/Network/PeerSelection/Governor/EstablishedPeers.hs
@@ -23,6 +23,7 @@ import           System.Random (randomR)
 import           Ouroboros.Network.PeerSelection.Governor.Types
 import           Ouroboros.Network.PeerSelection.LedgerPeers
                      (IsBigLedgerPeer (..))
+import           Ouroboros.Network.PeerSelection.PeerAdvertise (PeerAdvertise)
 import qualified Ouroboros.Network.PeerSelection.State.EstablishedPeers as EstablishedPeers
 import qualified Ouroboros.Network.PeerSelection.State.KnownPeers as KnownPeers
 import           Ouroboros.Network.PeerSelection.State.LocalRootPeers
@@ -125,8 +126,9 @@ belowTargetLocal actions
                           inProgressPromoteCold = inProgressPromoteCold
                                                <> selectedToPromote
                         },
-        decisionJobs  = [ jobPromoteColdPeer actions peer IsNotBigLedgerPeer
-                        | peer <- Set.toList selectedToPromote ]
+        decisionJobs  = [ jobPromoteColdPeer actions peer mbPeerAdvertise IsNotBigLedgerPeer
+                        | peer <- Set.toList selectedToPromote
+                        , let mbPeerAdvertise = KnownPeers.lookupPeerAdvertise peer knownPeers ]
       }
 
     -- If we could promote except that there are no peers currently available
@@ -215,8 +217,10 @@ belowTargetOther actions
                           inProgressPromoteCold = inProgressPromoteCold
                                                <> selectedToPromote
                         },
-        decisionJobs  = [ jobPromoteColdPeer actions peer IsNotBigLedgerPeer
-                        | peer <- Set.toList selectedToPromote ]
+        decisionJobs  = [ jobPromoteColdPeer actions peer mbPeerAdvertise IsNotBigLedgerPeer
+                        | peer <- Set.toList selectedToPromote
+                        , let mbPeerAdvertise = KnownPeers.lookupPeerAdvertise peer knownPeers
+                        ]
       }
 
     -- If we could connect except that there are no peers currently available
@@ -294,8 +298,10 @@ belowTargetBigLedgerPeers actions
                           inProgressPromoteCold = inProgressPromoteCold
                                                <> selectedToPromote
                         },
-        decisionJobs  = [ jobPromoteColdPeer actions peer IsBigLedgerPeer
-                        | peer <- Set.toList selectedToPromote ]
+        decisionJobs  = [ jobPromoteColdPeer actions peer mbPeerAdvertise IsBigLedgerPeer
+                        | peer <- Set.toList selectedToPromote
+                        , let mbPeerAdvertise = KnownPeers.lookupPeerAdvertise peer knownPeers
+                        ]
       }
 
     -- If we could connect except that there are no peers currently available
@@ -335,12 +341,13 @@ jobPromoteColdPeer :: forall peeraddr peerconn m.
                        (Monad m, Ord peeraddr)
                    => PeerSelectionActions peeraddr peerconn m
                    -> peeraddr
+                   -> Maybe PeerAdvertise
                    -> IsBigLedgerPeer
                    -> Job () m (Completion m peeraddr peerconn)
 jobPromoteColdPeer PeerSelectionActions {
                      peerStateActions = PeerStateActions {establishPeerConnection},
                      peerConnToPeerSharing
-                   } peeraddr isBigLedgerPeer =
+                   } peeraddr mbPeerAdvertise isBigLedgerPeer =
     Job job handler () "promoteColdPeer"
   where
     handler :: SomeException -> m (Completion m peeraddr peerconn)
@@ -415,7 +422,7 @@ jobPromoteColdPeer PeerSelectionActions {
             -- Update PeerSharing value in KnownPeers
             -- This will compute the appropriate peer sharing value using
             -- 'combinePeerInformation'
-            knownPeers'       = KnownPeers.insert (Map.singleton peeraddr (Just peerSharing, Nothing, Nothing))
+            knownPeers'       = KnownPeers.insert (Map.singleton peeraddr (Just peerSharing, mbPeerAdvertise, Nothing))
                               $ KnownPeers.clearTepidFlag peeraddr $
                                     KnownPeers.resetFailCount
                                         peeraddr

diff --git a/ouroboros-network/src/Ouroboros/Network/PeerSelection/State/KnownPeers.hs b/ouroboros-network/src/Ouroboros/Network/PeerSelection/State/KnownPeers.hs
@@ -19,6 +19,7 @@ module Ouroboros.Network.PeerSelection.State.KnownPeers
   , resetFailCount
   , lookupFailCount
   , lookupTepidFlag
+  , lookupPeerAdvertise
   , setTepidFlag
   , clearTepidFlag
     -- ** Tracking when we can (re)connect
@@ -312,6 +313,13 @@ lookupTepidFlag :: Ord peeraddr
 lookupTepidFlag peeraddr KnownPeers{allPeers} =
     knownPeerTepid <$> Map.lookup peeraddr allPeers
 
+lookupPeerAdvertise :: Ord peeraddr
+                    => peeraddr
+                    -> KnownPeers peeraddr
+                    -> Maybe PeerAdvertise
+lookupPeerAdvertise peeraddr KnownPeers{allPeers} =
+    knownPeerAdvertise <$> Map.lookup peeraddr allPeers
+
 setTepidFlag' :: Ord peeraddr
              => Bool
              -> peeraddr