Merge pull request #186 from well-typed/finley/override-ping-on-server

Override ping rate limit on server
well-typed · Jul 11, 2024 · 4c1f20d · 4c1f20d
2 parents c341134 + 44d36dc
commit 4c1f20d
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 5 deletions.
diff --git a/kvstore/KVStore/Client.hs b/kvstore/KVStore/Client.hs
@@ -93,6 +93,7 @@ client Cmdline{
            cmdJSON
          , cmdSecure
          , cmdDisableTcpNoDelay
+         , cmdPingRateLimit
          } statsVar = do
     knownKeys <- RandomAccessSet.new
     random    <- RandomGen.new
@@ -126,7 +127,8 @@ client Cmdline{
     params :: ConnParams
     params = def {
           connHTTP2Settings = def {
-              http2TcpNoDelay = not cmdDisableTcpNoDelay
+              http2TcpNoDelay            = not cmdDisableTcpNoDelay
+            , http2OverridePingRateLimit = cmdPingRateLimit
             }
         }
 

diff --git a/kvstore/KVStore/Cmdline.hs b/kvstore/KVStore/Cmdline.hs
@@ -19,6 +19,7 @@ data Cmdline = Cmdline {
     , cmdJSON              :: Bool
     , cmdSecure            :: Bool
     , cmdDisableTcpNoDelay :: Bool
+    , cmdPingRateLimit     :: Maybe Int
     }
 
 data Mode =
@@ -69,6 +70,13 @@ parseCmdline =
               Opt.long "disable-tcp-nodelay"
             , Opt.help "Disable the TCP_NODELAY option"
             ])
+      <*> (Opt.optional $
+            Opt.option Opt.auto $ mconcat [
+              Opt.long "ping-rate-limit"
+            , Opt.metavar "PINGs/sec"
+            , Opt.help "Allow at most this many pings per second from the peer"
+            ]
+          )
 
 parseMode :: Parser Mode
 parseMode = asum [

diff --git a/kvstore/KVStore/Server.hs b/kvstore/KVStore/Server.hs
@@ -26,6 +26,7 @@ withKeyValueServer cmdline@Cmdline{
                        cmdJSON
                      , cmdSecure
                      , cmdDisableTcpNoDelay
+                     , cmdPingRateLimit
                      } k = do
     store <- Store.new
 
@@ -61,7 +62,8 @@ withKeyValueServer cmdline@Cmdline{
     params :: ServerParams
     params = def {
           serverHTTP2Settings = def {
-              http2TcpNoDelay = not cmdDisableTcpNoDelay
+              http2TcpNoDelay            = not cmdDisableTcpNoDelay
+            , http2OverridePingRateLimit = cmdPingRateLimit
             }
 
           -- The Java benchmark does not use compression (unclear if the Java

diff --git a/util/Network/GRPC/Util/HTTP2.hs b/util/Network/GRPC/Util/HTTP2.hs
@@ -1,6 +1,6 @@
 {-# LANGUAGE CPP #-}
 
-#include "MachDeps.h"
+
 
 module Network.GRPC.Util.HTTP2 (
     -- * General auxiliary
@@ -138,16 +138,21 @@ mkServerConfig ::
 mkServerConfig http2Settings numberOfWorkers =
     Server.defaultServerConfig {
         Server.numberOfWorkers =
-          fromMaybe
+          maybe
             (Server.numberOfWorkers Server.defaultServerConfig)
-            (fromIntegral <$> numberOfWorkers)
+            fromIntegral
+            numberOfWorkers
       , Server.connectionWindowSize = fromIntegral $
           http2ConnectionWindowSize http2Settings
       , Server.settings = Server.defaultSettings {
             Server.initialWindowSize = fromIntegral $
               http2StreamWindowSize http2Settings
           , Server.maxConcurrentStreams = Just . fromIntegral $
               http2MaxConcurrentStreams http2Settings
+          , Server.pingRateLimit =
+              fromMaybe
+                (Server.pingRateLimit Server.defaultSettings)
+                (http2OverridePingRateLimit http2Settings)
           }
       }