How to add distributed tracing support to SqsTemplate

[jkuipers]

Hey,

I'm planning to upgrade from 2.4 to 3.0. For 2.4, I added some additional functionality to the QueueMessagingTemplate and I'm now looking into if and if so, how I need to port that for the 3.0 SqsTemplate.

Two things I did that are now built-in are support for async sending (I needed that because without this the underlying AmazonSQSBufferedAsyncClient would block for some time when sending multiple messages from a single thread, creating huge delays) and adding a header containing the payload type (to support multiple message payload types), so that's nice!

Another thing is that I added support to propagate the current trace information. This support is very important to us: it allows us to correlate logging between senders and receivers and also to quickly find the receiver's logging for dead-lettered messages so that we can understand why the processing failed.
In 2.4 I did that by subclassing the QueueMessagingTemplate and overriding its doSend(QueueMessageChannel destination, Message<?> message) method. There I could "inject" the message with the right header (nowadays that would be the "traceparent" header for W3C Trace Context-based propagation, before that it was a Brave b3 header).
I'm looking at the SqsTemplate, but it seems much trickier there. I don't think I can subclass the type as it has a private constructor, making it effectively final. The only hook I see is to configure an instance of the template with my own subtype of SqsMessagingMessageConverter that would in turn return a custom subtype of the SqsHeaderMapper which could add the necessary message attribute.
I'm also not very experienced with all the async future-driven code, so I'm also wondering if an approach like that would even work because I'm afraid that it's possible that the code might not even be running in a thread that has access to the current trace context to propagate.

What are your thoughts regarding an extension like this? Is there a better way that I'm overlooking to enrich the message that's being sent transparently while ensuring that I can access the current tracing context? I could also let my client code call an adapter that would create a full Message with the required header to then pass that to the template, but I'd much rather let clients use the SqsTemplate (or the interfaces that it implements) directly.

Thanks in advance for any advice!

[tomazfernandes]

Hi @jkuipers, nice to see you around.

I believe the SqsMessageMapper should be the proper place to handle this, and IIUC you should be able to use the setAdditionalHeadersFunction to fetch the header from ThreadLocal / MDC and create a MessageHeader out of it.

If you use the blocking send methods the thread you shouldn't run into any issues - the thread that reaches the converter will be the same that calls the method, and after the message is sent the same thread should continue execution.

If you use the async send methods though it might be trickier - you should be able to reach the converter on the calling thread, but after that you'll hop threads so you must be mindful of cleaning up any ThreadLocal information.

How does that sound so far? I'd be interested in having a clean way of solving this, or it this is clean enough we should at least document it properly.

Thanks.

[jkuipers]

Thanks @tomazfernandes, I'll give it a go then to see if I can get this to work via the mapper. I noticed that there's an open issue to add tracing support; if I can come up with a clean solution I'll try to contribute something there.

[jkuipers]

Something like this seems to do the trick (Tracing is a brave.Tracing type):

    @Bean
    SqsTemplate sqsTemplate(SqsAsyncClient sqsAsyncClient, ObjectMapper objectMapper, Tracing tracing) {
        var injector = tracing.propagation().injector(
                (Propagation.Setter<Map<String, Object>, String>) (headersMap, key, value) -> {
                    // only propagate the traceparent attr, no additional baggage, since SQS only supports a max of 10 attrs
                    if (key.startsWith("trace")) headersMap.put(key, value);
                });
        var converter = new SqsMessagingMessageConverter() {
            @Override
            public Message fromMessagingMessage(org.springframework.messaging.Message<?> message, MessageConversionContext context) {
                Span span = tracing.tracer().nextSpan()
                        .kind(Span.Kind.PRODUCER)
                        .name("send")
                        .remoteServiceName("sqs")
                        .start();
                Map<String, Object> newHeaders = new HashMap<>();
                injector.inject(span.context(), newHeaders);
                span.finish();
                return super.fromMessagingMessage(MessageHeaderUtils.addHeadersIfAbsent(message, newHeaders), context);
            }
        };
        converter.setObjectMapper(objectMapper);

        return SqsTemplate.builder()
                .sqsAsyncClient(sqsAsyncClient)
                .messageConverter(converter)
                .build();
    }

It might need some cleanup, but for my use will do for now. I'll first continue to investigate how to restore the trace context on receiving a message using a listener (rather than the template)

[jkuipers]

For receiving I came up with a solution that lets the endpoint wrap the MessagingMessageListenerAdapter with an adapter that calls the delegate in the context of a span created based on the attributes of the message. To ensure that this endpoint is used I extended the SqsListenerAnnotationBeanPostProcessor. Overall this then looks like this.

Bean definition for our customBPP:

    @Bean(name = SqsBeanNames.SQS_LISTENER_ANNOTATION_BEAN_POST_PROCESSOR_BEAN_NAME)
    TracingSqsListenerAnnotationBeanPostProcessor tracingSqsLABPP(Tracing tracing) {
        return new TracingSqsListenerAnnotationBeanPostProcessor(tracing);
    }

Code for that type:

public class TracingSqsListenerAnnotationBeanPostProcessor extends SqsListenerAnnotationBeanPostProcessor {
    private final Tracing tracing;

    public TracingSqsListenerAnnotationBeanPostProcessor(Tracing tracing) {
        this.tracing = tracing;
    }

    /**
     * Overrides parent method to ensure that our custom endpoint with tracing support is returned.
     *
     * @param sqsListenerAnnotation
     * @return
     */
    @Override
    protected Endpoint createEndpoint(SqsListener sqsListenerAnnotation) {
        return new TracingSqsEndpoint.TracingSqsEndpointBuilder(tracing)
                .queueNames(resolveEndpointNames(sqsListenerAnnotation.value()))
                .factoryBeanName(resolveAsString(sqsListenerAnnotation.factory(), "factory"))
                .id(getEndpointId(sqsListenerAnnotation.id()))
                .pollTimeoutSeconds(resolveAsInteger(sqsListenerAnnotation.pollTimeoutSeconds(), "pollTimeoutSeconds"))
                .maxMessagesPerPoll(resolveAsInteger(sqsListenerAnnotation.maxMessagesPerPoll(), "maxMessagesPerPoll"))
                .maxConcurrentMessages(
                        resolveAsInteger(sqsListenerAnnotation.maxConcurrentMessages(), "maxConcurrentMessages"))
                .messageVisibility(
                        resolveAsInteger(sqsListenerAnnotation.messageVisibilitySeconds(), "messageVisibility"))
                .build();
    }
}

code for the custom endpoint:

public class TracingSqsEndpoint extends SqsEndpoint {

    private final Tracing tracing;

    protected TracingSqsEndpoint(SqsEndpointBuilder builder, Tracing tracing) {
        super(builder);
        this.tracing = tracing;
    }

    @Override
    protected <T> MessageListener<T> createMessageListenerInstance(InvocableHandlerMethod handlerMethod) {
        return new TracingMessageListenerWrapper<>(super.createMessageListenerInstance(handlerMethod), tracing);
    }

    public static class TracingSqsEndpointBuilder extends SqsEndpointBuilder {

        private final Tracing tracing;

        public TracingSqsEndpointBuilder(Tracing tracing) {
            this.tracing = tracing;
        }

        @Override
        public SqsEndpoint build() {
            return new TracingSqsEndpoint(this, tracing);
        }
    }
}

and finally the code for the wrapper that performs the actual tracing logic:

public class TracingMessageListenerWrapper<T> implements MessageListener<T> {

    private static final Propagation.Getter<MessageHeaders, String> GETTER =
            (headers, key) -> (String) headers.get(key);

    private final MessageListener<T> delegate;
    private final TraceContext.Extractor<MessageHeaders> extractor;
    private final Tracer tracer;

    public TracingMessageListenerWrapper(MessageListener<T> delegate, Tracing tracing) {
        this.delegate = delegate;
        this.extractor = tracing.propagation().extractor(GETTER);
        this.tracer = tracing.tracer();
    }

    @Override
    public void onMessage(Message<T> message) {
        TraceContextOrSamplingFlags extracted = extractor.extract(message.getHeaders());
        Span span = tracer.nextSpan(extracted)
                .kind(CONSUMER)
                .name("on-message")
                .remoteServiceName("sqs")
                .start();

        try (Tracer.SpanInScope ws = tracer.withSpanInScope(span)) {
            delegate.onMessage(message);
        } catch (Throwable t) {
            span.error(t);
            throw t;
        } finally {
            span.finish();
        }
    }

    @Override
    public void onMessage(Collection<Message<T>> messages) {
        delegate.onMessage(messages);
    }
}

I don't think it makes sense to support tracing when handling a collection of messages, since each will typically correspond to a different trace.
This setup seems to work. The framework might make it a bit easier to get this going, but I think for a point solution like mine (instead of a generic framework-provided solution) this is OK.

[tomazfernandes]

Yup, sure looks solid, nice work!

Just curious, why are you using Brave instead of Micrometer 2.0 and the Observability API?

[jkuipers]

Good question ;) It's because this is a port of my current code that was written before the new Micrometer Observability API existed, in the Spring Cloud Sleuth days.
I am considering to figure out how to port this, but this code still works perfectly fine with the latest Boot and Micrometer as it still configures the Brave Tracing bean. Also, I do not want/need support for metrics here: SQS provides its own metrics already, I don't need the application to produce info about messages sent/received in addition to that. I really only care about propagating the trace context so that I can correlate logging.

FYI, we're using the log correlation in two ways: obviously to correlate the logging of the sender and receiver, but also to see the receiver's logging when a message cannot be processed and is dead-lettered. I've written a custom application to see and manage dead-lettered messages, which allows the user to move individual or all messages from a DLQ back to the source queue for retry, or to purge the DLQ. This tool also allows you to click on a dead-lettered message, which then takes you straight to the related logging in Datadog based on the trace ID that we parse from the message's traceparent attribute. Works really well and is very important for our functional administration of the integration platform that we provide, which is why I need this functionality.

[maciekgaszynski]

My (a little bit rough) attempt using Micrometer Observability API:

class TracingSqsEndpoint extends SqsEndpoint {

  private final ObservationRegistry observationRegistry;

  protected TracingSqsEndpoint(SqsEndpointBuilder builder, ObservationRegistry tracing) {
    super(builder);
    this.observationRegistry = tracing;
  }

  @Override
  protected <T> MessageListener<T> createMessageListenerInstance(
      InvocableHandlerMethod handlerMethod) {
    return new TracingMessageListenerWrapper<>(
        super.createMessageListenerInstance(handlerMethod), observationRegistry);
  }

  @RequiredArgsConstructor
  public static class TracingSqsEndpointBuilder extends SqsEndpointBuilder {

    private final ObservationRegistry observationRegistry;

    @Override
    public SqsEndpoint build() {
      return new TracingSqsEndpoint(this, observationRegistry);
    }
  }
}

Wrapper adapted to use ObservationRegistry with the most important change of introducing SqsReceiverContext:

@RequiredArgsConstructor
class TracingMessageListenerWrapper<T> implements MessageListener<T> {
  private final MessageListener<T> delegate;
  private final ObservationRegistry observationRegistry;

  @Override
  public void onMessage(Message<T> message) {

    SqsReceiverContext sqsReceiverContext = new SqsReceiverContext(message);
    Observation observation =
        Observation.start("sqs.receive", () -> sqsReceiverContext, observationRegistry);

    observation.observeChecked(() -> delegate.onMessage(message));
  }

  static class SqsReceiverContext extends ReceiverContext<Message<?>> {

    public SqsReceiverContext(Message<?> message) {
      super(((carrier, key) -> carrier.getHeaders().get(key, String.class)));
      setCarrier(message);
    }
  }
}

Post processor taken as is with ObservationRegistry.

@RequiredArgsConstructor
class TracingSqsListenerAnnotationBeanPostProcessor extends SqsListenerAnnotationBeanPostProcessor {

  private final ObjectProvider<ObservationRegistry> observationRegistry;

  /**
   * Overrides parent method to ensure that our custom endpoint with tracing support is returned.
   */
  @Override
  protected Endpoint createEndpoint(SqsListener sqsListenerAnnotation) {
    return new TracingSqsEndpoint.TracingSqsEndpointBuilder(observationRegistry.getIfAvailable())
        .queueNames(resolveEndpointNames(sqsListenerAnnotation.value()))
        .factoryBeanName(resolveAsString(sqsListenerAnnotation.factory(), "factory"))
        .id(getEndpointId(sqsListenerAnnotation.id()))
        .pollTimeoutSeconds(
            resolveAsInteger(sqsListenerAnnotation.pollTimeoutSeconds(), "pollTimeoutSeconds"))
        .maxMessagesPerPoll(
            resolveAsInteger(sqsListenerAnnotation.maxMessagesPerPoll(), "maxMessagesPerPoll"))
        .maxConcurrentMessages(
            resolveAsInteger(
                sqsListenerAnnotation.maxConcurrentMessages(), "maxConcurrentMessages"))
        .messageVisibility(
            resolveAsInteger(sqsListenerAnnotation.messageVisibilitySeconds(), "messageVisibility"))
        .build();
  }
}

and sender part with ugly traceparent generation:

  @Bean(name = SqsBeanNames.SQS_LISTENER_ANNOTATION_BEAN_POST_PROCESSOR_BEAN_NAME)
  TracingSqsListenerAnnotationBeanPostProcessor tracingSqsLABPP(
      ObjectProvider<ObservationRegistry> observationRegistry) {
    return new TracingSqsListenerAnnotationBeanPostProcessor(observationRegistry);
  }

  @Bean
  SqsTemplate sqsTemplate(
      SqsAsyncClient sqsAsyncClient, ObjectMapper objectMapper, Tracer tracing) {
    var converter =
        new SqsMessagingMessageConverter() {
          @Override
          public software.amazon.awssdk.services.sqs.model.Message fromMessagingMessage(
              Message<?> message, MessageConversionContext context) {
            Span span = tracing.nextSpan().name("send").remoteServiceName("sqs").start();
            Map<String, Object> newHeaders =
                Map.of(
                    "traceparent",
                    "00-" + span.context().traceId() + "-" + span.context().spanId() + "-00");
            span.end();
            return super.fromMessagingMessage(
                MessageHeaderUtils.addHeadersIfAbsent(message, newHeaders), context);
          }
        };
    converter.setObjectMapper(objectMapper);

    return SqsTemplate.builder().sqsAsyncClient(sqsAsyncClient).messageConverter(converter).build();
  }