io.opentelemetry.instrumentation.spring.kafka.v2_7.InstrumentedBatchInterceptor Maven / Gradle / Ivy

Go to download
/*
 * Copyright The OpenTelemetry Authors
 * SPDX-License-Identifier: Apache-2.0
 */

package io.opentelemetry.instrumentation.spring.kafka.v2_7;

import io.opentelemetry.context.Context;
import io.opentelemetry.context.Scope;
import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
import io.opentelemetry.instrumentation.api.util.VirtualField;
import io.opentelemetry.instrumentation.kafka.internal.KafkaConsumerContext;
import io.opentelemetry.instrumentation.kafka.internal.KafkaConsumerContextUtil;
import io.opentelemetry.instrumentation.kafka.internal.KafkaReceiveRequest;
import io.opentelemetry.javaagent.tooling.muzzle.NoMuzzle;
import java.lang.ref.WeakReference;
import javax.annotation.Nullable;
import org.apache.kafka.clients.consumer.Consumer;
import org.apache.kafka.clients.consumer.ConsumerRecords;
import org.springframework.kafka.listener.BatchInterceptor;

final class InstrumentedBatchInterceptor implements BatchInterceptor {

  private static final VirtualField, State> stateField =
      VirtualField.find(ConsumerRecords.class, State.class);
  private static final ThreadLocal>> lastProcessed =
      new ThreadLocal<>();

  private final Instrumenter batchProcessInstrumenter;
  @Nullable private final BatchInterceptor decorated;

  InstrumentedBatchInterceptor(
      Instrumenter batchProcessInstrumenter,
      @Nullable BatchInterceptor decorated) {
    this.batchProcessInstrumenter = batchProcessInstrumenter;
    this.decorated = decorated;
  }

  @Override
  public ConsumerRecords intercept(ConsumerRecords records, Consumer consumer) {
    Context parentContext = getParentContext(records);

    KafkaReceiveRequest request = KafkaReceiveRequest.create(records, consumer);
    if (batchProcessInstrumenter.shouldStart(parentContext, request) && !skipProcessing(records)) {
      Context context = batchProcessInstrumenter.start(parentContext, request);
      Scope scope = context.makeCurrent();
      stateField.set(records, State.create(request, context, scope));
    }

    return decorated == null ? records : decorated.intercept(records, consumer);
  }

  private static boolean skipProcessing(ConsumerRecords records) {
    // When retrying failed listener interceptors work as expected only in the earlier versions that
    // we test (e.g spring-kafka:2.7.1). In later versions interceptor isn't called at all during
    // the retry, which results in missing process span, or worse, the intercept method is called,
    // but neither success nor failure is called, which results in a context leak. Here we attempt
    // to prevent the context leak by observing whether intercept is called with the same
    // ConsumerRecords as on previous call, and if it is, we skip creating the process span.
    WeakReference> reference = lastProcessed.get();
    return reference != null && reference.get() == records;
  }

  private static Context getParentContext(ConsumerRecords records) {
    KafkaConsumerContext consumerContext = KafkaConsumerContextUtil.get(records);
    Context receiveContext = consumerContext.getContext();

    // use the receive CONSUMER span as parent if it's available
    return receiveContext != null ? receiveContext : Context.current();
  }

  @Override
  public void success(ConsumerRecords records, Consumer consumer) {
    try {
      if (decorated != null) {
        decorated.success(records, consumer);
      }
    } finally {
      end(records, null);
    }
  }

  @Override
  public void failure(ConsumerRecords records, Exception exception, Consumer consumer) {
    try {
      if (decorated != null) {
        decorated.failure(records, exception, consumer);
      }
    } finally {
      end(records, exception);
    }
  }

  private void end(ConsumerRecords records, @Nullable Throwable error) {
    State state = stateField.get(records);
    stateField.set(records, null);
    if (state != null) {
      KafkaReceiveRequest request = state.request();
      state.scope().close();
      batchProcessInstrumenter.end(state.context(), request, null, error);
      lastProcessed.set(new WeakReference<>(records));
    }
  }

  @NoMuzzle // method was added in 2.8.0
  @Override
  public void setupThreadState(Consumer consumer) {
    if (decorated != null) {
      decorated.setupThreadState(consumer);
    }
  }

  @NoMuzzle // method was added in 2.8.0
  @Override
  public void clearThreadState(Consumer consumer) {
    if (decorated != null) {
      decorated.clearThreadState(consumer);
    }
  }
}