Rework the StatsDManager

Make sure we can't dispose a stats consumer that's in use (as it will throw)
Rework to use a "lease" mechanism to track usages
Make passing in a statsmanager required

Author: andrewlock

tracer/missing-nullability-files.csv

@@ -34,10 +34,8 @@ src/Datadog.Trace/Tracer.cs
 src/Datadog.Trace/TracerConstants.cs
 src/Datadog.Trace/TracerManager.cs
 src/Datadog.Trace/TracerManagerFactory.cs
-src/Datadog.Trace/Agent/AgentWriter.cs
 src/Datadog.Trace/Agent/Api.cs
 src/Datadog.Trace/Agent/ClientStatsPayload.cs
-src/Datadog.Trace/Agent/IAgentWriter.cs
 src/Datadog.Trace/Agent/IApi.cs
 src/Datadog.Trace/Agent/IApiRequest.cs
 src/Datadog.Trace/Agent/IApiRequestFactory.cs

tracer/src/Datadog.Trace/Agent/AgentWriter.cs

@@ -3,6 +3,8 @@
 // This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2017 Datadog, Inc.
 // </copyright>
 
+#nullable enable
+
 using System;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
@@ -12,10 +14,8 @@
 using Datadog.Trace.Configuration;
 using Datadog.Trace.DogStatsd;
 using Datadog.Trace.Logging;
-using Datadog.Trace.Tagging;
 using Datadog.Trace.Telemetry;
 using Datadog.Trace.Telemetry.Metrics;
-using Datadog.Trace.Vendors.StatsdClient;
 
 namespace Datadog.Trace.Agent
 {
@@ -25,10 +25,10 @@ internal class AgentWriter : IAgentWriter
 
         private static readonly IDatadogLogger Log = DatadogLogging.GetLoggerFor<AgentWriter>();
 
-        private static readonly ArraySegment<byte> EmptyPayload = new(new byte[] { 0x90 });
+        private static readonly ArraySegment<byte> EmptyPayload = new([0x90]);
 
         private readonly ConcurrentQueue<WorkItem> _pendingTraces = new ConcurrentQueue<WorkItem>();
-        private readonly IDogStatsd _statsd;
+        private readonly IStatsdManager _statsd;
         private readonly Task _flushTask;
         private readonly Task _serializationTask;
         private readonly TaskCompletionSource<bool> _processExit = new TaskCompletionSource<bool>();
@@ -43,7 +43,7 @@ internal class AgentWriter : IAgentWriter
         private readonly int _batchInterval;
         private readonly IKeepRateCalculator _traceKeepRateCalculator;
 
-        private readonly IStatsAggregator _statsAggregator;
+        private readonly IStatsAggregator? _statsAggregator;
 
         private readonly bool _apmTracingEnabled;
 
@@ -67,7 +67,7 @@ internal class AgentWriter : IAgentWriter
 
         private bool _traceMetricsEnabled;
 
-        public AgentWriter(IApi api, IStatsAggregator statsAggregator, IDogStatsd statsd, TracerSettings settings)
+        public AgentWriter(IApi api, IStatsAggregator? statsAggregator, IStatsdManager statsd, TracerSettings settings)
             : this(api, statsAggregator, statsd, maxBufferSize: settings.TraceBufferSize, batchInterval: settings.TraceBatchInterval, apmTracingEnabled: settings.ApmTracingEnabled, initialTracerMetricsEnabled: settings.Manager.InitialMutableSettings.TracerMetricsEnabled)
         {
             settings.Manager.SubscribeToChanges(changes =>
@@ -76,16 +76,17 @@ public AgentWriter(IApi api, IStatsAggregator statsAggregator, IDogStatsd statsd
                  && mutable.TracerMetricsEnabled != changes.PreviousMutable.TracerMetricsEnabled)
                 {
                     Volatile.Write(ref _traceMetricsEnabled, mutable.TracerMetricsEnabled);
+                    _statsd.SetRequired(StatsdConsumer.AgentWriter, mutable.TracerMetricsEnabled);
                 }
             });
         }
 
-        public AgentWriter(IApi api, IStatsAggregator statsAggregator, IDogStatsd statsd, bool automaticFlush = true, int maxBufferSize = 1024 * 1024 * 10, int batchInterval = 100, bool apmTracingEnabled = true, bool initialTracerMetricsEnabled = false)
+        public AgentWriter(IApi api, IStatsAggregator? statsAggregator, IStatsdManager statsd, bool automaticFlush = true, int maxBufferSize = 1024 * 1024 * 10, int batchInterval = 100, bool apmTracingEnabled = true, bool initialTracerMetricsEnabled = false)
         : this(api, statsAggregator, statsd, MovingAverageKeepRateCalculator.CreateDefaultKeepRateCalculator(), automaticFlush, maxBufferSize, batchInterval, apmTracingEnabled, initialTracerMetricsEnabled)
         {
         }
 
-        internal AgentWriter(IApi api, IStatsAggregator statsAggregator, IDogStatsd statsd, IKeepRateCalculator traceKeepRateCalculator, bool automaticFlush, int maxBufferSize, int batchInterval, bool apmTracingEnabled, bool initialTracerMetricsEnabled)
+        internal AgentWriter(IApi api, IStatsAggregator? statsAggregator, IStatsdManager statsd, IKeepRateCalculator traceKeepRateCalculator, bool automaticFlush, int maxBufferSize, int batchInterval, bool apmTracingEnabled, bool initialTracerMetricsEnabled)
         {
             _statsAggregator = statsAggregator;
 
@@ -104,6 +105,7 @@ internal AgentWriter(IApi api, IStatsAggregator statsAggregator, IDogStatsd stat
 
             _apmTracingEnabled = apmTracingEnabled;
             _traceMetricsEnabled = initialTracerMetricsEnabled;
+            _statsd.SetRequired(StatsdConsumer.AgentWriter, initialTracerMetricsEnabled);
 
             _serializationTask = automaticFlush ? Task.Factory.StartNew(SerializeTracesLoop, TaskCreationOptions.LongRunning) : Task.CompletedTask;
             _serializationTask.ContinueWith(t => Log.Error(t.Exception, "Error in serialization task"), TaskContinuationOptions.OnlyOnFaulted);
@@ -114,7 +116,7 @@ internal AgentWriter(IApi api, IStatsAggregator statsAggregator, IDogStatsd stat
             _backBufferFlushTask = _frontBufferFlushTask = Task.CompletedTask;
         }
 
-        internal event Action Flushed;
+        internal event Action? Flushed;
 
         internal SpanBuffer ActiveBuffer => _activeBuffer;
 
@@ -151,8 +153,12 @@ public void WriteTrace(ArraySegment<Span> trace)
 
             if (Volatile.Read(ref _traceMetricsEnabled))
             {
-                _statsd.Increment(TracerMetricNames.Queue.EnqueuedTraces);
-                _statsd.Increment(TracerMetricNames.Queue.EnqueuedSpans, trace.Count);
+                using var lease = _statsd.TryGetClientLease();
+                if (lease.Client is { } statsd)
+                {
+                    statsd.Increment(TracerMetricNames.Queue.EnqueuedTraces);
+                    statsd.Increment(TracerMetricNames.Queue.EnqueuedSpans, trace.Count);
+                }
             }
         }
 
@@ -255,7 +261,7 @@ private async Task FlushBuffersTaskLoopAsync()
             {
                 tasks[2] = Task.Delay(TimeSpan.FromSeconds(1));
                 await Task.WhenAny(tasks).ConfigureAwait(false);
-                tasks[2] = null;
+                tasks[2] = null!;
 
                 if (_forceFlush.Task.IsCompleted)
                 {
@@ -327,8 +333,12 @@ async Task InternalBufferFlush()
                 {
                     if (Volatile.Read(ref _traceMetricsEnabled))
                     {
-                        _statsd.Increment(TracerMetricNames.Queue.DequeuedTraces, buffer.TraceCount);
-                        _statsd.Increment(TracerMetricNames.Queue.DequeuedSpans, buffer.SpanCount);
+                        using var lease = _statsd.TryGetClientLease();
+                        if (lease.Client is { } statsd)
+                        {
+                            statsd.Increment(TracerMetricNames.Queue.DequeuedTraces, buffer.TraceCount);
+                            statsd.Increment(TracerMetricNames.Queue.DequeuedSpans, buffer.SpanCount);
+                        }
                     }
 
                     var droppedTraces = Interlocked.Exchange(ref _droppedTraces, 0);
@@ -347,7 +357,7 @@ async Task InternalBufferFlush()
                         {
                             droppedP0Traces = Interlocked.Exchange(ref _droppedP0Traces, 0);
                             droppedP0Spans = Interlocked.Exchange(ref _droppedP0Spans, 0);
-                            Log.Debug<int, int, long, long>("Flushing {Spans} spans across {Traces} traces. CanComputeStats is enabled with {DroppedP0Traces} droppedP0Traces and {DroppedP0Spans} droppedP0Spans", buffer.SpanCount, buffer.TraceCount, droppedP0Traces, droppedP0Spans);
+                            Log.Debug("Flushing {Spans} spans across {Traces} traces. CanComputeStats is enabled with {DroppedP0Traces} droppedP0Traces and {DroppedP0Spans} droppedP0Spans", buffer.SpanCount, buffer.TraceCount, droppedP0Traces, droppedP0Spans);
                             // Metrics for unsampled traces/spans already recorded
                         }
                         else
@@ -388,7 +398,7 @@ async Task InternalBufferFlush()
         private void SerializeTrace(ArraySegment<Span> spans)
         {
             // Declaring as inline method because only safe to invoke in the context of SerializeTrace
-            SpanBuffer SwapBuffers()
+            SpanBuffer? SwapBuffers()
             {
                 if (_activeBuffer == _frontBuffer)
                 {
@@ -525,8 +535,12 @@ private void DropTrace(ArraySegment<Span> spans)
 
             if (Volatile.Read(ref _traceMetricsEnabled))
             {
-                _statsd.Increment(TracerMetricNames.Queue.DroppedTraces);
-                _statsd.Increment(TracerMetricNames.Queue.DroppedSpans, spans.Count);
+                using var lease = _statsd.TryGetClientLease();
+                if (lease.Client is { } statsd)
+                {
+                    statsd.Increment(TracerMetricNames.Queue.DroppedTraces);
+                    statsd.Increment(TracerMetricNames.Queue.DroppedSpans, spans.Count);
+                }
             }
         }
 
@@ -589,7 +603,7 @@ private void SerializeTracesLoop()
         private readonly struct WorkItem
         {
             public readonly ArraySegment<Span> Trace;
-            public readonly Action Callback;
+            public readonly Action? Callback;
 
             public WorkItem(ArraySegment<Span> trace)
             {

tracer/src/Datadog.Trace/Agent/Api.cs

@@ -33,7 +33,7 @@ internal class Api : IApi
 
         private readonly IDatadogLogger _log;
         private readonly IApiRequestFactory _apiRequestFactory;
-        private readonly IDogStatsd _originalStatsd;
+        private readonly IStatsdManager _statsd;
         private readonly string _containerId;
         private readonly string _entityId;
         private readonly Uri _tracesEndpoint;
@@ -42,13 +42,13 @@ internal class Api : IApi
         private readonly bool _partialFlushEnabled;
         private readonly SendCallback<SendStatsState> _sendStats;
         private readonly SendCallback<SendTracesState> _sendTraces;
-        private IDogStatsd _statsd;
         private string _cachedResponse;
         private string _agentVersion;
+        private bool _healthMetricsEnabled;
 
         public Api(
             IApiRequestFactory apiRequestFactory,
-            IDogStatsd statsd,
+            IStatsdManager statsd,
             Action<Dictionary<string, float>> updateSampleRates,
             bool partialFlushEnabled,
             bool healthMetricsEnabled,
@@ -60,12 +60,13 @@ public Api(
             _sendStats = SendStatsAsyncImpl;
             _sendTraces = SendTracesAsyncImpl;
             _updateSampleRates = updateSampleRates;
-            _originalStatsd = statsd;
+            _statsd = statsd;
             ToggleTracerHealthMetrics(healthMetricsEnabled);
             _containerId = ContainerMetadata.GetContainerId();
             _entityId = ContainerMetadata.GetEntityId();
             _apiRequestFactory = apiRequestFactory;
             _partialFlushEnabled = partialFlushEnabled;
+            _healthMetricsEnabled = healthMetricsEnabled;
             _tracesEndpoint = _apiRequestFactory.GetEndpoint(TracesPath);
             _log.Debug("Using traces endpoint {TracesEndpoint}", _tracesEndpoint.ToString());
             _statsEndpoint = _apiRequestFactory.GetEndpoint(StatsPath);
@@ -84,7 +85,8 @@ private enum SendResult
         [MemberNotNull(nameof(_statsd))]
         public void ToggleTracerHealthMetrics(bool enabled)
         {
-            Interlocked.Exchange(ref _statsd, enabled ? _originalStatsd : null);
+            Volatile.Write(ref _healthMetricsEnabled, enabled);
+            _statsd.SetRequired(StatsdConsumer.TraceApi, enabled);
         }
 
         public Task<bool> SendStatsAsync(StatsBuffer stats, long bucketDuration)
@@ -309,7 +311,9 @@ private async Task<SendResult> SendTracesAsyncImpl(IApiRequest request, bool fin
 
             try
             {
-                var healthStats = Volatile.Read(ref _statsd);
+                var healthMetricsEnabled = Volatile.Read(ref _healthMetricsEnabled);
+                using var lease = healthMetricsEnabled ? _statsd.TryGetClientLease() : default;
+                var healthStats = healthMetricsEnabled ? lease.Client : null;
                 try
                 {
                     TelemetryFactory.Metrics.RecordCountTraceApiRequests();
@@ -332,7 +336,7 @@ private async Task<SendResult> SendTracesAsyncImpl(IApiRequest request, bool fin
                     string[] tags = { $"status:{response.StatusCode}" };
 
                     // count every response, grouped by status code
-                    healthStats?.Increment(TracerMetricNames.Api.Responses, tags: tags);
+                    healthStats.Increment(TracerMetricNames.Api.Responses, tags: tags);
                 }
 
                 TelemetryFactory.Metrics.RecordCountTraceApiResponses(response.GetTelemetryStatusCodeMetricTag());

tracer/src/Datadog.Trace/Agent/IAgentWriter.cs

@@ -3,6 +3,8 @@
 // This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2017 Datadog, Inc.
 // </copyright>
 
+#nullable enable
+
 using System;
 using System.Threading.Tasks;
 

tracer/src/Datadog.Trace/Agent/ManagedApi.cs

@@ -11,6 +11,7 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Datadog.Trace.Configuration;
+using Datadog.Trace.DogStatsd;
 using Datadog.Trace.Vendors.StatsdClient;
 
 namespace Datadog.Trace.Agent;
@@ -24,7 +25,7 @@ internal class ManagedApi : IApi
 
     public ManagedApi(
         TracerSettings.SettingsManager settings,
-        IDogStatsd statsd,
+        IStatsdManager statsd,
         Action<Dictionary<string, float>> updateSampleRates,
         bool partialFlushEnabled)
     {

tracer/src/Datadog.Trace/Ci/Agent/ApmAgentWriter.cs

@@ -11,6 +11,7 @@
 using Datadog.Trace.Agent.DiscoveryService;
 using Datadog.Trace.Ci.EventModel;
 using Datadog.Trace.Configuration;
+using Datadog.Trace.DogStatsd;
 
 namespace Datadog.Trace.Ci.Agent;
 
@@ -30,16 +31,17 @@ public ApmAgentWriter(TracerSettings settings, Action<Dictionary<string, float>>
         var partialFlushEnabled = settings.PartialFlushEnabled;
         // CI Vis doesn't allow reconfiguration, so don't need to subscribe to changes
         var apiRequestFactory = TracesTransportStrategy.Get(settings.Manager.InitialExporterSettings);
-        var api = new Api(apiRequestFactory, null, updateSampleRates, partialFlushEnabled, healthMetricsEnabled: false);
+        var statsdManager = new StatsdManager(settings);
+        var api = new Api(apiRequestFactory, statsdManager, updateSampleRates, partialFlushEnabled, healthMetricsEnabled: false);
         var statsAggregator = StatsAggregator.Create(api, settings, discoveryService);
 
-        _agentWriter = new AgentWriter(api, statsAggregator, null, maxBufferSize: maxBufferSize, apmTracingEnabled: settings.ApmTracingEnabled, initialTracerMetricsEnabled: settings.Manager.InitialMutableSettings.TracerMetricsEnabled);
+        _agentWriter = new AgentWriter(api, statsAggregator, statsdManager, maxBufferSize: maxBufferSize, apmTracingEnabled: settings.ApmTracingEnabled, initialTracerMetricsEnabled: settings.Manager.InitialMutableSettings.TracerMetricsEnabled);
     }
 
     // Internal for testing
-    internal ApmAgentWriter(IApi api, int maxBufferSize = DefaultMaxBufferSize)
+    internal ApmAgentWriter(IApi api, IStatsdManager statsdManager, int maxBufferSize = DefaultMaxBufferSize)
     {
-        _agentWriter = new AgentWriter(api, null, null, maxBufferSize: maxBufferSize);
+        _agentWriter = new AgentWriter(api, null, statsdManager, maxBufferSize: maxBufferSize);
     }
 
     public void WriteEvent(IEvent @event)

tracer/src/Datadog.Trace/Ci/TestOptimizationTracerManager.cs

@@ -13,6 +13,7 @@
 using Datadog.Trace.Ci.EventModel;
 using Datadog.Trace.Configuration;
 using Datadog.Trace.DataStreamsMonitoring;
+using Datadog.Trace.DogStatsd;
 using Datadog.Trace.Logging;
 using Datadog.Trace.Logging.DirectSubmission;
 using Datadog.Trace.Logging.TracerFlare;
@@ -32,7 +33,7 @@ public TestOptimizationTracerManager(
             TracerSettings settings,
             IAgentWriter agentWriter,
             IScopeManager scopeManager,
-            IDogStatsd statsd,
+            IStatsdManager statsd,
             RuntimeMetricsWriter runtimeMetricsWriter,
             DirectLogSubmissionManager logSubmissionManager,
             ITelemetryController telemetry,
@@ -148,7 +149,7 @@ public LockedManager(
                 TracerSettings settings,
                 IAgentWriter agentWriter,
                 IScopeManager scopeManager,
-                IDogStatsd statsd,
+                IStatsdManager statsd,
                 RuntimeMetricsWriter runtimeMetricsWriter,
                 DirectLogSubmissionManager logSubmissionManager,
                 ITelemetryController telemetry,

tracer/src/Datadog.Trace/Ci/TestOptimizationTracerManagerFactory.cs

@@ -14,6 +14,7 @@
 using Datadog.Trace.Ci.Sampling;
 using Datadog.Trace.Configuration;
 using Datadog.Trace.DataStreamsMonitoring;
+using Datadog.Trace.DogStatsd;
 using Datadog.Trace.Logging.DirectSubmission;
 using Datadog.Trace.Logging.TracerFlare;
 using Datadog.Trace.RemoteConfigurationManagement;
@@ -41,7 +42,7 @@ protected override TracerManager CreateTracerManagerFrom(
             TracerSettings settings,
             IAgentWriter agentWriter,
             IScopeManager scopeManager,
-            IDogStatsd statsd,
+            IStatsdManager statsd,
             RuntimeMetricsWriter runtimeMetrics,
             DirectLogSubmissionManager logSubmissionManager,
             ITelemetryController telemetry,
@@ -87,7 +88,7 @@ protected override ITraceSampler GetSampler(TracerSettings settings)
 
         protected override bool ShouldEnableRemoteConfiguration(TracerSettings settings) => false;
 
-        protected override IAgentWriter GetAgentWriter(TracerSettings settings, IDogStatsd statsd, Action<Dictionary<string, float>> updateSampleRates, IDiscoveryService discoveryService, TelemetrySettings telemetrySettings)
+        protected override IAgentWriter GetAgentWriter(TracerSettings settings, IStatsdManager statsd, Action<Dictionary<string, float>> updateSampleRates, IDiscoveryService discoveryService, TelemetrySettings telemetrySettings)
         {
             // Check for agentless scenario
             if (_settings.Agentless)

tracer/src/Datadog.Trace/DogStatsd/IStatsdManager.cs

@@ -0,0 +1,26 @@
+// <copyright file="IStatsdManager.cs" company="Datadog">
+// Unless explicitly stated otherwise all files in this repository are licensed under the Apache 2 License.
+// This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2017 Datadog, Inc.
+// </copyright>
+
+#nullable enable
+using System;
+using Datadog.Trace.Vendors.StatsdClient;
+
+namespace Datadog.Trace.DogStatsd;
+
+internal interface IStatsdManager : IDisposable
+{
+    /// <summary>
+    /// Obtain a <see cref="StatsdManager.StatsdClientLease"/> for accessing a <see cref="IDogStatsd"/> instance.
+    /// The lease must be disposed after all references to the client have gone.
+    /// </summary>
+    StatsdManager.StatsdClientLease TryGetClientLease();
+
+    /// <summary>
+    /// Called by users of <see cref="StatsdManager"/> to indicate that a "live" client is required.
+    /// Each unique consumer of <see cref="StatsdManager"/> should set a different
+    /// <see cref="StatsdConsumer"/> value.
+    /// </summary>
+    void SetRequired(StatsdConsumer consumer, bool enabled);
+}

tracer/src/Datadog.Trace/DogStatsd/NoOpStatsd.cs

@@ -8,7 +8,7 @@
 
 namespace Datadog.Trace.DogStatsd
 {
-    internal class NoOpStatsd : IDogStatsd
+    internal sealed class NoOpStatsd : IDogStatsd
     {
         public static readonly NoOpStatsd Instance = new();