Allow to set default request queue length at runtime

3 files changed, 22 insertions, 20 deletions

diff --git a/src/Network/BitTorrent/Exchange/Message.hs b/src/Network/BitTorrent/Exchange/Message.hs
index d0d2bb03..7748f843 100644
--- a/src/Network/BitTorrent/Exchange/Message.hs
+++ b/src/Network/BitTorrent/Exchange/Message.hs
@@ -62,6 +62,7 @@ module Network.BitTorrent.Exchange.Message
        , StatusUpdate   (..)
        , Available      (..)
        , Transfer       (..)
+       , defaultRequestQueueLength
 
          -- ** Fast extension
        , FastMessage    (..)
@@ -483,6 +484,11 @@ instance PeerMessage Transfer where
   stats (Piece   p  ) = ByteStats (4 + 1) (4 + 4 + blockSize p) 0
   stats (Cancel  _  ) = ByteStats (4 + 1) (3 * 4) 0
 
+-- TODO increase
+-- | Max number of pending 'Request's inflight.
+defaultRequestQueueLength :: Int
+defaultRequestQueueLength = 1
+
 {-----------------------------------------------------------------------
 --  Fast messages
 -----------------------------------------------------------------------}
diff --git a/src/Network/BitTorrent/Exchange/Session.hs b/src/Network/BitTorrent/Exchange/Session.hs
index 416e00fd..1e72ba96 100644
--- a/src/Network/BitTorrent/Exchange/Session.hs
+++ b/src/Network/BitTorrent/Exchange/Session.hs
@@ -191,7 +191,7 @@ sendBroadcast msg = do
 
 fillRequestQueue :: Wire Session ()
 fillRequestQueue = do
-  maxN <- lift $ getAdvertisedQueueLength
+  maxN <- lift getMaxQueueLength
   rbf  <- use connBitfield
   addr <- asks connRemoteAddr
   blks <- withStatusUpdates $ do
diff --git a/src/Network/BitTorrent/Exchange/Wire.hs b/src/Network/BitTorrent/Exchange/Wire.hs
index e88b3ae5..4aebdd24 100644
--- a/src/Network/BitTorrent/Exchange/Wire.hs
+++ b/src/Network/BitTorrent/Exchange/Wire.hs
@@ -52,7 +52,7 @@ module Network.BitTorrent.Exchange.Wire
        , recvMessage
        , sendMessage
        , filterQueue
-       , getAdvertisedQueueLength
+       , getMaxQueueLength
 
          -- * Query
        , getMetadata
@@ -404,6 +404,8 @@ data Options = Options
     -- send any message for this period of time.
   , keepaliveTimeout    :: {-# UNPACK #-} !Int
 
+  , requestQueueLength  :: {-# UNPACK #-} !Int
+
     -- | Used to protect against flood attacks.
   , floodDetector       :: FloodDetector
 
@@ -434,11 +436,12 @@ data Options = Options
 -- change them.
 instance Default Options where
   def = Options
-    { keepaliveInterval = defaultKeepAliveInterval
-    , keepaliveTimeout  = defaultKeepAliveTimeout
-    , floodDetector     = def
-    , metadataFactor    = defaultMetadataFactor
-    , maxInfoDictSize   = defaultMaxInfoDictSize
+    { keepaliveInterval  = defaultKeepAliveInterval
+    , keepaliveTimeout   = defaultKeepAliveTimeout
+    , requestQueueLength = defaultRequestQueueLength
+    , floodDetector      = def
+    , metadataFactor     = defaultMetadataFactor
+    , maxInfoDictSize    = defaultMaxInfoDictSize
     }
 
 {-----------------------------------------------------------------------
@@ -591,19 +594,6 @@ instance MonadState ConnectionState (Connected s) where
 type Wire s a = ConduitM Message Message (Connected s) a
 
 {-----------------------------------------------------------------------
---  Query
------------------------------------------------------------------------}
-
--- TODO configurable
-defQueueLength :: Int
-defQueueLength = 1
-
-getAdvertisedQueueLength :: Connected s Int
-getAdvertisedQueueLength = do
-  ExtendedHandshake {..} <- use connRemoteEhs
-  return $ fromMaybe defQueueLength ehsQueueLength
-
-{-----------------------------------------------------------------------
 --  Wrapper
 -----------------------------------------------------------------------}
 
@@ -661,6 +651,12 @@ sendMessage msg = do
   ecaps <- use connExtCaps
   yield $ envelop ecaps msg
 
+getMaxQueueLength :: Connected s Int
+getMaxQueueLength = do
+  advertisedLen <- ehsQueueLength <$> use connRemoteEhs
+  defaultLen    <- asks (requestQueueLength . connOptions)
+  return $ fromMaybe defaultLen advertisedLen
+
 -- | Filter pending messages from send buffer.
 filterQueue :: (Message -> Bool) -> Wire s ()
 filterQueue p = lift $ do