From b11f233a607bf1b41a1e8b00b2f84f5d14ca5322 Mon Sep 17 00:00:00 2001
From: voidlizard <dzuikov@gmail.com>
Date: Fri, 1 Nov 2024 11:42:53 +0300
Subject: [PATCH] BlockDownload probe

---
 hbs2-peer/app/BlockDownload.hs | 24 ++++++++++++++++--------
 hbs2-peer/app/PeerMain.hs      |  4 ++++
 hbs2-peer/app/PeerTypes.hs     | 10 ++++++++++
 3 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/hbs2-peer/app/BlockDownload.hs b/hbs2-peer/app/BlockDownload.hs
index d7e7fc0a..37651e0b 100644
--- a/hbs2-peer/app/BlockDownload.hs
+++ b/hbs2-peer/app/BlockDownload.hs
@@ -3,7 +3,6 @@
 module BlockDownload where
 
 import HBS2.Peer.Prelude
-import HBS2.Base58
 import HBS2.Actors.Peer
 import HBS2.Data.Types.Peer
 import HBS2.Data.Detect
@@ -22,20 +21,16 @@ import HBS2.Misc.PrettyStuff
 
 import PeerTypes
 import PeerInfo
-import Brains
-import DownloadMon
 
 import Control.Concurrent.STM qualified as STM
 import Control.Monad.Trans.Cont
 import Control.Monad.Reader
 import Control.Monad.Trans.Maybe
-import Data.Cache qualified as Cache
 import Data.HashMap.Strict (HashMap)
 import Data.HashMap.Strict qualified as HashMap
 import Data.HashMap.Strict qualified as HM
 import Data.HashSet (HashSet)
 import Data.HashSet qualified as HS
-import Data.IntMap.Strict (IntMap)
 import Data.IntMap.Strict qualified as IntMap
 import Data.IntSet qualified as IntSet
 import Data.Maybe
@@ -44,9 +39,6 @@ import Data.ByteString.Lazy (ByteString)
 import Data.List qualified as L
 import Lens.Micro.Platform
 import Codec.Serialise
-import Data.Hashable
-import System.Random.Shuffle (shuffleM)
-import Control.Concurrent (getNumCapabilities)
 import Streaming.Prelude qualified as S
 import System.Random
 
@@ -630,6 +622,22 @@ blockDownloadLoop env0 = do
                 debug $ red "WAIT FOR PEERS TIMEOUT" <+> pretty blk
                 atomically $ writeTVar busy mempty
 
+    void $ ContT $ withAsync $ forever do
+      pause @'Seconds 10
+      let DownloadEnv{..} = env0
+      let DownloadMonEnv{..} = _downloadMon
+      p <- readTVarIO _downloadProbe
+      acceptReport p =<< S.toList_ do
+        S.yield =<< liftIO (readTVarIO _blockInQ <&> ("blockInQ",) . fromIntegral . HashMap.size)
+        S.yield =<< liftIO (readTVarIO _downloads <&> ("downloads",) . fromIntegral . HashMap.size)
+        S.yield =<< liftIO (readTVarIO nonces <&> ("nonces",) . fromIntegral . HashMap.size)
+        S.yield =<< liftIO (readTVarIO busy <&> ("busy",) . fromIntegral . HashMap.size)
+        S.yield =<< liftIO (readTVarIO rates <&> ("rates",) . fromIntegral . IntMap.size)
+        S.yield =<< liftIO (readTVarIO fetchH <&> ("fetchH",) . fromIntegral . HS.size)
+        S.yield =<< liftIO (readTVarIO sizes <&> ("sizes",) . fromIntegral . HashMap.size)
+        S.yield =<< liftIO (readTVarIO sizeReq <&> ("sizeReq",) . fromIntegral . HashMap.size)
+        S.yield =<< liftIO (readTVarIO seen <&> ("seen",) . fromIntegral . HashMap.size)
+
     forever do
       withPeerM e $ withDownload env0 do
         pause @'Seconds 5
diff --git a/hbs2-peer/app/PeerMain.hs b/hbs2-peer/app/PeerMain.hs
index df6a6a1f..801acbc7 100644
--- a/hbs2-peer/app/PeerMain.hs
+++ b/hbs2-peer/app/PeerMain.hs
@@ -815,6 +815,10 @@ runPeer opts = respawnOnError opts $ runResourceT do
 
   denv <- newDownloadEnv brains
 
+  dProbe <- newSimpleProbe "BlockDownload"
+  addProbe dProbe
+  downloadEnvSetProbe denv dProbe
+
   pl <- AnyPeerLocator <$> newBrainyPeerLocator @e (SomeBrains @e brains) mempty
 
   -- FIXME: messaing-watchdog
diff --git a/hbs2-peer/app/PeerTypes.hs b/hbs2-peer/app/PeerTypes.hs
index 695ac2b2..6f6bf2b9 100644
--- a/hbs2-peer/app/PeerTypes.hs
+++ b/hbs2-peer/app/PeerTypes.hs
@@ -4,6 +4,7 @@
 {-# Language AllowAmbiguousTypes #-}
 {-# Language MultiWayIf #-}
 {-# Language FunctionalDependencies #-}
+{-# Language RecordWildCards #-}
 module PeerTypes
   ( module PeerTypes
   , module PeerLogger
@@ -257,11 +258,19 @@ data DownloadEnv e =
   -- , _blockProposed  :: Cache (Hash HbSync, Peer e) ()
   , _downloadMon    :: DownloadMonEnv
   , _downloadBrains :: SomeBrains e
+  , _downloadProbe  :: TVar AnyProbe
   }
 
 makeLenses 'DownloadEnv
 
 
+downloadEnvSetProbe :: forall e m . (MonadIO m, MyPeer e)
+                    => DownloadEnv e
+                    -> AnyProbe
+                    -> m ()
+downloadEnvSetProbe DownloadEnv{..} p = do
+  atomically $ writeTVar _downloadProbe p
+
 newDownloadEnv :: (MonadIO m, MyPeer e, HasBrains e brains) => brains -> m (DownloadEnv e)
 newDownloadEnv brains = liftIO do
   DownloadEnv <$> newTVarIO mempty
@@ -271,6 +280,7 @@ newDownloadEnv brains = liftIO do
               -- <*> Cache.newCache (Just (toTimeSpec (2 :: Timeout 'Seconds)))
               <*> downloadMonNew
               <*> pure (SomeBrains brains)
+              <*> newTVarIO (AnyProbe ())
 
 newtype BlockDownloadM e m a =
   BlockDownloadM { fromBlockDownloadM :: ReaderT (DownloadEnv e) m a }