Add ability to observe the number of connections halibut opens for polling connections

source/Halibut.Tests/Support/TestConnectionsObserver.cs

@@ -10,12 +10,16 @@ public class TestConnectionsObserver : IConnectionsObserver
     {
         readonly ConcurrentBag<bool> connectionAcceptedAuthorized = new();
         readonly ConcurrentBag<bool> connectionClosedAuthorized = new();
-
+        readonly ConcurrentBag<Uri> connectionAcceptedForSubscriptions = new();
+        readonly ConcurrentBag<Uri> connectionClosedForSubscriptions = new();
+
         public long ConnectionAcceptedCount => connectionAcceptedAuthorized.Count;
         public long ConnectionClosedCount => connectionClosedAuthorized.Count;
 
         public IReadOnlyList<bool> ConnectionAcceptedAuthorized => connectionAcceptedAuthorized.ToList();
         public IReadOnlyList<bool> ConnectionClosedAuthorized => connectionClosedAuthorized.ToList();
+        public IReadOnlyList<Uri> ConnectionAcceptedForSubscriptions => connectionAcceptedForSubscriptions.ToList();
+        public IReadOnlyList<Uri> ConnectionClosedForSubscriptions => connectionClosedForSubscriptions.ToList();
 
         public void ConnectionAccepted(bool authorized)
         {
@@ -26,5 +30,15 @@ public void ConnectionClosed(bool authorized)
         {
             connectionClosedAuthorized.Add(authorized);
         }
+
+        public void ConnectionAcceptedFor(Uri subscriptionId)
+        {
+            connectionAcceptedForSubscriptions.Add(subscriptionId);
+        }
+
+        public void ConnectionClosedFor(Uri subscriptionId)
+        {
+            connectionClosedForSubscriptions.Add(subscriptionId);
+        }
     }
 }

source/Halibut.Tests/Transport/Protocol/ProtocolFixture.cs

@@ -27,7 +27,7 @@ public void SetUp()
             stream.SetRemoteIdentity(new RemoteIdentity(RemoteIdentityType.Server));
             var limits = new HalibutTimeoutsAndLimitsForTestsBuilder().Build();
             var activeConnectionsLimiter = new ActiveTcpConnectionsLimiter(limits);
-            protocol = new MessageExchangeProtocol(stream, new HalibutTimeoutsAndLimitsForTestsBuilder().Build(), activeConnectionsLimiter, Substitute.For<ILog>());
+            protocol = new MessageExchangeProtocol(stream, new HalibutTimeoutsAndLimitsForTestsBuilder().Build(), activeConnectionsLimiter, NoOpConnectionsObserver.Instance, Substitute.For<ILog>());
         }
 
         // TODO - ASYNC ME UP! ExchangeAsClientAsync cancellation

source/Halibut.Tests/Transport/SecureClientFixture.cs

@@ -75,7 +75,7 @@ public async Task SecureClientClearsPoolWhenAllConnectionsCorrupt()
                 var connection = Substitute.For<IConnection>();
                 var limits = new HalibutTimeoutsAndLimitsForTestsBuilder().Build();
                 var activeConnectionLimiter = new ActiveTcpConnectionsLimiter(limits);
-                connection.Protocol.Returns(new MessageExchangeProtocol(stream, limits, activeConnectionLimiter, log));
+                connection.Protocol.Returns(new MessageExchangeProtocol(stream, limits, activeConnectionLimiter, NoOpConnectionsObserver.Instance, log));
 
                 await connectionManager.ReleaseConnectionAsync(endpoint, connection, CancellationToken.None);
             }
@@ -109,7 +109,7 @@ static MessageExchangeProtocol GetProtocol(Stream stream, ILog logger)
         {
             var limits = new HalibutTimeoutsAndLimitsForTestsBuilder().Build();
             var activeConnectionLimiter = new ActiveTcpConnectionsLimiter(limits);
-            return new MessageExchangeProtocol(new MessageExchangeStream(stream, new MessageSerializerBuilder(new LogFactory()).Build(), new NoOpControlMessageObserver(), limits, logger), limits, activeConnectionLimiter, logger);
+            return new MessageExchangeProtocol(new MessageExchangeStream(stream, new MessageSerializerBuilder(new LogFactory()).Build(), new NoOpControlMessageObserver(), limits, logger), limits, activeConnectionLimiter, NoOpConnectionsObserver.Instance, logger);
         }
     }
 }

source/Halibut/HalibutRuntime.cs

@@ -119,7 +119,7 @@ public int Listen(int port)
 
         ExchangeProtocolBuilder ExchangeProtocolBuilder()
         {
-            return (stream, log) => new MessageExchangeProtocol(new MessageExchangeStream(stream, messageSerializer, controlMessageObserver, TimeoutsAndLimits, log), TimeoutsAndLimits, activeTcpConnectionsLimiter, log);
+            return (stream, log) => new MessageExchangeProtocol(new MessageExchangeStream(stream, messageSerializer, controlMessageObserver, TimeoutsAndLimits, log), TimeoutsAndLimits, activeTcpConnectionsLimiter, connectionsObserver, log);
         }
 
         public int Listen(IPEndPoint endpoint)

source/Halibut/Transport/ActiveTcpConnectionsLimiter.cs

@@ -10,7 +10,6 @@ public interface IActiveTcpConnectionsLimiter
     {
         IDisposable LeaseActiveTcpConnection(Uri subscriptionId);
 
-        IDisposable CreateUnlimitedLease();
     }
 
     public class ActiveTcpConnectionsLimiter : IActiveTcpConnectionsLimiter
@@ -35,7 +34,7 @@ public IDisposable LeaseActiveTcpConnection(Uri subscriptionId)
             return new LimitingAuthorizedTcpConnectionLease(subscriptionId, activeConnectionCountPerSubscriptionId, timeoutsAndLimits.MaximumActiveTcpConnectionsPerPollingSubscription.Value);
         }
 
-        public IDisposable CreateUnlimitedLease()
+        IDisposable CreateUnlimitedLease()
         {
             return new UnlimitedAuthorizedTcpConnectionLease();
         }

source/Halibut/Transport/Observability/IConnectionsObserver.cs

@@ -1,3 +1,5 @@
+using System;
+
 namespace Halibut.Transport.Observability
 {
     public interface IConnectionsObserver
@@ -6,7 +8,7 @@ public interface IConnectionsObserver
         /// The connection has been accepted and no bytes have been read from the wire.
         ///
         /// In this context server is anything that listens on a port.
-        /// 
+        ///
         /// This is called when any of the following occurs:
         /// - When a "server" accepts a connection from a polling service (either websocket or regular)
         /// - When a "server" accepts a connection from a listening client (so in this case the server is the service)
@@ -16,8 +18,23 @@ public interface IConnectionsObserver
         /// <summary>
         /// A previously accepted connection has been closed.
         ///
-        /// For every call to ConnectionClosed() their can be at most one call to this method. 
+        /// For every call to ConnectionClosed() their can be at most one call to this method.
         /// </summary>
         public void ConnectionClosed(bool authorized);
+
+        /// <summary>
+        /// Called once the connection is known to be for a
+        /// polling subscriber (i.e. after the subscription id has been read off the wire), and only
+        /// for connections that were not rejected for exceeding the active connection limit.
+        ///
+        /// For every call to this method there will be at most one matching call to ConnectionClosedFor()
+        /// with the same subscriptionId.
+        /// </summary>
+        public void ConnectionAcceptedFor(Uri subscriptionId);
+
+        /// <summary>
+        /// A previously accepted polling subscriber connection has been closed.
+        /// </summary>
+        public void ConnectionClosedFor(Uri subscriptionId);
     }
 }

source/Halibut/Transport/Observability/NoOpConnectionsObserver.cs

@@ -1,3 +1,5 @@
+using System;
+
 namespace Halibut.Transport.Observability
 {
     public class NoOpConnectionsObserver : IConnectionsObserver
@@ -13,5 +15,13 @@ public void ConnectionAccepted(bool authorized)
         public void ConnectionClosed(bool authorized)
         {
         }
+
+        public void ConnectionAcceptedFor(Uri subscriptionId)
+        {
+        }
+
+        public void ConnectionClosedFor(Uri subscriptionId)
+        {
+        }
     }
 }

source/Halibut/Transport/Protocol/MessageExchangeProtocol.cs

@@ -5,6 +5,7 @@
 using Halibut.Diagnostics;
 using Halibut.Exceptions;
 using Halibut.ServiceModel;
+using Halibut.Transport.Observability;
 
 namespace Halibut.Transport.Protocol
 {
@@ -20,15 +21,17 @@ public class MessageExchangeProtocol
         readonly IMessageExchangeStream stream;
         readonly HalibutTimeoutsAndLimits halibutTimeoutsAndLimits;
         readonly IActiveTcpConnectionsLimiter activeTcpConnectionsLimiter;
+        readonly IConnectionsObserver connectionsObserver;
         readonly ILog log;
         bool identified;
         volatile bool acceptClientRequests = true;
 
-        public MessageExchangeProtocol(IMessageExchangeStream stream, HalibutTimeoutsAndLimits halibutTimeoutsAndLimits, IActiveTcpConnectionsLimiter activeTcpConnectionsLimiter, ILog log)
+        public MessageExchangeProtocol(IMessageExchangeStream stream, HalibutTimeoutsAndLimits halibutTimeoutsAndLimits, IActiveTcpConnectionsLimiter activeTcpConnectionsLimiter, IConnectionsObserver connectionsObserver, ILog log)
         {
             this.stream = stream;
             this.halibutTimeoutsAndLimits = halibutTimeoutsAndLimits;
             this.activeTcpConnectionsLimiter = activeTcpConnectionsLimiter;
+            this.connectionsObserver = connectionsObserver;
             this.log = log;
         }
 
@@ -106,32 +109,30 @@ public async Task ExchangeAsServerAsync(Func<RequestMessage, Task<ResponseMessag
         {
             var identity = await GetRemoteIdentityAsync(cancellationToken);
 
-            //We might need to limit the connection, so by default, we create an unlimited connection lease
-            var limitedConnectionLease = activeTcpConnectionsLimiter.CreateUnlimitedLease();
-
-            //if the remote identity is a subscriber, we might need to limit their active TCP connections
-            if (identity.IdentityType == RemoteIdentityType.Subscriber)
+            switch (identity.IdentityType)
             {
-                limitedConnectionLease = activeTcpConnectionsLimiter.LeaseActiveTcpConnection(identity.SubscriptionId);
-            }
-
-            using (limitedConnectionLease)
-            {
-                await IdentifyAsServerAsync(identity, cancellationToken);
-
-                switch (identity.IdentityType)
-                {
-                    case RemoteIdentityType.Client:
-                        await ProcessClientRequestsAsync(incomingRequestProcessor, cancellationToken);
-                        break;
-                    case RemoteIdentityType.Subscriber:
+                case RemoteIdentityType.Client:
+                    await IdentifyAsServerAsync(identity, cancellationToken);
+                    await ProcessClientRequestsAsync(incomingRequestProcessor, cancellationToken);
+                    break;
+                case RemoteIdentityType.Subscriber:
+                    var limitedConnectionLease = activeTcpConnectionsLimiter.LeaseActiveTcpConnection(identity.SubscriptionId);
+                    try
+                    {
+                        connectionsObserver.ConnectionAcceptedFor(identity.SubscriptionId);
+                        await IdentifyAsServerAsync(identity, cancellationToken);
                         var pendingRequestQueue = pendingRequests(identity);
                         await ProcessSubscriberAsync(pendingRequestQueue, cancellationToken);
                         break;
-                    default:
-                        log.Write(EventType.ErrorInIdentify, $"Remote with identify {identity.SubscriptionId} identified itself with an unknown identity type {identity.IdentityType}");
-                        throw new ProtocolException("Unexpected remote identity: " + identity.IdentityType);
-                }
+                    }
+                    finally
+                    {
+                        connectionsObserver.ConnectionClosedFor(identity.SubscriptionId);
+                        limitedConnectionLease.Dispose();
+                    }
+                default:
+                    log.Write(EventType.ErrorInIdentify, $"Remote with identify {identity.SubscriptionId} identified itself with an unknown identity type {identity.IdentityType}");
+                    throw new ProtocolException("Unexpected remote identity: " + identity.IdentityType);
             }
         }