From 0d4ed9400e17492cc46b968dbadc71ae36605c13 Mon Sep 17 00:00:00 2001 From: Marcin Szamotulski Date: Thu, 21 May 2020 15:27:29 +0200 Subject: [PATCH] connection-manager: numberOfConnections For connection manager which allows to include inbound connections we need a method to count how many connections are currently included (for rate limiting). We count all connection, weather inbound or outbound. In future we could count only inbound connections, but this will require to know the state of mux (to know when on-demand responders actaully started). --- .../Network/ConnectionManager/Core.hs | 17 +++++++- .../Network/ConnectionManager/Types.hs | 40 ++++++++++++++++--- 2 files changed, 50 insertions(+), 7 deletions(-) diff --git a/ouroboros-network-framework/src/Ouroboros/Network/ConnectionManager/Core.hs b/ouroboros-network-framework/src/Ouroboros/Network/ConnectionManager/Core.hs index 49e5f545610..16e4c76de4a 100644 --- a/ouroboros-network-framework/src/Ouroboros/Network/ConnectionManager/Core.hs +++ b/ouroboros-network-framework/src/Ouroboros/Network/ConnectionManager/Core.hs @@ -114,13 +114,23 @@ withConnectionManager ConnectionManagerArguments { ConnectionHandler (WithResponderMode inboundHandler) -> ConnectionManager (WithResponderMode - (includeConnection stateVar mainThreadId inboundHandler Inbound)) + InboundConnectionManager { + icmIncludeConnection = + includeConnection stateVar mainThreadId inboundHandler Inbound, + icmNumberOfConnections = + countConnections stateVar + }) ConnectionHandler (WithInitiatorResponderMode outboundHandler inboundHandler) -> ConnectionManager (WithInitiatorResponderMode (connectAndInclude stateVar mainThreadId outboundHandler) - (includeConnection stateVar mainThreadId inboundHandler Inbound)) + InboundConnectionManager { + icmIncludeConnection = + includeConnection stateVar mainThreadId inboundHandler Inbound, + icmNumberOfConnections = + countConnections stateVar + }) k connectionManager `finally` do @@ -132,6 +142,9 @@ withConnectionManager ConnectionManagerArguments { >> close connectionSnocket chSocket ) state where + countConnections :: StrictTMVar m (State peerAddr socket muxPromise m) -> STM m Int + countConnections stateVar = Map.size <$> readTMVar stateVar + -- Include a connection in the 'State'; we use this for both inbound and -- outbound (via 'connectAndInclude' below) connections. -- diff --git a/ouroboros-network-framework/src/Ouroboros/Network/ConnectionManager/Types.hs b/ouroboros-network-framework/src/Ouroboros/Network/ConnectionManager/Types.hs index 9e62b6f187e..6d62c4c2095 100644 --- a/ouroboros-network-framework/src/Ouroboros/Network/ConnectionManager/Types.hs +++ b/ouroboros-network-framework/src/Ouroboros/Network/ConnectionManager/Types.hs @@ -19,10 +19,12 @@ module Ouroboros.Network.ConnectionManager.Types , ConnectionManagerArguments (..) -- * 'ConnectionManager' , ConnectionManager (..) + , InboundConnectionManager (..) , IncludeOutboundConnection , includeOutboundConnection , IncludeInboundConnection , includeInboundConnection + , numberOfConnections -- * Exceptions , ExceptionInHandler (..) -- * Mux types @@ -188,6 +190,19 @@ type IncludeInboundConnection socket peerAddr muxPromise m = socket -> peerAddr -> m (STM m muxPromise) +-- | Inbound connection manager. For a server implementation we also need to +-- know how many connections are now managed by the connection manager. +-- +-- This type is an internal detail of 'Ouroboros.Network.ConnectionManager' +-- +data InboundConnectionManager (muxMode :: MuxMode) socket peerAddr muxPromise m where + InboundConnectionManager + :: HasResponder muxMode ~ True + => { icmIncludeConnection :: IncludeInboundConnection socket peerAddr muxPromise m + , icmNumberOfConnections :: STM m Int + } + -> InboundConnectionManager muxMode socket peerAddr muxPromise m + -- | 'ConnectionManager'. -- -- We identify resources (e.g. 'Network.Socket.Socket') by their address. It @@ -198,10 +213,14 @@ type IncludeInboundConnection socket peerAddr muxPromise m -- newtype ConnectionManager (muxMode :: MuxMode) socket peerAddr muxPromise m = ConnectionManager { runConnectionManager - :: WithMuxMode muxMode (IncludeOutboundConnection peerAddr muxPromise m) - (IncludeInboundConnection socket peerAddr muxPromise m) + :: WithMuxMode muxMode (IncludeOutboundConnection peerAddr muxPromise m) + (InboundConnectionManager muxMode socket peerAddr muxPromise m) } +-- +-- ConnectionManager API +-- + -- | Include outbound connection into 'ConnectionManager'. -- includeOutboundConnection :: HasInitiator muxMode ~ True @@ -212,10 +231,21 @@ includeOutboundConnection = withInitiatorMode . runConnectionManager -- | Include an inbound connection into 'ConnectionManager'. -- includeInboundConnection :: HasResponder muxMode ~ True - => ConnectionManager muxMode socket peerAddr muxPromise m - -> IncludeInboundConnection socket peerAddr muxPromise m -includeInboundConnection = withResponderMode . runConnectionManager + => ConnectionManager muxMode socket peerAddr muxPromise m + -> IncludeInboundConnection socket peerAddr muxPromise m +includeInboundConnection = icmIncludeConnection . withResponderMode . runConnectionManager +-- | Number of currently included connections. +-- +-- Note: we count all connection: both inbound and outbound. In a future +-- version we could count only inbound connections, but this will require +-- tracking state inside mux if the responder side has started running (through +-- on-demand interface). This is currently not exposed by mux. +-- +numberOfConnections :: HasResponder muxMode ~ True + => ConnectionManager muxMode socket peerAddr muxPromise m + -> STM m Int +numberOfConnections = icmNumberOfConnections . withResponderMode . runConnectionManager -- -- Tracing