Rebus.TraceContext.W3C.Incoming

RebusTraceContextIncoming.cs

namespace RebusTraceContext
{
    using System.Diagnostics;
    using System.Threading.Tasks;
    using Microsoft.ApplicationInsights;
    using Microsoft.ApplicationInsights.DataContracts;
    using Microsoft.Extensions.DependencyInjection;
    using Rebus.Handlers;
    using Rebus.Pipeline;

    public static class ServiceCollectionExtensions
    {
        public static IServiceCollection AddRebusTraceContextDecorator(IServiceCollection services)
        {
            services.Decorate(typeof(IHandleMessages<>), typeof(AddTraceContextDecorator<>));
            return services;
        }
    }

    public class AddTraceContextDecorator<TStep> : IHandleMessages<TStep>
    {
        private readonly IHandleMessages<TStep> _handleMessages;
        private readonly TelemetryClient _telemetryClient;

        public AddTraceContextDecorator(IHandleMessages<TStep> handleMessages, TelemetryClient telemetryClient)
        {
            _handleMessages = handleMessages;
            _telemetryClient = telemetryClient;
        }

        public async Task Handle(TStep message)
        {
            var context = MessageContext.Current;
            var activity = new Activity(typeof(TStep).Name);

            if (context.Message.Headers.TryGetValue("traceparent", out var requestId)
                || context.Message.Headers.TryGetValue("Request-Id", out requestId))
            {
                activity.SetParentId(requestId);

                if (context.Message.Headers.TryGetValue("tracestate", out var traceState))
                {
                    activity.TraceStateString = traceState;
                }
            }

            using (var operation = _telemetryClient.StartOperation<RequestTelemetry>(activity))
            {
                await _handleMessages.Handle(message);
            }
        }
    }
}