gh-56 fix export pipeline stuck after Transform block

mocsharp · mocsharp · commit b10600eded48 · 2022-03-10T09:56:20.000-08:00
diff --git a/src/Configuration/MessageBrokerConfigurationKeys.cs b/src/Configuration/MessageBrokerConfigurationKeys.cs
@@ -35,12 +35,5 @@ public class MessageBrokerConfigurationKeys
         /// </summary>
         [JsonProperty(PropertyName = "exportRequestPrefix")]
         public string ExportRequestPrefix { get; set; } = "md.export.request";
-
-        /// <summary>
-        /// Gets or sets the name of the request queue.
-        /// Defaults to `export_tasks`
-        /// </summary>
-        [JsonProperty(PropertyName = "exportRequestQueue")]
-        public string ExportRequestQueue { get; set; } = "export_tasks";
     }
 }
diff --git a/src/InformaticsGateway/Program.cs b/src/InformaticsGateway/Program.cs
@@ -95,6 +95,7 @@ public static IHostBuilder CreateHostBuilder(string[] args) =>
                     services.AddTransient<IDicomToolkit, DicomToolkit>();
 
                     services.AddScoped(typeof(IInformaticsGatewayRepository<>), typeof(InformaticsGatewayRepository<>));
+                    services.AddScoped<IInferenceRequestRepository, InferenceRequestRepository>();
 
                     services.AddSingleton<MinIoStorageService>();
                     services.AddSingleton<IStorageService>(implementationFactory =>
diff --git a/src/InformaticsGateway/Services/Export/ExportServiceBase.cs b/src/InformaticsGateway/Services/Export/ExportServiceBase.cs
@@ -54,8 +54,9 @@ public abstract class ExportServiceBase : IHostedService, IMonaiService
     {
         private static readonly object SyncRoot = new();
 
-        internal event EventHandler ReportActionStarted;
+        internal event EventHandler ReportActionCompleted;
 
+        private readonly CancellationTokenSource _cancellationTokenSource;
         private readonly ILogger _logger;
         private readonly IServiceScopeFactory _serviceScopeFactory;
         private readonly IStorageInfoProvider _storageInfoProvider;
@@ -64,8 +65,6 @@ public abstract class ExportServiceBase : IHostedService, IMonaiService
         private readonly IMessageBrokerPublisherService _messagePublisher;
         private readonly IServiceScope _scope;
         private readonly Dictionary<string, ExportRequestMessage> _exportRequests;
-        private readonly string _exportQueueName;
-        private TransformManyBlock<ExportRequestMessage, ExportRequestDataMessage> _exportFlow;
 
         public abstract string RoutingKey { get; }
         protected abstract int Concurrency { get; }
@@ -87,6 +86,7 @@ public ExportServiceBase(
             IServiceScopeFactory serviceScopeFactory,
             IStorageInfoProvider storageInfoProvider)
         {
+            _cancellationTokenSource = new CancellationTokenSource();
             _logger = logger ?? throw new ArgumentNullException(nameof(logger));
             _serviceScopeFactory = serviceScopeFactory ?? throw new ArgumentNullException(nameof(serviceScopeFactory));
             _scope = _serviceScopeFactory.CreateScope();
@@ -103,12 +103,10 @@ public ExportServiceBase(
             _messagePublisher = _scope.ServiceProvider.GetRequiredService<IMessageBrokerPublisherService>();
 
             _exportRequests = new Dictionary<string, ExportRequestMessage>();
-            _exportQueueName = _configuration.Messaging.Topics.ExportRequestQueue;
         }
 
         public Task StartAsync(CancellationToken cancellationToken)
         {
-            SetupExportFlow(cancellationToken);
             SetupPolling();
 
             Status = ServiceStatus.Running;
@@ -118,44 +116,73 @@ public Task StartAsync(CancellationToken cancellationToken)
 
         public Task StopAsync(CancellationToken cancellationToken)
         {
+            _cancellationTokenSource.Cancel();
             _logger.LogInformation($"{ServiceName} is stopping.");
-            _exportFlow.Complete();
-            _exportFlow.Completion.Wait(cancellationToken);
             Status = ServiceStatus.Stopped;
             return Task.CompletedTask;
         }
 
-        private void SetupExportFlow(CancellationToken cancellationToken)
+        private void SetupPolling()
         {
+            _messageSubscriber.Subscribe(RoutingKey, String.Empty, OnMessageReceivedCallback);
+            _logger.Log(LogLevel.Information, $"{ServiceName} subscribed to {RoutingKey} messages.");
+        }
+
+        private void OnMessageReceivedCallback(MessageReceivedEventArgs eventArgs)
+        {
+            if (!_storageInfoProvider.HasSpaceAvailableForExport)
+            {
+                _logger.Log(LogLevel.Warning, $"Export service paused due to insufficient storage space.  Available storage space: {_storageInfoProvider.AvailableFreeSpace:D}.");
+                _messageSubscriber.Reject(eventArgs.Message);
+                return;
+            }
+
             try
             {
                 var executionOptions = new ExecutionDataflowBlockOptions
                 {
                     MaxDegreeOfParallelism = Concurrency,
                     MaxMessagesPerTask = 1,
-                    CancellationToken = cancellationToken
+                    CancellationToken = _cancellationTokenSource.Token
                 };
 
-                _exportFlow = new TransformManyBlock<ExportRequestMessage, ExportRequestDataMessage>(
-                    (exportRequest) => DownloadPayloadActionCallback(exportRequest, cancellationToken),
+                var exportFlow = new TransformManyBlock<ExportRequestMessage, ExportRequestDataMessage>(
+                    (exportRequest) => DownloadPayloadActionCallback(exportRequest, _cancellationTokenSource.Token),
                     executionOptions);
 
                 var exportActionBlock = new TransformBlock<ExportRequestDataMessage, ExportRequestDataMessage>(
                     async (exportDataRequest) =>
                     {
                         if (exportDataRequest.IsFailed) return exportDataRequest;
-                        return await ExportDataBlockCallback(exportDataRequest, cancellationToken);
+                        return await ExportDataBlockCallback(exportDataRequest, _cancellationTokenSource.Token);
                     },
                     executionOptions);
 
                 var reportingActionBlock = new ActionBlock<ExportRequestDataMessage>(ReportingActionBlock, executionOptions);
 
                 var linkOptions = new DataflowLinkOptions { PropagateCompletion = true };
 
-                _exportFlow.LinkTo(exportActionBlock, linkOptions);
+                exportFlow.LinkTo(exportActionBlock, linkOptions);
                 exportActionBlock.LinkTo(reportingActionBlock, linkOptions);
 
-                _logger.Log(LogLevel.Information, $"{ServiceName} completed workflow setup.");
+                lock (SyncRoot)
+                {
+                    var exportRequest = eventArgs.Message.ConvertTo<ExportRequestMessage>();
+                    if (_exportRequests.ContainsKey(exportRequest.ExportTaskId))
+                    {
+                        _logger.Log(LogLevel.Warning, $"The export request {exportRequest.ExportTaskId} is already queued for export.");
+                        return;
+                    }
+
+                    exportRequest.MessageId = eventArgs.Message.MessageId;
+                    exportRequest.DeliveryTag = eventArgs.Message.DeliveryTag;
+
+                    _exportRequests.Add(exportRequest.ExportTaskId, exportRequest);
+                    exportFlow.Post(exportRequest);
+                }
+
+                exportFlow.Complete();
+                reportingActionBlock.Completion.Wait(_cancellationTokenSource.Token);
             }
             catch (AggregateException ex)
             {
@@ -170,38 +197,6 @@ private void SetupExportFlow(CancellationToken cancellationToken)
             }
         }
 
-        private void SetupPolling()
-        {
-            _messageSubscriber.Subscribe(RoutingKey, _exportQueueName, OnMessageReceivedCallback);
-            _logger.Log(LogLevel.Information, $"{ServiceName} subscribed to {RoutingKey} messages.");
-        }
-
-        private void OnMessageReceivedCallback(MessageReceivedEventArgs eventArgs)
-        {
-            if (!_storageInfoProvider.HasSpaceAvailableForExport)
-            {
-                _logger.Log(LogLevel.Warning, $"Export service paused due to insufficient storage space.  Available storage space: {_storageInfoProvider.AvailableFreeSpace:D}.");
-                _messageSubscriber.Reject(eventArgs.Message);
-                return;
-            }
-
-            lock (SyncRoot)
-            {
-                var exportRequest = eventArgs.Message.ConvertTo<ExportRequestMessage>();
-                if (_exportRequests.ContainsKey(exportRequest.ExportTaskId))
-                {
-                    _logger.Log(LogLevel.Warning, $"The export request {exportRequest.ExportTaskId} is already queued for export.");
-                    return;
-                }
-
-                exportRequest.MessageId = eventArgs.Message.MessageId;
-                exportRequest.DeliveryTag = eventArgs.Message.DeliveryTag;
-
-                _exportRequests.Add(exportRequest.ExportTaskId, exportRequest);
-                _exportFlow.Post(exportRequest);
-            }
-        }
-
         private IEnumerable<ExportRequestDataMessage> DownloadPayloadActionCallback(ExportRequestMessage exportRequest, CancellationToken cancellationToken)
         {
             Guard.Against.Null(exportRequest, nameof(exportRequest));
@@ -277,12 +272,6 @@ private void ReportingActionBlock(ExportRequestDataMessage exportRequestData)
 
             _logger.Log(LogLevel.Information, $"Export task completed with {exportRequest.FailedFiles} failures out of {exportRequest.Files.Count()}");
 
-            if (ReportActionStarted != null)
-            {
-                _logger.Log(LogLevel.Debug, $"Calling ReportActionStarted callback.");
-                ReportActionStarted(this, null);
-            }
-
             var exportCompleteMessage = new ExportCompleteMessage(exportRequest);
             var jsonMessage = new JsonMessage<ExportCompleteMessage>(exportCompleteMessage, exportRequest.CorrelationId, exportRequest.DeliveryTag);
 
@@ -318,6 +307,12 @@ private void ReportingActionBlock(ExportRequestDataMessage exportRequestData)
             {
                 _exportRequests.Remove(exportRequestData.ExportTaskId);
             }
+
+            if (ReportActionCompleted != null)
+            {
+                _logger.Log(LogLevel.Debug, $"Calling ReportActionCompleted callback.");
+                ReportActionCompleted(this, null);
+            }
         }
     }
 }
diff --git a/src/InformaticsGateway/Test/Services/Export/DicomWebExportServiceTest.cs b/src/InformaticsGateway/Test/Services/Export/DicomWebExportServiceTest.cs
@@ -143,7 +143,7 @@ public async Task ExportDataBlockCallback_ReturnsNullIfInferenceRequestCannotBeF
                 _dicomToolkit.Object);
 
             var dataflowCompleted = new ManualResetEvent(false);
-            service.ReportActionStarted += (sender, args) =>
+            service.ReportActionCompleted += (sender, args) =>
             {
                 dataflowCompleted.Set();
             };
@@ -203,7 +203,7 @@ public async Task ExportDataBlockCallback_ReturnsNullIfInferenceRequestContainsN
                 _dicomToolkit.Object);
 
             var dataflowCompleted = new ManualResetEvent(false);
-            service.ReportActionStarted += (sender, args) =>
+            service.ReportActionCompleted += (sender, args) =>
             {
                 dataflowCompleted.Set();
             };
@@ -288,7 +288,7 @@ public async Task ExportDataBlockCallback_RecordsStowFailuresAndReportFailure()
                 _dicomToolkit.Object);
 
             var dataflowCompleted = new ManualResetEvent(false);
-            service.ReportActionStarted += (sender, args) =>
+            service.ReportActionCompleted += (sender, args) =>
             {
                 dataflowCompleted.Set();
             };
@@ -382,7 +382,7 @@ public async Task CompletesDataflow(HttpStatusCode httpStatusCode)
                 _dicomToolkit.Object);
 
             var dataflowCompleted = new ManualResetEvent(false);
-            service.ReportActionStarted += (sender, args) =>
+            service.ReportActionCompleted += (sender, args) =>
             {
                 dataflowCompleted.Set();
             };
diff --git a/src/InformaticsGateway/Test/Services/Export/ExportServiceBaseTest.cs b/src/InformaticsGateway/Test/Services/Export/ExportServiceBaseTest.cs
@@ -109,7 +109,6 @@ public async Task DataflowTest_StartStop()
             await StopAndVerify(service);
 
             _logger.VerifyLogging($"{service.ServiceName} subscribed to {service.RoutingKey} messages.", LogLevel.Information, Times.Once());
-            _logger.VerifyLogging($"{service.ServiceName} completed workflow setup.", LogLevel.Information, Times.Once());
             _logger.VerifyLogging($"{service.ServiceName} is stopping.", LogLevel.Information, Times.Once());
         }
 
@@ -167,7 +166,7 @@ public async Task DataflowTest_PayloadDownlaodFailure()
                 });
             var countdownEvent = new CountdownEvent(1);
             var service = new TestExportService(_logger.Object, _configuration, _serviceScopeFactory.Object, _storageInfoProvider.Object);
-            service.ReportActionStarted += (sender, e) =>
+            service.ReportActionCompleted += (sender, e) =>
             {
                 countdownEvent.Signal();
             };
@@ -186,9 +185,10 @@ public async Task DataflowTest_PayloadDownlaodFailure()
                                                               It.IsAny<ushort>()), Times.Once());
         }
 
-        [RetryFact(10, 10, DisplayName = "Data flow test - end to end workflow")]
+        [RetryFact(1, 10, DisplayName = "Data flow test - end to end workflow")]
         public async Task DataflowTest_EndToEnd()
         {
+            var messageCount = 5;
             var testData = "this is a test";
             _storageInfoProvider.Setup(p => p.HasSpaceAvailableForExport).Returns(true);
 
@@ -202,17 +202,20 @@ public async Task DataflowTest_EndToEnd()
                                  It.IsAny<ushort>()))
                 .Callback<string, string, Action<MessageReceivedEventArgs>, ushort>((topic, queue, messageReceivedCallback, prefetchCount) =>
                 {
-                    messageReceivedCallback(CreateMessageReceivedEventArgs());
+                    while (messageCount-- > 0)
+                    {
+                        messageReceivedCallback(CreateMessageReceivedEventArgs());
+                    }
                 });
 
             _storageService.Setup(p => p.GetObject(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<Action<Stream>>(), It.IsAny<CancellationToken>()))
                 .Callback<string, string, Action<Stream>, CancellationToken>((bucketName, objectName, callback, cancellationToken) =>
                 {
                     callback(new MemoryStream(Encoding.UTF8.GetBytes(testData)));
                 });
-            var countdownEvent = new CountdownEvent(3);
+            var countdownEvent = new CountdownEvent(5  * 3);
             var service = new TestExportService(_logger.Object, _configuration, _serviceScopeFactory.Object, _storageInfoProvider.Object);
-            service.ReportActionStarted += (sender, e) =>
+            service.ReportActionCompleted += (sender, e) =>
             {
                 countdownEvent.Signal();
             };
@@ -223,13 +226,13 @@ public async Task DataflowTest_EndToEnd()
                 countdownEvent.Signal();
             };
             await service.StartAsync(_cancellationTokenSource.Token);
-            Assert.True(countdownEvent.Wait(3000));
+            Assert.True(countdownEvent.Wait(1000000));
             await StopAndVerify(service);
 
             _messagePublisherService.Verify(
                 p => p.Publish(It.IsAny<string>(),
-                               It.Is<Message>(match => (match.ConvertTo<ExportCompleteMessage>()).Status == ExportStatus.Success)), Times.Once());
-            _messageSubscriberService.Verify(p => p.Acknowledge(It.IsAny<MessageBase>()), Times.Once());
+                               It.Is<Message>(match => (match.ConvertTo<ExportCompleteMessage>()).Status == ExportStatus.Success)), Times.Exactly(5));
+            _messageSubscriberService.Verify(p => p.Acknowledge(It.IsAny<MessageBase>()), Times.Exactly(5));
             _messageSubscriberService.Verify(p => p.Reject(It.IsAny<MessageBase>()), Times.Never());
             _messageSubscriberService.Verify(p => p.Subscribe(It.IsAny<string>(),
                                                               It.IsAny<string>(),
diff --git a/src/InformaticsGateway/Test/Services/Export/ScuExportServiceTest.cs b/src/InformaticsGateway/Test/Services/Export/ScuExportServiceTest.cs
@@ -129,7 +129,7 @@ public async Task ShallFailWhenNoDestinationIsDefined()
             var service = new ScuExportService(_logger.Object, _serviceScopeFactory.Object, _configuration, _storageInfoProvider.Object, _dicomToolkit.Object);
 
             var dataflowCompleted = new ManualResetEvent(false);
-            service.ReportActionStarted += (sender, args) =>
+            service.ReportActionCompleted += (sender, args) =>
             {
                 dataflowCompleted.Set();
             };
@@ -179,7 +179,7 @@ public async Task ShallFailWhenDestinationIsNotConfigured()
             var service = new ScuExportService(_logger.Object, _serviceScopeFactory.Object, _configuration, _storageInfoProvider.Object, _dicomToolkit.Object);
 
             var dataflowCompleted = new ManualResetEvent(false);
-            service.ReportActionStarted += (sender, args) =>
+            service.ReportActionCompleted += (sender, args) =>
             {
                 dataflowCompleted.Set();
             };
@@ -233,7 +233,7 @@ public async Task AssociationRejected()
             var service = new ScuExportService(_logger.Object, _serviceScopeFactory.Object, _configuration, _storageInfoProvider.Object, _dicomToolkit.Object);
 
             var dataflowCompleted = new ManualResetEvent(false);
-            service.ReportActionStarted += (sender, args) =>
+            service.ReportActionCompleted += (sender, args) =>
             {
                 dataflowCompleted.Set();
             };
@@ -289,7 +289,7 @@ public async Task SimulateAbort()
             _dicomToolkit.Setup(p => p.Load(It.IsAny<byte[]>())).Returns(InstanceGenerator.GenerateDicomFile(sopInstanceUid: sopInstanceUid));
 
             var dataflowCompleted = new ManualResetEvent(false);
-            service.ReportActionStarted += (sender, args) =>
+            service.ReportActionCompleted += (sender, args) =>
             {
                 dataflowCompleted.Set();
             };
@@ -345,7 +345,7 @@ public async Task CStoreFailure()
             _dicomToolkit.Setup(p => p.Load(It.IsAny<byte[]>())).Returns(InstanceGenerator.GenerateDicomFile(sopInstanceUid: sopInstanceUid));
 
             var dataflowCompleted = new ManualResetEvent(false);
-            service.ReportActionStarted += (sender, args) =>
+            service.ReportActionCompleted += (sender, args) =>
             {
                 dataflowCompleted.Set();
             };
@@ -401,7 +401,7 @@ public async Task ErrorLoadingDicomContent()
             _dicomToolkit.Setup(p => p.Load(It.IsAny<byte[]>())).Throws(new Exception("error"));
 
             var dataflowCompleted = new ManualResetEvent(false);
-            service.ReportActionStarted += (sender, args) =>
+            service.ReportActionCompleted += (sender, args) =>
             {
                 dataflowCompleted.Set();
             };
@@ -457,7 +457,7 @@ public async Task UnreachableServer()
             _dicomToolkit.Setup(p => p.Load(It.IsAny<byte[]>())).Returns(InstanceGenerator.GenerateDicomFile(sopInstanceUid: sopInstanceUid));
 
             var dataflowCompleted = new ManualResetEvent(false);
-            service.ReportActionStarted += (sender, args) =>
+            service.ReportActionCompleted += (sender, args) =>
             {
                 dataflowCompleted.Set();
             };
@@ -512,7 +512,7 @@ public async Task ExportCompletes()
             _dicomToolkit.Setup(p => p.Load(It.IsAny<byte[]>())).Returns(InstanceGenerator.GenerateDicomFile(sopInstanceUid: sopInstanceUid));
 
             var dataflowCompleted = new ManualResetEvent(false);
-            service.ReportActionStarted += (sender, args) =>
+            service.ReportActionCompleted += (sender, args) =>
             {
                 dataflowCompleted.Set();
             };
diff --git a/src/MessageBroker/RabbitMq/RabbitMqMessageSubscriberService.cs b/src/MessageBroker/RabbitMq/RabbitMqMessageSubscriberService.cs

Original file line number	Diff line number	Diff line change
`@@ -35,12 +35,5 @@ public class MessageBrokerConfigurationKeys`
`35`	`35`	`/// </summary>`
`36`	`36`	`[JsonProperty(PropertyName = "exportRequestPrefix")]`
`37`	`37`	`public string ExportRequestPrefix { get; set; } = "md.export.request";`
`38`		`-`
`39`		`- /// <summary>`
`40`		`- /// Gets or sets the name of the request queue.`
`41`		- /// Defaults to `export_tasks`
`42`		`- /// </summary>`
`43`		`- [JsonProperty(PropertyName = "exportRequestQueue")]`
`44`		`- public string ExportRequestQueue { get; set; } = "export_tasks";`
`45`	`38`	`}`
`46`	`39`	`}`