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

source/Halibut.Tests/Support/TestConnectionsObserver.cs

@@ -10,13 +10,16 @@ public class TestConnectionsObserver : IConnectionsObserver
     {
         readonly ConcurrentBag<bool> connectionAcceptedAuthorized = new();
         readonly ConcurrentBag<bool> connectionClosedAuthorized = new();
-
+        readonly ConcurrentQueue<(Uri SubscriptionId, int PreviousCount, int CurrentCount)> connectionsCountChangedForSubscription = 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 SubscriptionId, int PreviousCount, int CurrentCount)> ConnectionsCountChangedForSubscription => connectionsCountChangedForSubscription.ToList();
+
         public void ConnectionAccepted(bool authorized)
         {
             connectionAcceptedAuthorized.Add(authorized);
@@ -26,5 +29,10 @@ public void ConnectionClosed(bool authorized)
         {
             connectionClosedAuthorized.Add(authorized);
         }
+
+        public void ConnectionsCountChangedFor(Uri subscriptionId, int previousCount, int currentCount)
+        {
+            connectionsCountChangedForSubscription.Enqueue((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/Observability/ConnectionObserverFixture.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using FluentAssertions;
@@ -7,7 +8,6 @@
 using Halibut.Tests.Support.TestAttributes;
 using Halibut.Tests.Support.TestCases;
 using Halibut.Tests.TestServices.Async;
-using Halibut.Tests.TestServices.SyncClientWithOptions;
 using Halibut.TestUtils.Contracts;
 using NUnit.Framework;
 
@@ -51,6 +51,69 @@ public async Task ObserveAuthorizedConnections(ClientAndServiceTestCase clientAn
             connectionsObserver.ConnectionClosedAuthorized.Should().AllSatisfy(a => a.Should().BeTrue());
         }
 
+        [Test]
+        [LatestClientAndLatestServiceTestCases(testNetworkConditions: false, testListening: false)]
+        public async Task ConnectionsCountForAPollingSubscriptionChangesOneConnectionAtATime(ClientAndServiceTestCase clientAndServiceTestCase)
+        {
+            var connectionsObserver = new TestConnectionsObserver();
+            await using (var clientAndService = await clientAndServiceTestCase.CreateTestCaseBuilder()
+                             .WithStandardServices()
+                             .AsLatestClientAndLatestServiceBuilder()
+                             .WithConnectionObserverOnTcpServer(connectionsObserver)
+                             .WithPortForwarding(out var portForwarderRef)
+                             .Build(CancellationToken))
+            {
+                var echo = clientAndService.CreateAsyncClient<IEchoService, IAsyncClientEchoService>();
+                await echo.SayHelloAsync("hello");
+
+                var openFirstConnection = connectionsObserver.ConnectionsCountChangedForSubscription.First();
+                openFirstConnection.PreviousCount.Should().Be(0);
+                openFirstConnection.CurrentCount.Should().Be(1);
+                openFirstConnection.SubscriptionId.Should().Be(clientAndService.ServiceUri);
+
+
+                portForwarderRef.Value.CloseExistingConnections();
+
+                await Try.CatchingError(() => echo.SayHelloAsync("hello"));
+
+                var closeFirstConnection = connectionsObserver.ConnectionsCountChangedForSubscription.Skip(1).First();
+                closeFirstConnection.PreviousCount.Should().Be(1);
+                closeFirstConnection.CurrentCount.Should().Be(0);
+                closeFirstConnection.SubscriptionId.Should().Be(clientAndService.ServiceUri);
+
+                await echo.SayHelloAsync("hello");
+
+                var openSecondConnection = connectionsObserver.ConnectionsCountChangedForSubscription.Skip(2).First();
+                openSecondConnection.PreviousCount.Should().Be(0);
+                openSecondConnection.CurrentCount.Should().Be(1);
+                openSecondConnection.SubscriptionId.Should().Be(clientAndService.ServiceUri);
+            }
+
+            Wait.UntilActionSucceeds(() =>
+            {
+                var closeSecondConnection = connectionsObserver.ConnectionsCountChangedForSubscription.Skip(3).First();
+                closeSecondConnection.PreviousCount.Should().Be(1);
+                closeSecondConnection.CurrentCount.Should().Be(0);
+            }, TimeSpan.FromSeconds(30), Logger, CancellationToken);
+        }
+
+        [Test]
+        [LatestClientAndLatestServiceTestCases(testNetworkConditions: false, testWebSocket: false, testPolling: false)]
+        public async Task ConnectionsCountIsNotChangedForListeningConnections(ClientAndServiceTestCase clientAndServiceTestCase)
+        {
+            var connectionsObserver = new TestConnectionsObserver();
+            await using var clientAndService = await clientAndServiceTestCase.CreateTestCaseBuilder()
+                .WithStandardServices()
+                .AsLatestClientAndLatestServiceBuilder()
+                .WithConnectionObserverOnTcpServer(connectionsObserver)
+                .Build(CancellationToken);
+
+            var echo = clientAndService.CreateAsyncClient<IEchoService, IAsyncClientEchoService>();
+            await echo.SayHelloAsync("hello");
+
+            connectionsObserver.ConnectionsCountChangedForSubscription.Should().BeEmpty("only polling subscriptions lease a counted connection");
+        }
+
         [Test]
         [LatestClientAndLatestServiceTestCases(testNetworkConditions: false, testWebSocket: false, testPolling: false)]
         public async Task ObserveUnauthorizedListeningConnections(ClientAndServiceTestCase clientAndServiceTestCase)

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,101 +1,122 @@
-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;
+
+                var (previousCount, currentCount) = IncrementCount( maximumAcceptedTcpConnectionsPerThumbprint);
+
+                // Calling observer outside the lock might lead to surprising short term values but it will reduce the impact of a slow observer on the connection acceptance.
+                // E.g. (1, 2) might be processed before (0, 1) which will result in (-1, 1) -> (0, 1) values in the bucket.
+                // (-1, 1) should not really last for long as the connections are rather long-lived.
+                connectionsObserver.ConnectionsCountChangedFor(subscriptionId, previousCount, currentCount);
+            }
 
-                lock (this.activeConnectionCountPerSubscriptionId)
+            (int previousCount, int currentCount) IncrementCount(int maximumAcceptedTcpConnectionsPerThumbprint)
+            {
+                lock (activeConnectionCountPerSubscriptionId)
                 {
-                    if (!this.activeConnectionCountPerSubscriptionId.TryGetValue(subscriptionId, out var count))
+                    if (!activeConnectionCountPerSubscriptionId.TryGetValue(subscriptionId, out var count))
                     {
                         count = new StrongBox<int>(0);
-                        this.activeConnectionCountPerSubscriptionId.Add(subscriptionId, count);
+                        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}");
+                        throw new ActiveTcpConnectionsExceededException(subscriptionId, $"Exceeded the maximum number ({maximumAcceptedTcpConnectionsPerThumbprint}) of active TCP connections for subscription {subscriptionId}");
                     }
+
+                    count.Value++;
+
+                    return (previousCount, count.Value);
                 }
             }
 
             public void Dispose()
+            {
+                var counts = DecrementCount();
+                if (counts == null) return;
+
+                var (previousCount, currentCount) = counts.Value;
+                connectionsObserver.ConnectionsCountChangedFor(subscriptionId, previousCount, currentCount);
+            }
+
+            (int previousCount, int currentCount)? DecrementCount()
             {
                 lock (activeConnectionCountPerSubscriptionId)
                 {
                     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);
                         }
+
+                        return (previousCount, count.Value);
                     }
                 }
+                return null;
             }
         }
     }
-}
+}

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);
     }
 }