Fix the classic engine's publish pipeline: refill after keep-alive recovery, and bound every producer

src/Opc.Ua.Client/Session/Subscription/ClassicSubscriptionEngine.cs

@@ -134,14 +134,11 @@ public void StartPublishing(int timeout, bool fullQueue)
                 return;
             }
 
-            // Refill the pipeline up to the desired count. Every request is
-            // bounded by the in flight reservation: StartPublishing runs once
-            // per subscription as the subscriptions are created, so an
-            // unbounded send stacks a pipeline worth of requests on every call
-            // and the outstanding count grows far past the desired count. A
-            // drained pipeline always has room, so this still refills a
-            // pipeline whose requests are outstanding but are no longer
-            // expected to return.
+            // Refill the pipeline up to the desired count. The reservation is
+            // the real bound: it counts the requests the session still expects
+            // plus the ones sent but not yet recorded, so this both refuses to
+            // overshoot and still refills a pipeline whose requests the session
+            // has written off.
             int startCount = fullQueue
                 ? 0
                 : GoodPublishRequestCount;
@@ -154,7 +151,7 @@ public void StartPublishing(int timeout, bool fullQueue)
 
                 if (!BeginPublishCore(timeout))
                 {
-                    Interlocked.Decrement(ref m_publishRequestsInFlight);
+                    Interlocked.Decrement(ref m_unrecordedPublishRequests);
                     break;
                 }
             }
@@ -225,48 +222,55 @@ protected virtual void Dispose(bool disposing)
         /// <summary>
         /// Sends a publish request to the server.
         /// </summary>
+        /// <remarks>
+        /// This is the recovery nudge behind
+        /// <see cref="Subscription.HandleOnKeepAliveStopped"/>: a subscription
+        /// that has seen no notification asks for one more publish so the
+        /// server has a request to answer. It is bounded by the same
+        /// reservation as the automatic top up, because sending past the
+        /// desired count cannot help - the server already holds that many
+        /// requests and answers the surplus with
+        /// <see cref="StatusCodes.BadTooManyPublishRequests"/> - while every
+        /// subscription firing this nudge at once would otherwise multiply the
+        /// pipeline by the number of subscriptions. A drained pipeline always
+        /// has room, so the nudge still gets through when it is the outstanding
+        /// requests that have stopped coming back.
+        /// </remarks>
         /// <param name="timeout">The timeout for publish requests
         /// in milliseconds.</param>
-        /// <returns>True if the request was sent successfully.</returns>
+        /// <returns>
+        /// True if the pipeline holds the requested publish request when this
+        /// returns, whether it was sent here or was already outstanding.
+        /// </returns>
         internal bool BeginPublish(int timeout)
         {
-            // An explicit request to send a publish (e.g. keep-alive recovery).
-            // Respect the server-side limit when it has been learned from a
-            // BadTooManyPublishRequests response: if the pipeline is already at
-            // or above the desired count, delegate to the ordinary top-up path
-            // so the reservation check prevents overshooting.
-            int desiredCount = GetDesiredPublishRequestCount(false);
-
-            if (desiredCount > 0 &&
-                !TryReservePublishRequest(desiredCount, out _))
-            {
-                // Pipeline is already at capacity; queue a top-up instead so
-                // the next completed request will refill the slot naturally.
-                QueueBeginPublish();
-                return true;
-            }
+            // At least one, so an empty pipeline is always refillable even
+            // when no subscription has been created yet.
+            int limit = Math.Max(1, GetDesiredPublishRequestCount(false));
 
-            // Either no desired count yet (engine not yet started) or we
-            // successfully reserved a slot.  Proceed unconditionally so that
-            // a stalled pipeline can always be kick-started.
-            if (desiredCount == 0)
+            if (!TryReservePublishRequest(limit, out _))
             {
-                Interlocked.Increment(ref m_publishRequestsInFlight);
+                // The pipeline already holds the requests this nudge asks for,
+                // so the caller's intent is met without sending. Sending anyway
+                // cannot help: the server holds that many requests already and
+                // answers the surplus with BadTooManyPublishRequests.
+                return true;
             }
 
             if (!BeginPublishCore(timeout))
             {
-                Interlocked.Decrement(ref m_publishRequestsInFlight);
+                Interlocked.Decrement(ref m_unrecordedPublishRequests);
                 return false;
             }
 
             return true;
         }
 
         /// <summary>
-        /// Sends a publish request. The caller must have accounted for the
-        /// request in <see cref="m_publishRequestsInFlight"/> and must release
-        /// that reservation if this returns <c>false</c>.
+        /// Sends a publish request. The caller must have reserved the request
+        /// in <see cref="m_unrecordedPublishRequests"/>. The reservation is
+        /// released here once the session has recorded the request, or by the
+        /// caller if this returns <c>false</c>.
         /// </summary>
         private bool BeginPublishCore(int timeout)
         {
@@ -349,6 +353,13 @@ private bool BeginPublishCore(int timeout)
                     activity,
                     requestHeader.RequestHandle,
                     DataTypes.PublishRequest);
+
+                // The session now counts this request, so the reservation that
+                // stood in for it is released. Releasing after the record is
+                // published keeps the total conservative: the request is
+                // briefly counted twice rather than not at all.
+                Interlocked.Decrement(ref m_unrecordedPublishRequests);
+
                 task.ConfigureAwait(false)
                     .GetAwaiter()
                     .OnCompleted(() => OnPublishComplete(
@@ -383,14 +394,6 @@ private void OnPublishComplete(
                 requestHeader.RequestHandle,
                 DataTypes.PublishRequest);
 
-            // Release the reservation only once the session has retired the
-            // request from its outstanding list. Releasing it when the task
-            // completed would open a window in which a concurrent top up sees
-            // a free slot while GoodPublishRequestCount still counts this
-            // request, so the pipeline drifts above the limit by the number of
-            // completions that overlap.
-            Interlocked.Decrement(ref m_publishRequestsInFlight);
-
             m_eventLogger.ClientEventPublishStop(
                 (int)requestHeader.RequestHandle,
                 sessionId);
@@ -877,39 +880,56 @@ private void QueueBeginPublish()
 
             if (!BeginPublishCore(m_context.OperationTimeout))
             {
-                Interlocked.Decrement(ref m_publishRequestsInFlight);
+                Interlocked.Decrement(ref m_unrecordedPublishRequests);
             }
         }
 
         /// <summary>
         /// Atomically reserves capacity for one more publish request if fewer
-        /// than <paramref name="limit"/> are in flight.
+        /// than <paramref name="limit"/> are outstanding.
         /// </summary>
-        /// <param name="limit">The maximum number of requests in flight.</param>
-        /// <param name="inFlight">The number observed in flight.</param>
+        /// <remarks>
+        /// Outstanding means the requests the session still expects to return
+        /// (<see cref="ISubscriptionEngineContext.GoodPublishRequestCount"/>)
+        /// plus the requests this engine has sent that the session has not
+        /// recorded yet. The session's count is authoritative: it drops a
+        /// request as soon as the session writes it off, which is what
+        /// <see cref="Session.OnKeepAlive"/> does to the whole pipeline when
+        /// keep alives recover. Counting only what this engine has sent would
+        /// hold those write offs forever and the pipeline would never refill.
+        /// The unrecorded count closes the opposite gap: the session only
+        /// counts a request once <c>AsyncRequestStarted</c> has recorded it,
+        /// which happens after the request was issued, so concurrent callers
+        /// would otherwise all read the same lagging value and each send.
+        /// </remarks>
+        /// <param name="limit">The maximum number of outstanding requests.</param>
+        /// <param name="outstanding">The number observed outstanding.</param>
         /// <returns>True if a slot was reserved.</returns>
-        private bool TryReservePublishRequest(int limit, out int inFlight)
+        private bool TryReservePublishRequest(int limit, out int outstanding)
         {
-            int current = Volatile.Read(ref m_publishRequestsInFlight);
-
-            while (current < limit)
+            while (true)
             {
-                int prior = Interlocked.CompareExchange(
-                    ref m_publishRequestsInFlight,
-                    current + 1,
-                    current);
+                // Read the unrecorded count first. A request that moves from
+                // unrecorded to recorded between the two reads is then counted
+                // twice rather than missed, which errs towards sending less.
+                int unrecorded = Volatile.Read(ref m_unrecordedPublishRequests);
+                int current = m_context.GoodPublishRequestCount + unrecorded;
 
-                if (prior == current)
+                if (current >= limit)
                 {
-                    inFlight = current;
-                    return true;
+                    outstanding = current;
+                    return false;
                 }
 
-                current = prior;
+                if (Interlocked.CompareExchange(
+                        ref m_unrecordedPublishRequests,
+                        unrecorded + 1,
+                        unrecorded) == unrecorded)
+                {
+                    outstanding = current;
+                    return true;
+                }
             }
-
-            inFlight = current;
-            return false;
         }
 
         /// <summary>
@@ -1150,7 +1170,7 @@ private bool BelowPublishRequestLimit(int requestCount)
         private readonly Lock m_acknowledgementsToSendLock = new();
         private List<SubscriptionAcknowledgement> m_acknowledgementsToSend = [];
         internal uint PublishCounter;
-        private int m_publishRequestsInFlight;
+        private int m_unrecordedPublishRequests;
         private int m_tooManyPublishRequests;
         private int m_minPublishRequestCount;
         private int m_maxPublishRequestCount;