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/src/Datadog.Trace/Agent/AgentWriter.cs

@@ -28,7 +28,7 @@ internal class AgentWriter : IAgentWriter
         private static readonly ArraySegment<byte> EmptyPayload = new(new byte[] { 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>();
@@ -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);
@@ -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);
+                }
             }
         }
 
@@ -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);
@@ -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);
+                }
             }
         }
 

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,8 @@ private async Task<SendResult> SendTracesAsyncImpl(IApiRequest request, bool fin
 
             try
             {
-                var healthStats = Volatile.Read(ref _statsd);
+                using var lease = Volatile.Read(ref _healthMetricsEnabled) ? _statsd.TryGetClientLease() : default;
+                var healthStats = _healthMetricsEnabled ? lease.Client : null;
                 try
                 {
                     TelemetryFactory.Metrics.RecordCountTraceApiRequests();
@@ -332,7 +335,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/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/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();
 

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

@@ -0,0 +1,20 @@
+// <copyright file="StatsdConsumer.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;
+
+namespace Datadog.Trace.DogStatsd;
+
+[Flags]
+internal enum StatsdConsumer
+{
+    // None = 0, Must not use this
+    // Define bits per consumer:
+    RuntimeMetricsWriter = 1 << 0,
+    TraceApi = 1 << 1,
+    AgentWriter = 1 << 2,
+}