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