track total egress bytes on stow (#3337)

track total egress bytes on stow
microsoft · Feb 3, 2024 · 0743fda · 0743fda
1 parent 4f2bd5d
commit 0743fda
Show file tree

Hide file tree

Showing 5 changed files with 84 additions and 1 deletion.
diff --git a/src/Microsoft.Health.Dicom.Api.UnitTests/Features/Context/DefaultDicomRequestContext.cs b/src/Microsoft.Health.Dicom.Api.UnitTests/Features/Context/DefaultDicomRequestContext.cs
@@ -51,4 +51,9 @@ public class DefaultDicomRequestContext : IDicomRequestContext
     public IDictionary<string, StringValues> ResponseHeaders { get; set; }
 
     public int Version { get; set; }
+
+    /// <summary>
+    /// Egress bytes from Dicom server to other resources
+    /// </summary>
+    public long TotalDicomEgressToStorageBytes { get; set; }
 }
diff --git a/src/Microsoft.Health.Dicom.Core.UnitTests/Features/Store/DicomStoreServiceTests.cs b/src/Microsoft.Health.Dicom.Core.UnitTests/Features/Store/DicomStoreServiceTests.cs
@@ -155,12 +155,15 @@ public async Task GivenAValidDicomInstanceEntry_WhenProcessed_ThenSuccessfulEntr
         IDicomInstanceEntry dicomInstanceEntry = Substitute.For<IDicomInstanceEntry>();
 
         dicomInstanceEntry.GetDicomDatasetAsync(DefaultCancellationToken).Returns(_dicomDataset1);
+        long bytesStored = 100L;
+        _storeOrchestrator.StoreDicomInstanceEntryAsync(dicomInstanceEntry, DefaultCancellationToken).Returns(bytesStored);
 
         await ExecuteAndValidateAsync(dicomInstanceEntry);
 
         _storeResponseBuilder.Received(1).AddSuccess(_dicomDataset1, DefaultStoreValidationResult, _dicomRequestContext.DataPartition);
         _storeResponseBuilder.DidNotReceiveWithAnyArgs().AddFailure(default);
         Assert.Equal(1, _dicomRequestContextAccessor.RequestContext.PartCount);
+        Assert.Equal(bytesStored, _dicomRequestContextAccessor.RequestContext.TotalDicomEgressToStorageBytes);
     }
 
     [Fact]
@@ -174,6 +177,8 @@ public async Task GiveAnInvalidDicomDataset_WhenProcessed_ThenFailedEntryShouldB
 
         _storeResponseBuilder.DidNotReceiveWithAnyArgs().AddSuccess(default, DefaultStoreValidationResult, Partition.Default);
         _storeResponseBuilder.Received(1).AddFailure(null, TestConstants.ProcessingFailureReasonCode);
+        Assert.Equal(0, _dicomRequestContextAccessor.RequestContext.TotalDicomEgressToStorageBytes);
+        await _storeOrchestrator.DidNotReceiveWithAnyArgs().StoreDicomInstanceEntryAsync(Arg.Any<IDicomInstanceEntry>(), Arg.Any<CancellationToken>());
     }
 
     [Fact]
@@ -187,6 +192,8 @@ public async Task GivenADicomDatasetFailsToOpenDueToDicomValidationException_Whe
 
         _storeResponseBuilder.DidNotReceiveWithAnyArgs().AddSuccess(default, DefaultStoreValidationResult, Partition.Default);
         _storeResponseBuilder.Received(1).AddFailure(null, TestConstants.ValidationFailureReasonCode);
+        Assert.Equal(0, _dicomRequestContextAccessor.RequestContext.TotalDicomEgressToStorageBytes);
+        await _storeOrchestrator.DidNotReceiveWithAnyArgs().StoreDicomInstanceEntryAsync(Arg.Any<IDicomInstanceEntry>(), Arg.Any<CancellationToken>());
     }
 
     [Fact]
@@ -199,6 +206,8 @@ public async Task GivenAValidationError_WhenDropDataEnabled_ThenSucceedsWithErro
         dicomDataset.Add(DicomTag.StudyDate, "NotAValidStudyDate");
 
         dicomInstanceEntry.GetDicomDatasetAsync(DefaultCancellationToken).Returns(dicomDataset);
+        long bytesStored = 100L;
+        _storeOrchestrator.StoreDicomInstanceEntryAsync(dicomInstanceEntry, DefaultCancellationToken).Returns(bytesStored);
 
         // call
         StoreResponse response = await _storeServiceDropData.ProcessAsync(
@@ -235,6 +244,8 @@ await _storeOrchestrator
                 dicomInstanceEntry,
                 DefaultCancellationToken
             );
+
+        Assert.Equal(bytesStored, _dicomRequestContextAccessorLatestApi.RequestContext.TotalDicomEgressToStorageBytes);
     }
 
     [Fact]
@@ -247,6 +258,8 @@ public async Task GivenAValidationErrorOnNonCoreTag_WhenDropDataEnabledAndFullDi
         dicomDataset.Add(DicomTag.StudyDate, "NotAValidStudyDate");
 
         dicomInstanceEntry.GetDicomDatasetAsync(DefaultCancellationToken).Returns(dicomDataset);
+        long bytesStored = 100L;
+        _storeOrchestrator.StoreDicomInstanceEntryAsync(dicomInstanceEntry, DefaultCancellationToken).Returns(bytesStored);
 
         // call
         StoreResponse response = await _storeServiceDropData.ProcessAsync(
@@ -283,6 +296,8 @@ await _storeOrchestrator
                 dicomInstanceEntry,
                 DefaultCancellationToken
             );
+
+        Assert.Equal(bytesStored, _dicomRequestContextAccessorLatestApi.RequestContext.TotalDicomEgressToStorageBytes);
     }
 
     [Fact]
@@ -515,10 +530,12 @@ public async Task GivenAnExternalStoreExceptionWhenStoring_WhenProcessed_ThenExp
     }
 
     [Fact]
-    public async Task GivenMultipleDicomInstanceEntries_WhenProcessed_ThenCorrespondingEntryShouldBeAdded()
+    public async Task GivenMultipleDicomInstanceEntriesWithFailure_WhenProcessed_ThenCorrespondingEntryShouldBeAdded()
     {
         IDicomInstanceEntry dicomInstanceEntryToSucceed = Substitute.For<IDicomInstanceEntry>();
         IDicomInstanceEntry dicomInstanceEntryToFail = Substitute.For<IDicomInstanceEntry>();
+        long bytesStored = 100L;
+        _storeOrchestrator.StoreDicomInstanceEntryAsync(Arg.Any<IDicomInstanceEntry>(), DefaultCancellationToken).Returns(bytesStored);
 
         dicomInstanceEntryToSucceed.GetDicomDatasetAsync(DefaultCancellationToken).Returns(_dicomDataset1);
         dicomInstanceEntryToFail.GetDicomDatasetAsync(DefaultCancellationToken).Returns(_dicomDataset2);
@@ -535,6 +552,56 @@ public async Task GivenMultipleDicomInstanceEntries_WhenProcessed_ThenCorrespond
         _storeResponseBuilder.Received(1).AddSuccess(_dicomDataset1, DefaultStoreValidationResult, _dicomRequestContext.DataPartition);
         _storeResponseBuilder.Received(1).AddFailure(_dicomDataset2, TestConstants.ProcessingFailureReasonCode);
         Assert.Equal(2, _dicomRequestContextAccessor.RequestContext.PartCount);
+
+        await _storeOrchestrator
+            .Received(1)
+            .StoreDicomInstanceEntryAsync(
+                dicomInstanceEntryToSucceed,
+                DefaultCancellationToken
+            );
+
+        await _storeOrchestrator
+            .DidNotReceive()
+            .StoreDicomInstanceEntryAsync(
+                dicomInstanceEntryToFail,
+                DefaultCancellationToken
+            );
+
+        Assert.Equal(bytesStored, _dicomRequestContextAccessor.RequestContext.TotalDicomEgressToStorageBytes);
+    }
+
+    [Fact]
+    public async Task GivenMultipleDicomInstanceEntries_WhenProcessed_ThenTotalBytesTracked()
+    {
+        IDicomInstanceEntry dicomInstanceEntryToSucceed = Substitute.For<IDicomInstanceEntry>();
+        IDicomInstanceEntry dicomInstanceEntryToSucceed2 = Substitute.For<IDicomInstanceEntry>();
+        long bytesStored = 100L;
+        long totalExpectedBytesStored = bytesStored * 2;
+        _storeOrchestrator.StoreDicomInstanceEntryAsync(Arg.Any<IDicomInstanceEntry>(), DefaultCancellationToken).Returns(bytesStored);
+
+        dicomInstanceEntryToSucceed.GetDicomDatasetAsync(DefaultCancellationToken).Returns(_dicomDataset1);
+        dicomInstanceEntryToSucceed2.GetDicomDatasetAsync(DefaultCancellationToken).Returns(_dicomDataset2);
+
+        await ExecuteAndValidateAsync(dicomInstanceEntryToSucceed, dicomInstanceEntryToSucceed2);
+
+        _storeResponseBuilder.Received(1).AddSuccess(_dicomDataset1, DefaultStoreValidationResult, _dicomRequestContext.DataPartition);
+        _storeResponseBuilder.Received(1).AddSuccess(_dicomDataset2, DefaultStoreValidationResult, _dicomRequestContext.DataPartition);
+
+        await _storeOrchestrator
+            .Received(1)
+            .StoreDicomInstanceEntryAsync(
+                dicomInstanceEntryToSucceed,
+                DefaultCancellationToken
+            );
+
+        await _storeOrchestrator
+            .Received(1)
+            .StoreDicomInstanceEntryAsync(
+                dicomInstanceEntryToSucceed2,
+                DefaultCancellationToken
+            );
+
+        Assert.Equal(totalExpectedBytesStored, _dicomRequestContextAccessor.RequestContext.TotalDicomEgressToStorageBytes);
     }
 
     [Fact]

diff --git a/src/Microsoft.Health.Dicom.Core/Features/Context/DicomRequestContext.cs b/src/Microsoft.Health.Dicom.Core/Features/Context/DicomRequestContext.cs
@@ -86,4 +86,9 @@ public DicomRequestContext(
     public IDictionary<string, StringValues> RequestHeaders { get; }
 
     public IDictionary<string, StringValues> ResponseHeaders { get; }
+
+    /// <summary>
+    /// Egress bytes from Dicom server to other resources
+    /// </summary>
+    public long TotalDicomEgressToStorageBytes { get; set; }
 }
diff --git a/src/Microsoft.Health.Dicom.Core/Features/Context/IDicomRequestContext.cs b/src/Microsoft.Health.Dicom.Core/Features/Context/IDicomRequestContext.cs
@@ -28,4 +28,9 @@ public interface IDicomRequestContext : IRequestContext
 
     // Opportunity for the core to change based on the caller version
     int Version { get; set; }
+
+    /// <summary>
+    /// Egress bytes from Dicom server to other resources
+    /// </summary>
+    long TotalDicomEgressToStorageBytes { get; set; }
 }
diff --git a/src/Microsoft.Health.Dicom.Core/Features/Store/StoreService.cs b/src/Microsoft.Health.Dicom.Core/Features/Store/StoreService.cs
@@ -115,6 +115,7 @@ public async Task<StoreResponse> ProcessAsync(
                         long len = length.GetValueOrDefault();
                         // Update Telemetry
                         _storeMeter.InstanceLength.Record(len);
+                        _dicomRequestContextAccessor.RequestContext.TotalDicomEgressToStorageBytes += len;
                     }
                 }
                 finally