Add ability to observe the number of connections halibut opens for ea…

source/Halibut.Tests/Support/TestConnectionsObserver.cs

@@ -10,7 +10,8 @@ public class TestConnectionsObserver : IConnectionsObserver
     {
         readonly ConcurrentBag<bool> connectionAcceptedAuthorized = new();
         readonly ConcurrentBag<bool> connectionClosedAuthorized = new();
-
+        readonly ConcurrentBag<(Uri SubscriptionId, int PreviousCount, int CurrentCount)> connectionsCountChangedForSubscription = new();
+
         public long ConnectionAcceptedCount => connectionAcceptedAuthorized.Count;
         public long ConnectionClosedCount => connectionClosedAuthorized.Count;
 
@@ -26,5 +27,10 @@ public void ConnectionClosed(bool authorized)
         {
             connectionClosedAuthorized.Add(authorized);
         }
+
+        public void ConnectionsCountChangedFor(Uri subscriptionId, int previousCount, int currentCount)
+        {
+            connectionsCountChangedForSubscription.Add((subscriptionId, previousCount, currentCount));
+        }
     }
 }

source/Halibut.Tests/Transport/ActiveTcpConnectionsLimiterFixture.cs

@@ -6,6 +6,7 @@
 using Halibut.Diagnostics;
 using Halibut.Exceptions;
 using Halibut.Transport;
+using Halibut.Transport.Observability;
 using NUnit.Framework;
 
 namespace Halibut.Tests.Transport
@@ -22,7 +23,7 @@ public void LimitsConcurrentConnectionsForSingleSubscription()
             var limiter = new ActiveTcpConnectionsLimiter(new HalibutTimeoutsAndLimits
             {
                 MaximumActiveTcpConnectionsPerPollingSubscription = limit
-            });
+            }, NoOpConnectionsObserver.Instance);
 
             // Act
             //we create a new URI each time to make sure we aren't doing object reference checks
@@ -46,7 +47,7 @@ public void CompletedLeasesAreRemovedFromTheCount()
             var limiter = new ActiveTcpConnectionsLimiter(new HalibutTimeoutsAndLimits
             {
                 MaximumActiveTcpConnectionsPerPollingSubscription = limit
-            });
+            }, NoOpConnectionsObserver.Instance);
 
             // Act
             limiter.LeaseActiveTcpConnection(subscription);
@@ -76,7 +77,7 @@ public void DoesNotLimitConcurrentConnectionsForDifferentSubscriptions()
             var limiter = new ActiveTcpConnectionsLimiter(new HalibutTimeoutsAndLimits
             {
                 MaximumActiveTcpConnectionsPerPollingSubscription = limit
-            });
+            }, NoOpConnectionsObserver.Instance);
 
             // Act
             limiter.LeaseActiveTcpConnection(subscription1);
@@ -99,7 +100,7 @@ public async Task ShouldHandleMultiThreading()
             var limiter = new ActiveTcpConnectionsLimiter(new HalibutTimeoutsAndLimits
             {
                 MaximumActiveTcpConnectionsPerPollingSubscription = limit
-            });
+            }, NoOpConnectionsObserver.Instance);
 
             // Capture how many claims fail with the exception
             var failures = 0;
@@ -140,7 +141,7 @@ public async Task ShouldHandleMultiThreadingWithFakeWorkDuringLease()
             var limiter = new ActiveTcpConnectionsLimiter(new HalibutTimeoutsAndLimits
             {
                 MaximumActiveTcpConnectionsPerPollingSubscription = limit
-            });
+            }, NoOpConnectionsObserver.Instance);
 
             // Capture how many claims fail with the exception
             var failures = 0;

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

@@ -26,7 +26,7 @@ public void SetUp()
             stream = new DumpStream();
             stream.SetRemoteIdentity(new RemoteIdentity(RemoteIdentityType.Server));
             var limits = new HalibutTimeoutsAndLimitsForTestsBuilder().Build();
-            var activeConnectionsLimiter = new ActiveTcpConnectionsLimiter(limits);
+            var activeConnectionsLimiter = new ActiveTcpConnectionsLimiter(limits, NoOpConnectionsObserver.Instance);
             protocol = new MessageExchangeProtocol(stream, new HalibutTimeoutsAndLimitsForTestsBuilder().Build(), activeConnectionsLimiter, Substitute.For<ILog>());
         }
 

source/Halibut.Tests/Transport/SecureClientFixture.cs

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

source/Halibut/HalibutRuntime.cs

@@ -82,7 +82,7 @@ ISecureConnectionObserver secureConnectionObserver
 
             connectionManager = new ConnectionManagerAsync();
             tcpConnectionFactory = new TcpConnectionFactory(serverCertificate, TimeoutsAndLimits, streamFactory, secureConnectionObserver);
-            activeTcpConnectionsLimiter = new ActiveTcpConnectionsLimiter(TimeoutsAndLimits);
+            activeTcpConnectionsLimiter = new ActiveTcpConnectionsLimiter(TimeoutsAndLimits, connectionsObserver);
         }
 
         public ILogFactory Logs => logs;

source/Halibut/Transport/ActiveTcpConnectionsLimiter.cs

@@ -1,61 +1,58 @@
-using System;
+using System;
 using System.Collections.Generic;
 using System.Runtime.CompilerServices;
 using Halibut.Diagnostics;
 using Halibut.Exceptions;
+using Halibut.Transport.Observability;
 
 namespace Halibut.Transport
 {
     public interface IActiveTcpConnectionsLimiter
     {
         IDisposable LeaseActiveTcpConnection(Uri subscriptionId);
-
-        IDisposable CreateUnlimitedLease();
     }
 
     public class ActiveTcpConnectionsLimiter : IActiveTcpConnectionsLimiter
     {
         readonly HalibutTimeoutsAndLimits timeoutsAndLimits;
+        readonly IConnectionsObserver connectionsObserver;
 
         Dictionary<Uri, StrongBox<int>> activeConnectionCountPerSubscriptionId = new();
 
-        public ActiveTcpConnectionsLimiter(HalibutTimeoutsAndLimits timeoutsAndLimits)
+        public ActiveTcpConnectionsLimiter(HalibutTimeoutsAndLimits timeoutsAndLimits, IConnectionsObserver connectionsObserver)
         {
             this.timeoutsAndLimits = timeoutsAndLimits;
+            this.connectionsObserver = connectionsObserver;
         }
 
         public IDisposable LeaseActiveTcpConnection(Uri subscriptionId)
         {
-            //if there is no limit, then we return a NoOp lease (which doesn't limit anything)
+            //if there is no limit, then we still count the connection (the observer is told about every
+            //connection either way), we just never reject it
             if (!timeoutsAndLimits.MaximumActiveTcpConnectionsPerPollingSubscription.HasValue)
             {
-                return CreateUnlimitedLease();
+                return CreateUnlimitedLease(subscriptionId);
             }
 
-            return new LimitingAuthorizedTcpConnectionLease(subscriptionId, activeConnectionCountPerSubscriptionId, timeoutsAndLimits.MaximumActiveTcpConnectionsPerPollingSubscription.Value);
-        }
-
-        public IDisposable CreateUnlimitedLease()
-        {
-            return new UnlimitedAuthorizedTcpConnectionLease();
+            return new LimitingAuthorizedTcpConnectionLease(subscriptionId, activeConnectionCountPerSubscriptionId, timeoutsAndLimits.MaximumActiveTcpConnectionsPerPollingSubscription.Value, connectionsObserver);
         }
 
-        class UnlimitedAuthorizedTcpConnectionLease : IDisposable
+        IDisposable CreateUnlimitedLease(Uri subscriptionId)
         {
-            public void Dispose()
-            {
-            }
+            return new LimitingAuthorizedTcpConnectionLease(subscriptionId, activeConnectionCountPerSubscriptionId, int.MaxValue, connectionsObserver);
         }
 
         class LimitingAuthorizedTcpConnectionLease : IDisposable
         {
             readonly Uri subscriptionId;
             readonly Dictionary<Uri, StrongBox<int>> activeConnectionCountPerSubscriptionId;
+            readonly IConnectionsObserver connectionsObserver;
 
-            public LimitingAuthorizedTcpConnectionLease(Uri subscriptionId, Dictionary<Uri, StrongBox<int>> activeConnectionCountPerSubscriptionId, int maximumAcceptedTcpConnectionsPerThumbprint)
+            public LimitingAuthorizedTcpConnectionLease(Uri subscriptionId, Dictionary<Uri, StrongBox<int>> activeConnectionCountPerSubscriptionId, int maximumAcceptedTcpConnectionsPerThumbprint, IConnectionsObserver connectionsObserver)
             {
                 this.subscriptionId = subscriptionId;
                 this.activeConnectionCountPerSubscriptionId = activeConnectionCountPerSubscriptionId;
+                this.connectionsObserver = connectionsObserver;
 
                 lock (this.activeConnectionCountPerSubscriptionId)
                 {
@@ -65,17 +62,18 @@ public LimitingAuthorizedTcpConnectionLease(Uri subscriptionId, Dictionary<Uri,
                         this.activeConnectionCountPerSubscriptionId.Add(subscriptionId, count);
                     }
 
-                    count.Value++;
+                    var previousCount = count.Value;
 
                     //validate the new count. If this throws an exception, it'll kill the connection
-                    if (count.Value > maximumAcceptedTcpConnectionsPerThumbprint)
+                    if (count.Value + 1 > maximumAcceptedTcpConnectionsPerThumbprint)
                     {
-                        //decrement as this connection has been rejected
-                        count.Value--;
-
                         //throw an exception, bailing on the connection
                         throw new ActiveTcpConnectionsExceededException(this.subscriptionId, $"Exceeded the maximum number ({maximumAcceptedTcpConnectionsPerThumbprint}) of active TCP connections for subscription {subscriptionId}");
                     }
+
+                    count.Value++;
+
+                    connectionsObserver.ConnectionsCountChangedFor(subscriptionId, previousCount, count.Value);
                 }
             }
 
@@ -86,16 +84,19 @@ public void Dispose()
                     if (activeConnectionCountPerSubscriptionId.TryGetValue(subscriptionId, out var count))
                     {
                         //decrement the count of authorized connections
+                        var previousCount = count.Value;
                         count.Value--;
 
                         // Remove the key from the dictionary if the value is 0
                         if (count.Value == 0)
                         {
                             activeConnectionCountPerSubscriptionId.Remove(subscriptionId);
                         }
+
+                        connectionsObserver.ConnectionsCountChangedFor(subscriptionId, previousCount, count.Value);
                     }
                 }
             }
         }
     }
-}
+}

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,20 @@ 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>
+        /// A polling subscriber's connections' count has changed
+        /// </summary>
+        /// <param name="subscriptionId">The polling subscriber's subscription id.</param>
+        /// <param name="previousCount">
+        /// The number of active TCP connections for this subscriptionId immediately before the change
+        /// </param>
+        /// <param name="currentCount">
+        /// The number of active TCP connections for this subscriptionId immediately after the change
+        /// </param>
+        public void ConnectionsCountChangedFor(Uri subscriptionId, int previousCount, int currentCount);
     }
 }

source/Halibut/Transport/Observability/NoOpConnectionsObserver.cs

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

source/Halibut/Transport/Protocol/MessageExchangeProtocol.cs

@@ -106,32 +106,23 @@ 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:
+                    using (activeTcpConnectionsLimiter.LeaseActiveTcpConnection(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);
-                }
+                    }
+                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);
             }
         }