TCP probe

hbs2-core/lib/HBS2/Net/Messaging/TCP.hs

@@ -8,6 +8,7 @@ module HBS2.Net.Messaging.TCP
   , tcpPeerConn
   , tcpCookie
   , tcpOnClientStarted
+  , messagingTCPSetProbe
   ) where
 
 import HBS2.Clock
@@ -43,7 +44,7 @@ import UnliftIO.Async
 import UnliftIO.STM
 import UnliftIO.Exception qualified as U
 
-
+{-HLINT ignore "Functor law"-}
 
 -- FIXME: control-recv-capacity-to-avoid-leaks
 
@@ -62,11 +63,15 @@ data MessagingTCP =
   , _tcpRecv             :: TQueue (Peer L4Proto, ByteString)
   , _tcpDefer            :: TVar (HashMap (Peer L4Proto) [(TimeSpec, ByteString)])
   , _tcpDeferEv          :: TQueue ()
+  , _tcpProbe            :: TVar AnyProbe
   , _tcpOnClientStarted  :: PeerAddr L4Proto -> Word64 -> IO () -- ^ Cient TCP connection succeed
   }
 
 makeLenses 'MessagingTCP
 
+messagingTCPSetProbe :: MonadIO m => MessagingTCP -> AnyProbe -> m ()
+messagingTCPSetProbe MessagingTCP{..} p = atomically $ writeTVar _tcpProbe p
+
 newMessagingTCP :: ( MonadIO m
                    , FromSockAddr 'TCP (Peer L4Proto)
                    )
@@ -86,6 +91,7 @@ newMessagingTCP pa = liftIO do
     <*> newTQueueIO
     <*> newTVarIO mempty
     <*> newTQueueIO
+    <*> newTVarIO (AnyProbe ())
     <*> pure (\_ _ -> none) -- do nothing by default
 
 instance Messaging MessagingTCP L4Proto ByteString where
@@ -358,14 +364,16 @@ connectPeerTCP env peer = liftIO do
 -- FIXME: link-all-asyncs
 
 runMessagingTCP :: forall m . MonadIO m => MessagingTCP -> m ()
-runMessagingTCP env = liftIO do
+runMessagingTCP env@MessagingTCP{..} = liftIO do
+
+  void $ flip runContT pure do
 
     own <- toPeerAddr $ view tcpOwnPeer env
     let (L4Address _ (IPAddrPort (i,p))) = own
 
     let defs = view tcpDefer env
 
-  mon <- async $ forever do
+    mon <- ContT $ withAsync $ forever do
       pause @'Seconds 30
       now <- getTimeCoarse
 
@@ -378,7 +386,7 @@ runMessagingTCP env = liftIO do
                                [] -> Nothing
                                xs -> Just xs
 
-  con <- async $ forever do
+    con <- ContT $ withAsync $ forever do
 
       let ev = view tcpDeferEv env
 
@@ -412,7 +420,7 @@ runMessagingTCP env = liftIO do
 
           pure ()
 
-  stat <- async $ forever do
+    stat <- ContT $ withAsync $ forever do
       pause @'Seconds 120
       ps <- readTVarIO $ view tcpConnPeer env
       let peers = HashMap.toList ps
@@ -425,8 +433,24 @@ runMessagingTCP env = liftIO do
 
     mapM_ link [mon,con,stat]
 
-  liftIO (
-    listen (Host (show i)) (show p) $ \(sock, sa) -> do
+    probes <- ContT $ withAsync $ forever do
+      pause @'Seconds 10
+      probe <- readTVarIO _tcpProbe
+      values <- atomically do
+        a <- readTVar _tcpConnPeer <&> HashMap.size <&> (L.singleton . ("tcpConnPeer",))
+        b <- readTVar _tcpPeerConn <&> HashMap.size <&> (L.singleton . ("tcpPeerConn",))
+        c <- readTVar _tcpConnUsed <&> HashMap.size <&> (L.singleton . ("tcpConnUsed",))
+        d <- readTVar _tcpConnQ <&> HashMap.size <&> (L.singleton . ("tcpConnQ",))
+        e <- readTVar _tcpPeerPx <&> HashMap.size <&> (L.singleton . ("tcpPeerPx",))
+        f <- readTVar _tcpPeerXp <&> HashMap.size <&> (L.singleton . ("tcpPeerXp",))
+        g <- readTVar _tcpDefer <&> HashMap.size <&> (L.singleton . ("tcpPeerDefer",))
+        pure $ mconcat [a, b, c, d, e, f, g]
+      acceptReport probe (fmap (over _2 fromIntegral) values)
+
+
+    ContT $ bracket (pure ()) $ \_ -> mapM_ cancel [mon,con,stat,probes]
+
+    liftIO $ listen (Host (show i)) (show p) $ \(sock, sa) -> do
       withFdSocket sock setCloseOnExecIfNeeded
       debug $ "Listening on" <+> pretty sa
 
@@ -449,7 +473,7 @@ runMessagingTCP env = liftIO do
 
           debug $ "CLOSING CONNECTION" <+> pretty remote
           shutdown so ShutdownBoth
-          close so ) `U.finally` mapM_ cancel [mon,con,stat]
+          close so -- ) -- `U.finally` mapM_ cancel [mon,con,stat]
 
 
 traceCmd :: forall a ann b m . ( Pretty a

hbs2-peer/app/PeerMain.hs

@@ -857,6 +857,11 @@ runPeer opts = respawnOnError opts $ runResourceT do
                              <&> set tcpOnClientStarted (onClientTCPConnected brains)
                              <&> set tcpSOCKS5 socks5
 
+                lift do
+                  tcpProbe <- newSimpleProbe "MessagingTCP"
+                  addProbe tcpProbe
+                  messagingTCPSetProbe tcpEnv tcpProbe
+
                 void $ liftIO ( async do
                           runMessagingTCP tcpEnv
                             `U.withException` \(e :: SomeException) -> do
@@ -909,6 +914,10 @@ runPeer opts = respawnOnError opts $ runResourceT do
 
   rce <- refChanWorkerEnv conf penv denv refChanNotifySource
 
+  rcwProbe <- newSimpleProbe "RefChanWorker"
+  addProbe rcwProbe
+  refChanWorkerEnvSetProbe rce rcwProbe
+
   let refChanAdapter =
         RefChanAdapter
         { refChanOnHead = refChanOnHeadFn rce

hbs2-peer/app/RefChan.hs

@@ -12,6 +12,7 @@ module RefChan (
   , runRefChanRelyWorker
   , refChanWorkerEnv
   , refChanNotifyOnUpdated
+  , refChanWorkerEnvSetProbe
   ) where
 
 import HBS2.Prelude.Plated
@@ -105,10 +106,18 @@ data RefChanWorkerEnv e =
   , _refChanWorkerNotifiersInbox  :: TQueue (RefChanNotify e) -- ^ to rely messages from clients to gossip
   , _refChanWorkerNotifiersDone   :: Cache (Hash HbSync) ()
   , _refChanWorkerLocalRelyDone   :: Cache (Peer UNIX, Hash HbSync) ()
+  , _refChanWorkerProbe           :: TVar AnyProbe
   }
 
 makeLenses 'RefChanWorkerEnv
 
+refChanWorkerEnvSetProbe :: forall m e . (MonadIO m, ForRefChans e)
+                         => RefChanWorkerEnv e
+                         -> AnyProbe
+                         -> m ()
+refChanWorkerEnvSetProbe RefChanWorkerEnv{..} probe = do
+  liftIO $ atomically $ writeTVar _refChanWorkerProbe probe
+
 refChanWorkerEnv :: forall m e . (MonadIO m, ForRefChans e)
                  => PeerConfig
                  -> PeerEnv e
@@ -127,6 +136,8 @@ refChanWorkerEnv conf pe de nsource =
     <*> newTQueueIO
     <*> Cache.newCache (Just defRequestLimit)
     <*> Cache.newCache (Just defRequestLimit)
+    <*> newTVarIO (AnyProbe ())
+
 
 refChanOnHeadFn :: forall e m . (ForRefChans e, MonadIO m) => RefChanWorkerEnv e -> RefChanId e -> RefChanHeadBlockTran e -> m ()
 refChanOnHeadFn env chan tran = do
@@ -595,6 +606,15 @@ refChanWorker env@RefChanWorkerEnv{..} brains = do
 
     bullshit <- ContT $ withAsync $ forever do
      pause @'Seconds 10
+     probe <- readTVarIO _refChanWorkerProbe
+     values <- atomically do
+      refChanWorkerEnvDownloadSize <- readTVar _refChanWorkerEnvDownload <&> HashMap.size
+      refChanWorkerNotifiersSize   <-  readTVar _refChanWorkerNotifiers <&> HashMap.size
+      pure [ ("refChanWorkerEnvDownloadSize", fromIntegral refChanWorkerEnvDownloadSize)
+           ,  ("refChanWorkerNotifiersSize", fromIntegral refChanWorkerNotifiersSize)
+           ]
+
+     acceptReport probe values
      debug "I'm refchan worker"
 
     waitAnyCatchCancel [hw,downloads,polls,wtrans,merge,cleanup1,bullshit]