[Server] Bound the server shutdown drain and make request admission an extension

tests/Opc.Ua.Server.Tests/RequestManagerTests.cs

@@ -268,6 +268,64 @@ public void DisposeCancelsPendingRequests()
             Assert.That(requestLifetime.CancellationToken.IsCancellationRequested, Is.True);
         }
 
+        [Test]
+        [Category("NodeManagerLifecycle")]
+        public void RegisterLifecycleExtensionTwiceReturnsSameExtension()
+        {
+            RequestManagerLifecycleExtension first = RegisterLifecycleExtension();
+            RequestManagerLifecycleExtension second = RegisterLifecycleExtension();
+
+            Assert.That(second, Is.SameAs(first));
+        }
+
+        [Test]
+        [Category("NodeManagerLifecycle")]
+        public void RegisterLifecycleExtensionAfterDisposeThrowsObjectDisposedException()
+        {
+            m_requestManager.Dispose();
+
+            Assert.That(
+                () => m_requestManager.RegisterLifecycleExtension(),
+                Throws.TypeOf<ObjectDisposedException>());
+        }
+
+        [Test]
+        [Category("NodeManagerLifecycle")]
+        public void RequestManagerWithoutLifecycleExtensionAdmitsValidationAndRequests()
+        {
+            using var requestLifetime = new RequestLifetime();
+            OperationContext context = CreateOperationContext(15, requestLifetime);
+
+            using (m_requestManager.EnterValidationScope())
+            {
+                Assert.DoesNotThrow(() => m_requestManager.RequestReceived(context));
+            }
+
+            m_requestManager.RequestCompleted(context);
+            Assert.That(requestLifetime.TryCancel(StatusCodes.BadTimeout), Is.False);
+        }
+
+        [Test]
+        [Category("NodeManagerLifecycle")]
+        public async Task RequestManagerWithoutLifecycleExtensionDoesNotRepeatDrainAsync()
+        {
+            using var requestLifetime = new RequestLifetime();
+            OperationContext context = CreateOperationContext(16, requestLifetime);
+            IDisposable validationScope = m_requestManager.EnterValidationScope();
+
+            Task drain = m_requestManager.WaitForCurrentRequestsAsync().AsTask();
+            Assert.That(drain.IsCompleted, Is.False);
+
+            m_requestManager.RequestReceived(context);
+            validationScope.Dispose();
+
+            await AssertCompletesWithinTimeoutAsync(drain).ConfigureAwait(false);
+            Assert.That(
+                requestLifetime.TryCancel(StatusCodes.BadTimeout),
+                Is.True,
+                "Without the lifecycle extension, the drain must not resnapshot promoted requests.");
+        }
+
         [Test]
         [Category("NodeManagerLifecycle")]
         public async Task WaitForCurrentRequestsAsyncCompletesAfterAllSnapshotRequestsCompleteAsync()
@@ -363,6 +421,93 @@ public async Task WaitForCurrentRequestsAsyncWithNoRequestsCompletesImmediatelyA
             await waiter.ConfigureAwait(false);
         }
 
+        [Test]
+        [Category("NodeManagerLifecycle")]
+        public async Task WaitForCurrentRequestsAsyncExcludesEveryLifecycleWaiterAsync()
+        {
+            RequestManagerLifecycleExtension extension = RegisterLifecycleExtension();
+            using var requestLifetimeA = new RequestLifetime();
+            using var requestLifetimeB = new RequestLifetime();
+            using var requestLifetimeC = new RequestLifetime();
+            OperationContext contextA = CreateOperationContext(10, requestLifetimeA);
+            OperationContext contextB = CreateOperationContext(11, requestLifetimeB);
+            OperationContext contextC = CreateOperationContext(12, requestLifetimeC);
+
+            using IDisposable requestScopeA = m_requestManager.EnterRequestScope(contextA);
+            using RequestManagerLifecycleExtension.RequestLifecycleWaiterScope waiterScopeA =
+                extension.EnterLifecycleWaiter();
+            waiterScopeA.MarkSemaphoreWaitStarted();
+            using IDisposable requestScopeB = m_requestManager.EnterRequestScope(contextB);
+            using RequestManagerLifecycleExtension.RequestLifecycleWaiterScope waiterScopeB =
+                extension.EnterLifecycleWaiter();
+            waiterScopeB.MarkSemaphoreWaitStarted();
+            m_requestManager.RequestReceived(contextC);
+
+            Task drain = m_requestManager.WaitForCurrentRequestsAsync().AsTask();
+            Assert.That(drain.IsCompleted, Is.False);
+
+            m_requestManager.RequestCompleted(contextC);
+
+            await AssertCompletesWithinTimeoutAsync(drain).ConfigureAwait(false);
+            Assert.That(
+                requestLifetimeA.TryCancel(StatusCodes.BadTimeout),
+                Is.True,
+                "The first excluded request must still be executing.");
+            Assert.That(
+                requestLifetimeB.TryCancel(StatusCodes.BadTimeout),
+                Is.True,
+                "The second excluded request must still be executing.");
+        }
+
+        [Test]
+        [Category("NodeManagerLifecycle")]
+        public async Task SemaphoreWaitStartReleasesAlreadyActiveDrainAsync()
+        {
+            RequestManagerLifecycleExtension extension = RegisterLifecycleExtension();
+            using var requestLifetime = new RequestLifetime();
+            OperationContext context = CreateOperationContext(13, requestLifetime);
+            using IDisposable requestScope = m_requestManager.EnterRequestScope(context);
+
+            Task drain = m_requestManager.WaitForCurrentRequestsAsync().AsTask();
+            Assert.That(drain.IsCompleted, Is.False);
+
+            using RequestManagerLifecycleExtension.RequestLifecycleWaiterScope waiterScope =
+                extension.EnterLifecycleWaiter();
+
+            Assert.That(
+                drain.IsCompleted,
+                Is.False,
+                "Registration must not release a drain before the semaphore wait is queued.");
+
+            waiterScope.MarkSemaphoreWaitStarted();
+
+            await AssertCompletesWithinTimeoutAsync(drain).ConfigureAwait(false);
+            Assert.That(
+                requestLifetime.TryCancel(StatusCodes.BadTimeout),
+                Is.True,
+                "Waiting for the lifecycle semaphore must not complete the request itself.");
+        }
+
+        [Test]
+        [Category("NodeManagerLifecycle")]
+        public async Task DisposedLifecycleWaiterIsIncludedInNextDrainAsync()
+        {
+            RequestManagerLifecycleExtension extension = RegisterLifecycleExtension();
+            using var requestLifetime = new RequestLifetime();
+            OperationContext context = CreateOperationContext(14, requestLifetime);
+            IDisposable requestScope = m_requestManager.EnterRequestScope(context);
+            using (extension.EnterLifecycleWaiter())
+            {
+            }
+
+            Task drain = m_requestManager.WaitForCurrentRequestsAsync().AsTask();
+            Assert.That(drain.IsCompleted, Is.False);
+
+            requestScope.Dispose();
+
+            await AssertCompletesWithinTimeoutAsync(drain).ConfigureAwait(false);
+        }
+
         [Test]
         public void RequestReceivedCalledTwiceWithSameContextIsIdempotent()
         {
@@ -435,6 +580,58 @@ async Task HandleAsync()
             });
         }
 
+        [Test]
+        [Category("NodeManagerLifecycle")]
+        public async Task ServiceDispatchScopeReceivesRequestIdentityFromAsyncValidationAsync()
+        {
+            RequestManagerLifecycleExtension extension = RegisterLifecycleExtension();
+            using var requestLifetime = new RequestLifetime();
+            OperationContext context = CreateOperationContext(61, requestLifetime);
+
+            using (m_requestManager.EnterServiceDispatchScope())
+            {
+                await RegisterAsync().ConfigureAwait(false);
+                using RequestManagerLifecycleExtension.RequestLifecycleWaiterScope waiter =
+                    extension.EnterLifecycleWaiter();
+            }
+
+            m_requestManager.RequestCompleted(context);
+            Assert.That(requestLifetime.TryCancel(StatusCodes.BadTimeout), Is.False);
+
+            async Task RegisterAsync()
+            {
+                await Task.Yield();
+                m_requestManager.RequestReceived(context);
+            }
+        }
+
+        [Test]
+        [Category("NodeManagerLifecycle")]
+        public void NestedRequestRegistrationPreservesOuterDispatchIdentity()
+        {
+            RequestManagerLifecycleExtension extension = RegisterLifecycleExtension();
+            using var outerLifetime = new RequestLifetime();
+            using var nestedLifetime = new RequestLifetime();
+            OperationContext outer = CreateOperationContext(62, outerLifetime);
+            OperationContext nested = CreateOperationContext(63, nestedLifetime);
+
+            using (m_requestManager.EnterServiceDispatchScope())
+            using (m_requestManager.EnterRequestScope(outer))
+            {
+                m_requestManager.RequestReceived(nested);
+                m_requestManager.RequestCompleted(nested);
+
+                using RequestManagerLifecycleExtension.RequestLifecycleWaiterScope waiter =
+                    extension.EnterLifecycleWaiter();
+            }
+
+            Assert.Multiple(() =>
+            {
+                Assert.That(outerLifetime.TryCancel(StatusCodes.BadTimeout), Is.False);
+                Assert.That(nestedLifetime.TryCancel(StatusCodes.BadTimeout), Is.False);
+            });
+        }
+
         [Test]
         public void EnterRequestScopeDoesNotMarkTheFlowAsExecutingARequest()
         {
@@ -568,13 +765,97 @@ public void NestedValidationScopesLeaveRegisteredRequestsToTheirOwnScope()
             m_requestManager.RequestCompleted(innerContext);
         }
 
+        [Test]
+        public void CloseAdmissionRejectsNewValidationAndDirectRequests()
+        {
+            RequestManagerLifecycleExtension extension = RegisterLifecycleExtension();
+            using var requestLifetime = new RequestLifetime();
+            OperationContext context = CreateOperationContext(90, requestLifetime);
+
+            extension.CloseAdmission();
+
+            ServiceResultException validationException =
+                Assert.Throws<ServiceResultException>(
+                    () => m_requestManager.EnterValidationScope());
+            ServiceResultException requestException =
+                Assert.Throws<ServiceResultException>(
+                    () => m_requestManager.RequestReceived(context));
+
+            Assert.Multiple(() =>
+            {
+                Assert.That(
+                    validationException.StatusCode,
+                    Is.EqualTo(StatusCodes.BadServerHalted));
+                Assert.That(
+                    requestException.StatusCode,
+                    Is.EqualTo(StatusCodes.BadServerHalted));
+                Assert.That(
+                    requestLifetime.TryCancel(StatusCodes.BadTimeout),
+                    Is.True);
+            });
+        }
+
+        [Test]
+        public void PreAdmittedValidationCanPromoteRequestAfterAdmissionCloses()
+        {
+            RequestManagerLifecycleExtension extension = RegisterLifecycleExtension();
+            using var requestLifetime = new RequestLifetime();
+            OperationContext context = CreateOperationContext(91, requestLifetime);
+            IDisposable validationScope = m_requestManager.EnterValidationScope();
+
+            extension.CloseAdmission();
+
+            using (m_requestManager.EnterServiceDispatchScope())
+            using (m_requestManager.EnterRequestScope(context))
+            {
+                validationScope.Dispose();
+                Assert.That(m_requestManager.IsExecutingRequest, Is.True);
+            }
+
+            Assert.That(requestLifetime.TryCancel(StatusCodes.BadTimeout), Is.False);
+        }
+
+        [Test]
+        [Category("NodeManagerLifecycle")]
+        public async Task ClosedAdmissionDrainTracksRequestRegisteredAfterInitialSnapshotAsync()
+        {
+            RequestManagerLifecycleExtension extension = RegisterLifecycleExtension();
+            using var requestLifetime = new RequestLifetime();
+            OperationContext context = CreateOperationContext(92, requestLifetime);
+            IDisposable validationScope = m_requestManager.EnterValidationScope();
+            extension.CloseAdmission();
+
+            Task drain = m_requestManager.WaitForCurrentRequestsAsync().AsTask();
+            Assert.That(drain.IsCompleted, Is.False);
+            m_requestManager.RequestReceived(context);
+            m_requestManager.RequestReceived(context);
+            validationScope.Dispose();
+
+            Task completed = await Task.WhenAny(
+                drain,
+                Task.Delay(TimeSpan.FromMilliseconds(250))).ConfigureAwait(false);
+            Assert.That(
+                completed,
+                Is.Not.SameAs(drain),
+                "The closed-admission drain must resnapshot the request promoted by validation.");
+
+            m_requestManager.RequestCompleted(context);
+            await AssertCompletesWithinTimeoutAsync(drain).ConfigureAwait(false);
+            Assert.That(requestLifetime.TryCancel(StatusCodes.BadTimeout), Is.False);
+        }
+
         private static OperationContext CreateOperationContext(
             uint requestHandle,
             RequestLifetime requestLifetime)
         {
             return CreateOperationContext(requestHandle, requestLifetime, 0);
         }
 
+        private RequestManagerLifecycleExtension RegisterLifecycleExtension()
+        {
+            return m_requestManager.RegisterLifecycleExtension();
+        }
+
         private static OperationContext CreateOperationContext(
             uint requestHandle,
             RequestLifetime requestLifetime,

tests/Opc.Ua.Server.Tests/StandardServerCoverageTests.cs

@@ -111,6 +111,16 @@ public void ConstructorUsesProvidedTimeProvider()
             Assert.That(server.TimeProviderAccessor, Is.SameAs(custom));
         }
 
+        [Test]
+        public void DisposeCanBeCalledMoreThanOnce()
+        {
+            TestableStandardServer server = CreateServer();
+
+            server.Dispose();
+
+            Assert.DoesNotThrow(server.Dispose);
+        }
+
         [Test]
         public void LoadComplexTypesDefaultsToTrue()
         {