Rebus.TraceContext.W3C.Outgoing

RebusTraceContextOutgoing.cs

namespace RebusTraceContext
{
    using System;
    using System.Diagnostics;
    using System.Threading.Tasks;
    using Rebus.Config;
    using Rebus.Messages;
    using Rebus.Pipeline;
    using Rebus.Pipeline.Send;

    public static class OptionsConfigurerExtensions
    {
        public static void AddTraceContext(this OptionsConfigurer configurer)
        {
            configurer.Decorate<IPipeline>(c =>
            {
                var pipeline = c.Get<IPipeline>();
                var step = new AddTraceContextOutgoingStep();
                return new PipelineStepInjector(pipeline)
                    .OnSend(step, PipelineRelativePosition.Before, typeof(SerializeOutgoingMessageStep));
            });
        }
    }

    public class AddTraceContextOutgoingStep : IOutgoingStep
    {
        public async Task Process(OutgoingStepContext context, Func<Task> next)
        {
            var activity = Activity.Current;
            var message = context.Load<Message>();

            if (activity.IdFormat == ActivityIdFormat.W3C)
            {
                if (!message.Headers.ContainsKey("traceparent"))
                {
                    message.Headers["traceparent"] = activity.Id;
                    if (activity.TraceStateString != null)
                    {
                        message.Headers["tracestate"] = activity.TraceStateString;
                    }
                }
            }
            else
            {
                if (!message.Headers.ContainsKey("Request-Id"))
                {
                    message.Headers["Request-Id"] = activity.Id;
                }
            }

            await next();
        }
    }
}