KeepAlive Client collect samples on first packet

Update the KeepAlive client so that it collect RTT samples for the first
packet. Previusly two request where sent, one after the other because
the first couldn't be used to get a sample.

ouroboros-network-protocols/CHANGELOG.md

@@ -6,6 +6,8 @@
 
 ### Non-breaking changes
 
+* Make it possible for KeepAlive client to collect a rtt sample for the first packet.
+
 ## 0.5.3.0 -- 2023-10-26
 
 ### Non-breaking changes

ouroboros-network-protocols/src/Ouroboros/Network/Protocol/KeepAlive/Client.hs

@@ -1,8 +1,10 @@
-{-# LANGUAGE DataKinds #-}
-{-# LANGUAGE GADTs     #-}
+{-# LANGUAGE DataKinds           #-}
+{-# LANGUAGE GADTs               #-}
+{-# LANGUAGE ScopedTypeVariables #-}
 
 module Ouroboros.Network.Protocol.KeepAlive.Client
   ( KeepAliveClient (..)
+  , KeepAliveClientSt (..)
   , keepAliveClientPeer
   ) where
 
@@ -11,31 +13,40 @@ import           Network.TypedProtocol.Core
 import           Ouroboros.Network.Protocol.KeepAlive.Type
 
 
-data KeepAliveClient m a where
+newtype KeepAliveClient m a = KeepAliveClient (m (KeepAliveClientSt m a))
+
+data KeepAliveClientSt m a where
     SendMsgKeepAlive
       :: Cookie
-      -> (m (KeepAliveClient m a))
-      -> KeepAliveClient m a
+      -> m (KeepAliveClientSt m a)
+      -> KeepAliveClientSt m a
 
     SendMsgDone
       :: m a
-      -> KeepAliveClient m a
+      -> KeepAliveClientSt m a
 
 
 -- | Interpret a particular client action sequence into the client side of the
 -- 'KeepAlive' protocol.
 --
 keepAliveClientPeer
-  :: MonadThrow m
+  :: forall m a. MonadThrow m
   => KeepAliveClient m a
   -> Peer KeepAlive AsClient StClient m a
-
-keepAliveClientPeer (SendMsgDone mresult) =
-    Yield (ClientAgency TokClient) MsgDone $
-      Effect (Done TokDone <$> mresult)
-
-keepAliveClientPeer (SendMsgKeepAlive cookieReq next) =
-    Yield (ClientAgency TokClient) (MsgKeepAlive cookieReq) $
-      Await (ServerAgency TokServer) $ \(MsgKeepAliveResponse cookieRsp) ->
-        if cookieReq == cookieRsp then Effect $ keepAliveClientPeer <$> next
-                                  else Effect $ throwIO $ KeepAliveCookieMissmatch cookieReq cookieRsp
+keepAliveClientPeer (KeepAliveClient client) =
+   Effect $ keepAliveClientStPeer <$> client
+ where
+
+   keepAliveClientStPeer
+     :: KeepAliveClientSt m a
+     -> Peer KeepAlive AsClient StClient m a
+
+   keepAliveClientStPeer (SendMsgDone mresult) =
+     Yield (ClientAgency TokClient) MsgDone $
+       Effect (Done TokDone <$> mresult)
+
+   keepAliveClientStPeer (SendMsgKeepAlive cookieReq next) =
+     Yield (ClientAgency TokClient) (MsgKeepAlive cookieReq) $
+       Await (ServerAgency TokServer) $ \(MsgKeepAliveResponse cookieRsp) ->
+         if cookieReq == cookieRsp then Effect $ keepAliveClientStPeer <$> next
+                                   else Effect $ throwIO $ KeepAliveCookieMissmatch cookieReq cookieRsp

ouroboros-network-protocols/testlib/Ouroboros/Network/Protocol/KeepAlive/Direct.hs

@@ -1,19 +1,28 @@
 {-# LANGUAGE NamedFieldPuns #-}
+{-# LANGUAGE RankNTypes     #-}
 
 module Ouroboros.Network.Protocol.KeepAlive.Direct where
 
 import           Ouroboros.Network.Protocol.KeepAlive.Client
 import           Ouroboros.Network.Protocol.KeepAlive.Server
 
-direct :: Monad m
+
+direct :: forall a b m. Monad m
        => KeepAliveServer m a
        -> KeepAliveClient m b
        -> m (a, b)
-direct KeepAliveServer { recvMsgDone }
-       (SendMsgDone mdone) =
-    (,) <$> recvMsgDone <*> mdone
-direct KeepAliveServer { recvMsgKeepAlive }
-       (SendMsgKeepAlive _cookie mclient) = do
-    server <- recvMsgKeepAlive
-    client <- mclient
-    direct server client
+direct srv (KeepAliveClient clientM) = do
+  go srv =<< clientM
+ where
+   go :: Monad m
+      => KeepAliveServer m a
+      -> KeepAliveClientSt m b
+      -> m (a, b)
+   go KeepAliveServer { recvMsgDone }
+          (SendMsgDone mdone) =
+       (,) <$> recvMsgDone <*> mdone
+   go KeepAliveServer { recvMsgKeepAlive }
+          (SendMsgKeepAlive _cookie mclient) = do
+       server <- recvMsgKeepAlive
+       client <- mclient
+       go server client

ouroboros-network-protocols/testlib/Ouroboros/Network/Protocol/KeepAlive/Examples.hs

@@ -16,9 +16,9 @@ keepAliveClientApply :: forall acc m. Monad m
                     -> Int
                     -- ^ count of number of requests
                     -> KeepAliveClient m acc
-keepAliveClientApply f = go
+keepAliveClientApply f aa an = KeepAliveClient $ return (go aa an)
   where
-    go :: acc -> Int -> KeepAliveClient m acc
+    go :: acc -> Int -> KeepAliveClientSt m acc
     go acc n
       | n <= 0
       = SendMsgDone (pure acc)

ouroboros-network/CHANGELOG.md

@@ -12,6 +12,7 @@
 
 ### Non-breaking changes
 
+* Update KeepAlive client to collect a rtt sample for the first packet.
 * Less aggresive churning of established and known peers.
 * Add peer sharing to wireshark dissector.
 * Adds ledger peers to diffusion simulation

ouroboros-network/src/Ouroboros/Network/KeepAlive.hs

@@ -52,7 +52,9 @@ keepAliveClient
     -> KeepAliveClient m ()
 keepAliveClient tracer inRng controlMessageSTM peer dqCtx KeepAliveInterval { keepAliveInterval } =
     let (cookie, rng) = random inRng in
-    SendMsgKeepAlive (Cookie cookie) (go rng Nothing)
+    KeepAliveClient $ do
+      startTime <- getMonotonicTime
+      return $ SendMsgKeepAlive (Cookie cookie) (go rng startTime)
   where
     payloadSize = 2
 
@@ -70,40 +72,32 @@ keepAliveClient tracer inRng controlMessageSTM peer dqCtx KeepAliveInterval { ke
                  then return Continue
                  else retry
 
-    go :: StdGen -> Maybe Time -> m (KeepAliveClient m ())
-    go rng startTime_m = do
+    go :: StdGen -> Time -> m (KeepAliveClientSt m ())
+    go rng startTime = do
       endTime <- getMonotonicTime
-      case startTime_m of
-           Just startTime -> do
-               let rtt = diffTime endTime startTime
-                   sample = fromSample startTime endTime payloadSize
-               gsv' <- atomically $ do
-                   m <- readTVar dqCtx
-                   assert (peer `M.member` m) $ do
-                     let (gsv', m') = M.updateLookupWithKey
-                             (\_ a -> if sampleTime a == Time 0 -- Ignore the initial dummy value
-                                         then Just sample
-                                         else Just $ sample <> a
-                             ) peer m
-                     writeTVar dqCtx m'
-                     return $ fromJust gsv'
-               traceWith tracer $ AddSample peer rtt gsv'
-
-           Nothing        -> return ()
-
-      let keepAliveInterval' = case startTime_m of
-                                    Just _  -> keepAliveInterval
-                                    Nothing -> 0 -- The first time we send a packet directly.
-
-      delayVar <- registerDelay keepAliveInterval'
+      let rtt = diffTime endTime startTime
+          sample = fromSample startTime endTime payloadSize
+      gsv' <- atomically $ do
+          m <- readTVar dqCtx
+          assert (peer `M.member` m) $ do
+            let (gsv', m') = M.updateLookupWithKey
+                    (\_ a -> if sampleTime a == Time 0 -- Ignore the initial dummy value
+                                then Just sample
+                                else Just $ sample <> a
+                    ) peer m
+            writeTVar dqCtx m'
+            return $ fromJust gsv'
+      traceWith tracer $ AddSample peer rtt gsv'
+
+      delayVar <- registerDelay keepAliveInterval
       decision <- atomically (decisionSTM delayVar)
       now <- getMonotonicTime
       case decision of
         -- 'decisionSTM' above cannot return 'Quiesce'
         Quiesce   -> error "keepAlive: impossible happened"
         Continue  ->
             let (cookie, rng') = random rng in
-            pure (SendMsgKeepAlive (Cookie cookie) $ go rng' $ Just now)
+            pure (SendMsgKeepAlive (Cookie cookie) $ go rng' now)
         Terminate -> pure (SendMsgDone (pure ()))