Fixing DSM flush logic & CPU usage (#7827)

## Summary of changes
Moves DSM flush work to a dedicated thread to improve thread pool
performance. Also simplifies the flushing logic to more closely match
the AgentWriter logic.

(See below for test results)

## Reason for change

A = DSM enabled after the fix
B = DSM enabled before the fix
<img width="1168" height="573" alt="Screenshot 2025-12-01 at 5 23 44 pm"
src="https://github.com/user-attachments/assets/95d77f11-927c-4e7a-bfa6-451a1cf527a9"
/>

DSM flush logic was previously running entirely on the thread pool,
which includes the task to frequently move DSM checkpoints off of the
buffer and add them to the backlog. This was causing thread contention
shown in the above profile.

## Implementation details

Moves the DSM buffer flushing logic into it's own thread. Below is a
profile after the fix (DSM disabled on the right).

<img width="1422" height="622" alt="Screenshot 2025-12-01 at 5 26 53 pm"
src="https://github.com/user-attachments/assets/4a0cb28a-633f-47c8-9e04-cff9ce6829b9"
/>

## Test coverage

Mostly manual testing. I ran a set of three apps:
1. DSM enabled, after the fix
2. DSM disabled, after the fix
3. DSM enabled, before the fix

The test app just forwards a Kafka message every 0.5 seconds.

And run those via docker-compose with a restricted CPU count of 0.1 and
0.5 in the below screenshots. In both cases, DSM disabled is the lowest
CPU usage, and the line just above is DSM enabled after the fix. The
highest CPU is DSM enabled, prior to the fix. The performance is a lot
better, but I'm not sure if the gap can be closed further.

<img width="1046" height="523" alt="Screenshot 2025-12-01 at 6 30 10 pm"
src="https://github.com/user-attachments/assets/09cc8cdb-ba8d-4b44-be6e-b866cccb022d"
/>

<img width="694" height="379" alt="Screenshot 2025-12-01 at 5 55 28 pm"
src="https://github.com/user-attachments/assets/f4e1c24d-4800-46ed-8e86-42eec3def22d"
/>

## Other details

Author: robcarlan-datadog

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

@@ -21,26 +21,29 @@ namespace Datadog.Trace.DataStreamsMonitoring;
 
 internal class DataStreamsWriter : IDataStreamsWriter
 {
+    private const TaskCreationOptions TaskOptions = TaskCreationOptions.RunContinuationsAsynchronously;
+
     private static readonly IDatadogLogger Log = DatadogLogging.GetLoggerFor<DataStreamsWriter>();
 
     private readonly object _initLock = new();
     private readonly long _bucketDurationMs;
     private readonly BoundedConcurrentQueue<StatsPoint> _buffer = new(queueLimit: 10_000);
     private readonly BoundedConcurrentQueue<BacklogPoint> _backlogBuffer = new(queueLimit: 10_000);
     private readonly TimeSpan _waitTimeSpan = TimeSpan.FromMilliseconds(10);
+    private readonly TimeSpan _flushSemaphoreWaitTime = TimeSpan.FromSeconds(1);
     private readonly DataStreamsAggregator _aggregator;
     private readonly IDiscoveryService _discoveryService;
     private readonly IDataStreamsApi _api;
     private readonly bool _isInDefaultState;
 
-    private readonly TaskCompletionSource<bool> _processExit = new(TaskCreationOptions.RunContinuationsAsynchronously);
+    private readonly TaskCompletionSource<bool> _processExit = new(TaskOptions);
     private readonly SemaphoreSlim _flushSemaphore = new(1, 1);
     private MemoryStream? _serializationBuffer;
     private long _pointsDropped;
-    private int _flushRequested;
     private Task? _processTask;
+    private Task? _flushTask;
     private Timer? _flushTimer;
-    private TaskCompletionSource<bool>? _currentFlushTcs;
+    private TaskCompletionSource<bool> _forceFlush = new(TaskOptions);
 
     private int _isSupported = SupportState.Unknown;
     private bool _isInitialized;
@@ -93,8 +96,12 @@ private void Initialize()
                 return;
             }
 
-            _processTask = Task.Run(ProcessQueueLoopAsync);
+            _processTask = Task.Factory.StartNew(ProcessQueueLoop, TaskCreationOptions.LongRunning);
             _processTask.ContinueWith(t => Log.Error(t.Exception, "Error in processing task"), TaskContinuationOptions.OnlyOnFaulted);
+
+            _flushTask = Task.Run(FlushTaskLoopAsync);
+            _flushTask.ContinueWith(t => Log.Error(t.Exception, "Error in data streams flush task"), TaskContinuationOptions.OnlyOnFaulted);
+
             _flushTimer = new Timer(
                 x => ((DataStreamsWriter)x!).RequestFlush(),
                 this,
@@ -169,67 +176,106 @@ private async Task FlushAndCloseAsync()
             return;
         }
 
-        // request a final flush - as the _processExit flag is now set
-        // this ensures we will definitely flush all the stats
-        // (and sets the mutex if it isn't already set)
-        RequestFlush();
-
-        // wait for the processing loop to complete
         var completedTask = await Task.WhenAny(
                                            _processTask,
-                                           Task.Delay(TimeSpan.FromSeconds(20)))
+                                           Task.Delay(TimeSpan.FromSeconds(1)))
                                       .ConfigureAwait(false);
 
         if (completedTask != _processTask)
         {
             Log.Error("Could not flush all data streams stats before process exit");
         }
+
+        await FlushAsync().ConfigureAwait(false);
     }
 
-    public async Task FlushAsync()
+    private void RequestFlush()
     {
-        await _flushSemaphore.WaitAsync().ConfigureAwait(false);
-        try
+        _forceFlush.TrySetResult(true);
+    }
+
+    private async Task FlushTaskLoopAsync()
+    {
+        Task[] tasks = new Task[2];
+        tasks[0] = _processTask!;
+        tasks[1] = _forceFlush.Task;
+
+        while (true)
         {
-            var timeout = TimeSpan.FromSeconds(5);
+            await Task.WhenAny(tasks).ConfigureAwait(false);
 
-            if (_processExit.Task.IsCompleted)
+            if (_forceFlush.Task.IsCompleted)
             {
-                return;
+                _forceFlush = new TaskCompletionSource<bool>(TaskOptions);
+                tasks[1] = _forceFlush.Task;
             }
 
-            if (!Volatile.Read(ref _isInitialized) || _processTask == null)
+            await FlushAggregatorAsync().ConfigureAwait(false);
+
+            if (_processTask!.IsCompleted)
             {
                 return;
             }
+        }
+    }
 
-            var tcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
-            Interlocked.Exchange(ref _currentFlushTcs, tcs);
-
-            RequestFlush();
+    public async Task FlushAsync()
+    {
+        if (!Volatile.Read(ref _isInitialized) || _processTask == null)
+        {
+            return;
+        }
 
-            var completedTask = await Task.WhenAny(
-                tcs.Task,
-                _processExit.Task,
-                Task.Delay(timeout)).ConfigureAwait(false);
+        if (await _flushSemaphore.WaitAsync(TimeSpan.FromMilliseconds(100)).ConfigureAwait(false))
+        {
+            try
+            {
+                while (_buffer.TryDequeue(out var statsPoint))
+                {
+                    _aggregator.Add(in statsPoint);
+                }
 
-            if (completedTask != tcs.Task)
+                while (_backlogBuffer.TryDequeue(out var backlogPoint))
+                {
+                    _aggregator.AddBacklog(in backlogPoint);
+                }
+            }
+            catch (Exception ex)
+            {
+                Log.Error(ex, "An error occured while processing data streams buffers");
+            }
+            finally
             {
-                Log.Error("Data streams flush timeout after {Timeout}ms", timeout.TotalMilliseconds);
+                _flushSemaphore.Release();
             }
         }
+
+        await FlushAggregatorAsync().ConfigureAwait(false);
+    }
+
+    private async Task FlushAggregatorAsync()
+    {
+        if (!await _flushSemaphore.WaitAsync(_flushSemaphoreWaitTime).ConfigureAwait(false))
+        {
+            Log.Warning("Data streams flush timeout");
+            return;
+        }
+
+        try
+        {
+            await WriteToApiAsync().ConfigureAwait(false);
+            FlushComplete?.Invoke(this, EventArgs.Empty);
+        }
+        catch (Exception ex)
+        {
+            Log.Error(ex, "An error occured while writing data streams");
+        }
         finally
         {
-            _currentFlushTcs = null;
             _flushSemaphore.Release();
         }
     }
 
-    private void RequestFlush()
-    {
-        Interlocked.Exchange(ref _flushRequested, 1);
-    }
-
     private async Task WriteToApiAsync()
     {
         // This method blocks ingestion of new stats points into the aggregator,
@@ -268,11 +314,22 @@ private async Task WriteToApiAsync()
         }
     }
 
-    private async Task ProcessQueueLoopAsync()
+    private void ProcessQueueLoop()
     {
-        var isFinalFlush = false;
         while (true)
         {
+            Thread.Sleep(_waitTimeSpan);
+
+            if (!_flushSemaphore.Wait(_flushSemaphoreWaitTime))
+            {
+                if (_processExit.Task.IsCompleted)
+                {
+                    return;
+                }
+
+                continue;
+            }
+
             try
             {
                 while (_buffer.TryDequeue(out var statsPoint))
@@ -284,40 +341,19 @@ private async Task ProcessQueueLoopAsync()
                 {
                     _aggregator.AddBacklog(in backlogPoint);
                 }
-
-                var flushRequested = Interlocked.CompareExchange(ref _flushRequested, 0, 1);
-                if (flushRequested == 1)
-                {
-                    await WriteToApiAsync().ConfigureAwait(false);
-                    var currentFlushTcs = Volatile.Read(ref _currentFlushTcs);
-                    currentFlushTcs?.TrySetResult(true);
-                    FlushComplete?.Invoke(this, EventArgs.Empty);
-                }
             }
             catch (Exception ex)
             {
-                Log.Error(ex, "An error occured in the processing thread");
+                Log.Error(ex, "An error occured in the data streams processing thread");
             }
-
-            if (_processExit.Task.IsCompleted)
+            finally
             {
-                if (isFinalFlush)
-                {
-                    return;
-                }
-
-                // do one more loop to make sure everything is flushed
-                RequestFlush();
-                isFinalFlush = true;
-                continue;
+                _flushSemaphore.Release();
             }
 
-            // The logic is copied from https://github.com/dotnet/runtime/blob/main/src/libraries/Common/tests/System/Threading/Tasks/TaskTimeoutExtensions.cs#L26
-            // and modified to avoid dealing with exceptions
-            var tcs = new TaskCompletionSource<bool>();
-            using (new Timer(s => ((TaskCompletionSource<bool>)s!).SetResult(true), tcs, _waitTimeSpan, Timeout.InfiniteTimeSpan))
+            if (_processExit.Task.IsCompleted)
             {
-                await Task.WhenAny(_processExit.Task, tcs.Task).ConfigureAwait(false);
+                return;
             }
         }
     }

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

@@ -218,7 +218,7 @@ public void WhenEnabled_SetCheckpoint_SetsSpanTags()
         var dsm = GetDataStreamManager(true, out _);
         var span = new Span(new SpanContext(traceId: 123, spanId: 456), DateTimeOffset.UtcNow);
 
-        span.SetDataStreamsCheckpoint(dsm,  CheckpointKind.Produce, new[] { "direction:out" }, 100, 0);
+        span.SetDataStreamsCheckpoint(dsm, CheckpointKind.Produce, new[] { "direction:out" }, 100, 0);
         span.Tags.GetTag("pathway.hash").Should().NotBeNull();
     }