io.opentelemetry.javaagent.instrumentation.kafkaclients.v0_11.TracingIterable Maven / Gradle / Ivy

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

package io.opentelemetry.javaagent.instrumentation.kafkaclients.v0_11;

import io.opentelemetry.instrumentation.kafka.internal.KafkaConsumerContext;
import io.opentelemetry.javaagent.bootstrap.kafka.KafkaClientsConsumerProcessTracing;
import java.util.Iterator;
import org.apache.kafka.clients.consumer.ConsumerRecord;

public class TracingIterable implements Iterable> {
  private final Iterable> delegate;
  private final KafkaConsumerContext consumerContext;
  private boolean firstIterator = true;

  protected TracingIterable(
      Iterable> delegate, KafkaConsumerContext consumerContext) {
    this.delegate = delegate;
    this.consumerContext = consumerContext;
  }

  public static  Iterable> wrap(
      Iterable> delegate, KafkaConsumerContext consumerContext) {
    if (KafkaClientsConsumerProcessTracing.wrappingEnabled()) {
      return new TracingIterable<>(delegate, consumerContext);
    }
    return delegate;
  }

  @Override
  public Iterator> iterator() {
    Iterator> it;
    // We should only return one iterator with tracing.
    // However, this is not thread-safe, but usually the first (hopefully only) traversal of
    // ConsumerRecords is performed in the same thread that called poll()
    if (firstIterator) {
      it = TracingIterator.wrap(delegate.iterator(), consumerContext);
      firstIterator = false;
    } else {
      it = delegate.iterator();
    }

    return it;
  }
}