Add trace count header for statistics agent side (DataDog#465)

* Add trace count header for statistics agent side

* Add global assertion to MockTracerAgent

src/Datadog.Trace/Agent/Api.cs

@@ -14,16 +14,16 @@ internal class Api : IApi
     {
         private const string TracesPath = "/v0.4/traces";
 
-        private static readonly ILog _log = LogProvider.For<Api>();
-        private static readonly SerializationContext _serializationContext = new SerializationContext();
-        private static readonly SpanMessagePackSerializer Serializer = new SpanMessagePackSerializer(_serializationContext);
+        private static readonly ILog Log = LogProvider.For<Api>();
+        private static readonly SerializationContext SerializationContext = new SerializationContext();
+        private static readonly SpanMessagePackSerializer Serializer = new SpanMessagePackSerializer(SerializationContext);
 
         private readonly Uri _tracesEndpoint;
         private readonly HttpClient _client;
 
         static Api()
         {
-            _serializationContext.ResolveSerializer += (sender, eventArgs) =>
+            SerializationContext.ResolveSerializer += (sender, eventArgs) =>
             {
                 if (eventArgs.TargetType == typeof(Span))
                 {
@@ -54,7 +54,7 @@ public Api(Uri baseEndpoint, DelegatingHandler delegatingHandler = null)
 
         public async Task SendTracesAsync(IList<List<Span>> traces)
         {
-            // retry up to 5 times with exponential backoff
+            // retry up to 5 times with exponential back-off
             var retryLimit = 5;
             var retryCount = 1;
             var sleepDuration = 100; // in milliseconds
@@ -65,9 +65,12 @@ public async Task SendTracesAsync(IList<List<Span>> traces)
 
                 try
                 {
+                    var traceIds = GetUniqueTraceIds(traces);
+
                     // re-create content on every retry because some versions of HttpClient always dispose of it, so we can't reuse.
-                    using (var content = new MsgPackContent<IList<List<Span>>>(traces, _serializationContext))
+                    using (var content = new MsgPackContent<IList<List<Span>>>(traces, SerializationContext))
                     {
+                        content.Headers.Add(AgentHttpHeaderNames.TraceCount, traceIds.Count.ToString());
                         responseMessage = await _client.PostAsync(_tracesEndpoint, content).ConfigureAwait(false);
                         responseMessage.EnsureSuccessStatusCode();
                     }
@@ -77,7 +80,7 @@ public async Task SendTracesAsync(IList<List<Span>> traces)
                     if (retryCount >= retryLimit)
                     {
                         // stop retrying
-                        _log.ErrorException("An error occurred while sending traces to the agent at {Endpoint}", ex, _tracesEndpoint);
+                        Log.ErrorException("An error occurred while sending traces to the agent at {Endpoint}", ex, _tracesEndpoint);
                         return;
                     }
 
@@ -101,13 +104,28 @@ public async Task SendTracesAsync(IList<List<Span>> traces)
                 }
                 catch (Exception ex)
                 {
-                    _log.ErrorException("Traces sent successfully to the Agent at {Endpoint}, but an error occurred deserializing the response.", ex, _tracesEndpoint);
+                    Log.ErrorException("Traces sent successfully to the Agent at {Endpoint}, but an error occurred deserializing the response.", ex, _tracesEndpoint);
                 }
 
                 return;
             }
         }
 
+        private static HashSet<ulong> GetUniqueTraceIds(IList<List<Span>> traces)
+        {
+            var uniqueTraceIds = new HashSet<ulong>();
+
+            foreach (var trace in traces)
+            {
+                foreach (var span in trace)
+                {
+                    uniqueTraceIds.Add(span.TraceId);
+                }
+            }
+
+            return uniqueTraceIds;
+        }
+
         private static void GetFrameworkDescription(out string name, out string version)
         {
             // RuntimeInformation.FrameworkDescription returns string like ".NET Framework 4.7.2" or ".NET Core 2.1",

src/Datadog.Trace/AgentHttpHeaderNames.cs

@@ -29,5 +29,10 @@ internal static class AgentHttpHeaderNames
         /// The version of the tracer that generated this span.
         /// </summary>
         public const string TracerVersion = "Datadog-Meta-Tracer-Version";
+
+        /// <summary>
+        /// The number of unique traces per request.
+        /// </summary>
+        public const string TraceCount = "X-Datadog-Trace-Count";
     }
 }

test/Datadog.Trace.TestHelpers/MockTracerAgent.cs

@@ -62,6 +62,8 @@ public MockTracerAgent(int port = 8126, int retries = 5)
 
         public IImmutableList<Span> Spans { get; private set; } = ImmutableList<Span>.Empty;
 
+        public IImmutableList<HttpListenerRequest> RawRequests { get; private set; } = ImmutableList<HttpListenerRequest>.Empty;
+
         /// <summary>
         /// Wait for the given number of spans to appear.
         /// </summary>
@@ -100,6 +102,23 @@ public IImmutableList<Span> WaitForSpans(
                 Thread.Sleep(500);
             }
 
+            foreach (var request in RawRequests)
+            {
+                // This is the place to check against headers we expect
+                AssertHeader(
+                    request,
+                    "X-Datadog-Trace-Count",
+                    header =>
+                    {
+                        if (int.TryParse(header, out int traceCount))
+                        {
+                            return traceCount > 0;
+                        }
+
+                        return false;
+                    });
+            }
+
             if (!returnAllOperations)
             {
                 relevantSpans =
@@ -151,6 +170,24 @@ private static List<Span> ToSpans(dynamic data)
             return new List<Span>();
         }
 
+        private void AssertHeader(
+            HttpListenerRequest request,
+            string headerKey,
+            Func<string, bool> assertion)
+        {
+            var header = request.Headers.Get(headerKey);
+
+            if (string.IsNullOrEmpty(header))
+            {
+                throw new Exception($"Every submission to the agent should have a {headerKey} header.");
+            }
+
+            if (!assertion(header))
+            {
+                throw new Exception($"Failed assertion for {headerKey} on {header}");
+            }
+        }
+
         private void HandleHttpRequests()
         {
             while (_listener.IsListening)
@@ -166,6 +203,7 @@ private void HandleHttpRequests()
                         // we only need to lock when replacing the span collection,
                         // not when reading it because it is immutable
                         Spans = Spans.AddRange(spans);
+                        RawRequests = RawRequests.Add(ctx.Request);
                     }
 
                     ctx.Response.ContentType = "application/json";