DSMS-90: Flush DSM stats on Lambda Stop (#7054)

## Summary of changes
We intermittently drop DSM checkpoints when running via Lambda. We don't
flush DSM stats on Lambda end, which can cause us to lose that data if
the invocation doesn't occur during a flush interval.
It's the same logic as for trace submission.

## Reason for change

## Implementation details
Adds a FlushAsync call to the DataStreamsManager.

## Test coverage

Mostly manual testing, I was invoking a lambda at a regular cadence and
expecting to see a steady number of messages. With the changes in the
PR, the rate is constant. When the changes were disabled, then we
started to see gaps in DSM metrics.
There's probably a more thorough way of testing this, but I think this
is good enough.

**Bug:** Number of DSM tracked messages (Missing data points in the
middle invocation, which uses the old logic):
<img width="1511" height="366" alt="Screenshot 2025-07-28 at 3 12 45 pm"
src="https://github.com/user-attachments/assets/579af9dc-021e-4cb8-982a-4d80a5b40fcf"
/>

Lambda Invocations (Constant for both):
<img width="2382" height="616" alt="Screenshot 2025-07-28 at 3 12 57 pm"
src="https://github.com/user-attachments/assets/e29caec9-053b-450b-b4a4-564a6dba75be"
/>

## Other details
<!-- Fixes #{issue} -->

<!-- ⚠️ Note: where possible, please obtain 2 approvals prior to
merging. Unless CODEOWNERS specifies otherwise, for external teams it is
typically best to have one review from a team member, and one review
from apm-dotnet. Trivial changes do not require 2 reviews. -->

Author: robcarlan-datadog

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/LambdaCommon.cs

@@ -71,9 +71,12 @@ internal static async Task EndInvocationAsync(string returnValue, Exception exce
     {
         try
         {
-            await Tracer.Instance.TracerManager.AgentWriter.FlushTracesAsync()
-                        .WaitAsync(TimeSpan.FromSeconds(ServerlessMaxWaitingFlushTime))
-                        .ConfigureAwait(false);
+            await Task.WhenAll(
+                Tracer.Instance.TracerManager.AgentWriter.FlushTracesAsync()
+                    .WaitAsync(TimeSpan.FromSeconds(ServerlessMaxWaitingFlushTime)),
+                Tracer.Instance.TracerManager.DataStreamsManager.FlushAsync()
+                    .WaitAsync(TimeSpan.FromSeconds(ServerlessMaxWaitingFlushTime)))
+                .ConfigureAwait(false);
         }
         catch (Exception ex)
         {

tracer/src/Datadog.Trace/DataStreamsMonitoring/DataStreamsManager.cs

@@ -75,6 +75,22 @@ public async Task DisposeAsync()
         await writer.DisposeAsync().ConfigureAwait(false);
     }
 
+    public async Task FlushAsync()
+    {
+        if (!IsEnabled)
+        {
+            return;
+        }
+
+        var writer = Volatile.Read(ref _writer);
+        if (writer is null)
+        {
+            return;
+        }
+
+        await writer.FlushAsync().ConfigureAwait(false);
+    }
+
     /// <summary>
     /// Trys to extract a <see cref="PathwayContext"/>, from the provided <paramref name="headers"/>
     /// If data streams is disabled, or no pathway is present, returns null.

tracer/src/Datadog.Trace/DataStreamsMonitoring/DataStreamsWriter.cs

@@ -34,11 +34,13 @@ internal class DataStreamsWriter : IDataStreamsWriter
     private readonly bool _isInDefaultState;
 
     private readonly TaskCompletionSource<bool> _processExit = new(TaskCreationOptions.RunContinuationsAsynchronously);
+    private readonly SemaphoreSlim _flushSemaphore = new(1, 1);
     private MemoryStream? _serializationBuffer;
     private long _pointsDropped;
     private int _flushRequested;
     private Task? _processTask;
     private Timer? _flushTimer;
+    private TaskCompletionSource<bool>? _currentFlushTcs;
 
     private int _isSupported = SupportState.Unknown;
     private bool _isInitialized;
@@ -152,6 +154,7 @@ public async Task DisposeAsync()
         _flushTimer?.Dispose();
 #endif
         await FlushAndCloseAsync().ConfigureAwait(false);
+        _flushSemaphore.Dispose();
     }
 
     private async Task FlushAndCloseAsync()
@@ -184,6 +187,45 @@ private async Task FlushAndCloseAsync()
         }
     }
 
+    public async Task FlushAsync()
+    {
+        await _flushSemaphore.WaitAsync().ConfigureAwait(false);
+        try
+        {
+            var timeout = TimeSpan.FromSeconds(5);
+
+            if (_processExit.Task.IsCompleted)
+            {
+                return;
+            }
+
+            if (!Volatile.Read(ref _isInitialized) || _processTask == null)
+            {
+                return;
+            }
+
+            var tcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+            Interlocked.Exchange(ref _currentFlushTcs, tcs);
+
+            RequestFlush();
+
+            var completedTask = await Task.WhenAny(
+                tcs.Task,
+                _processExit.Task,
+                Task.Delay(timeout)).ConfigureAwait(false);
+
+            if (completedTask != tcs.Task)
+            {
+                Log.Error("Data streams flush timeout after {Timeout}ms", timeout.TotalMilliseconds);
+            }
+        }
+        finally
+        {
+            _currentFlushTcs = null;
+            _flushSemaphore.Release();
+        }
+    }
+
     private void RequestFlush()
     {
         Interlocked.Exchange(ref _flushRequested, 1);
@@ -248,6 +290,8 @@ private async Task ProcessQueueLoopAsync()
                 if (flushRequested == 1)
                 {
                     await WriteToApiAsync().ConfigureAwait(false);
+                    var currentFlushTcs = Volatile.Read(ref _currentFlushTcs);
+                    currentFlushTcs?.TrySetResult(true);
                     FlushComplete?.Invoke(this, EventArgs.Empty);
                 }
             }

tracer/src/Datadog.Trace/DataStreamsMonitoring/IDataStreamsWriter.cs

@@ -15,5 +15,7 @@ internal interface IDataStreamsWriter
 
     void AddBacklog(in BacklogPoint point);
 
+    Task FlushAsync();
+
     Task DisposeAsync();
 }

tracer/test/Datadog.Trace.Tests/DataStreamsMonitoring/DataStreamsManagerTests.cs

@@ -332,5 +332,10 @@ public async Task DisposeAsync()
             Interlocked.Increment(ref _disposeCount);
             await Task.Yield();
         }
+
+        public async Task FlushAsync()
+        {
+            await Task.Yield();
+        }
     }
 }