Fix kafka tests failing when consumer is not yet subscribed to topic
MODIFY
instrumentation/kafka-clients/kafka-clients-0.11/javaagent/src/test/groovy/KafkaClientBaseTest.groovy
instrumentation/kafka-clients/kafka-clients-0.11/javaagent/src/test/groovy/KafkaClientBaseTest.groovy
@@ -7,6 +7,7 @@ import io.opentelemetry.instrumentation.test.AgentInstrumentationSpecification
import org.apache.kafka.clients.admin.AdminClient
import org.apache.kafka.clients.admin.NewTopic
import org.apache.kafka.clients.consumer.Consumer
+import org.apache.kafka.clients.consumer.ConsumerRebalanceListener
import org.apache.kafka.clients.consumer.KafkaConsumer
import org.apache.kafka.clients.producer.KafkaProducer
import org.apache.kafka.clients.producer.Producer
@@ -21,6 +22,7 @@ import spock.lang.Shared
import spock.lang.Unroll
import java.time.Duration
+import java.util.concurrent.CountDownLatch
import java.util.concurrent.TimeUnit
abstract class KafkaClientBaseTest extends AgentInstrumentationSpecification {
@@ -36,6 +38,8 @@ abstract class KafkaClientBaseTest extends AgentInstrumentationSpecification {
static Producer<Integer, String> producer
@Shared
static Consumer<Integer, String> consumer
+ @Shared
+ static CountDownLatch consumerReady = new CountDownLatch(1)
def setupSpec() {
kafka = new KafkaContainer()
@@ -58,20 +62,27 @@ abstract class KafkaClientBaseTest extends AgentInstrumentationSpecification {
]
producer = new KafkaProducer<>(producerProps)
- // values copied from spring's KafkaTestUtils
def consumerProps = [
"bootstrap.servers" : kafka.bootstrapServers,
"group.id" : "test",
- "enable.auto.commit" : "false",
+ "enable.auto.commit" : "true",
"auto.commit.interval.ms": "10",
- "session.timeout.ms" : "60000",
+ "session.timeout.ms" : "30000",
"key.deserializer" : IntegerDeserializer,
"value.deserializer" : StringDeserializer
]
consumer = new KafkaConsumer<>(consumerProps)
- // assign only existing topic partition
- consumer.assign([new TopicPartition(SHARED_TOPIC, 0)])
+ consumer.subscribe([SHARED_TOPIC], new ConsumerRebalanceListener() {
+ @Override
+ void onPartitionsRevoked(Collection<TopicPartition> collection) {
+ }
+
+ @Override
+ void onPartitionsAssigned(Collection<TopicPartition> collection) {
+ consumerReady.countDown()
+ }
+ })
}
def cleanupSpec() {
@@ -85,12 +96,32 @@ abstract class KafkaClientBaseTest extends AgentInstrumentationSpecification {
when:
String message = "Testing without headers"
producer.send(new ProducerRecord<>(SHARED_TOPIC, message))
+ .get(5, TimeUnit.SECONDS)
then:
- // check that the message was received
- def records = consumer.poll(Duration.ofSeconds(5).toMillis())
+ awaitUntilConsumerIsReady()
+
+ def records = consumer.poll(Duration.ofSeconds(1).toMillis())
+ records.count() == 1
for (record in records) {
assert record.headers().iterator().hasNext() == propagationEnabled
}
}
+
+ // Kafka's eventual consistency behavior forces us to do a couple of empty poll() calls until it gets properly assigned a topic partition
+ static void awaitUntilConsumerIsReady() {
+ if (consumerReady.await(0, TimeUnit.SECONDS)) {
+ return
+ }
+ for (i in 0..<10) {
+ consumer.poll(0)
+ if (consumerReady.await(1, TimeUnit.SECONDS)) {
+ break
+ }
+ }
+ if (consumerReady.getCount() != 0) {
+ throw new AssertionError("Consumer wasn't assigned any partitions!")
+ }
+ consumer.seekToBeginning([])
+ }
}

MODIFY
instrumentation/kafka-clients/kafka-clients-0.11/javaagent/src/test/groovy/KafkaClientPropagationDisabledTest.groovy
instrumentation/kafka-clients/kafka-clients-0.11/javaagent/src/test/groovy/KafkaClientPropagationDisabledTest.groovy
@@ -35,7 +35,11 @@ class KafkaClientPropagationDisabledTest extends KafkaClientBaseTest {
}
when: "read message without context propagation"
+ awaitUntilConsumerIsReady()
def records = consumer.poll(Duration.ofSeconds(5).toMillis())
+ records.count() == 1
+
+ // iterate over records to generate spans
for (record in records) {
runWithSpan("processing") {}
}

MODIFY
instrumentation/kafka-clients/kafka-clients-0.11/javaagent/src/test/groovy/KafkaClientPropagationEnabledTest.groovy
instrumentation/kafka-clients/kafka-clients-0.11/javaagent/src/test/groovy/KafkaClientPropagationEnabledTest.groovy
@@ -9,6 +9,7 @@ import org.apache.kafka.clients.producer.ProducerRecord
import org.apache.kafka.common.TopicPartition
import java.time.Duration
+import java.util.concurrent.TimeUnit
import static io.opentelemetry.api.trace.SpanKind.CONSUMER
import static io.opentelemetry.api.trace.SpanKind.INTERNAL
@@ -26,12 +27,15 @@ class KafkaClientPropagationEnabledTest extends KafkaClientBaseTest {
} else {
runWithSpan("producer exception: " + ex) {}
}
- }
+ }.get(5, TimeUnit.SECONDS)
}
then:
- // check that the message was received
+ awaitUntilConsumerIsReady()
def records = consumer.poll(Duration.ofSeconds(5).toMillis())
+ records.count() == 1
+
+ // iterate over records to generate spans
for (record in records) {
runWithSpan("processing") {
assert record.value() == greeting
@@ -106,11 +110,14 @@ class KafkaClientPropagationEnabledTest extends KafkaClientBaseTest {
def "test pass through tombstone"() {
when:
- producer.send(new ProducerRecord<>(SHARED_TOPIC, null))
+ producer.send(new ProducerRecord<>(SHARED_TOPIC, null)).get(5, TimeUnit.SECONDS)
then:
- // check that the message was received
+ awaitUntilConsumerIsReady()
def records = consumer.poll(Duration.ofSeconds(5).toMillis())
+ records.count() == 1
+
+ // iterate over records to generate spans
for (record in records) {
assert record.value() == null
assert record.key() == null
@@ -175,14 +182,18 @@ class KafkaClientPropagationEnabledTest extends KafkaClientBaseTest {
when: "send message"
def greeting = "Hello from MockConsumer!"
- producer.send(new ProducerRecord<>(SHARED_TOPIC, partition, null, greeting))
+ producer.send(new ProducerRecord<>(SHARED_TOPIC, partition, null, greeting)).get(5, TimeUnit.SECONDS)
then: "wait for PRODUCER span"
waitForTraces(1)
when: "receive messages"
+ awaitUntilConsumerIsReady()
def consumerRecords = consumer.poll(Duration.ofSeconds(5).toMillis())
def recordsInPartition = consumerRecords.records(new TopicPartition(SHARED_TOPIC, partition))
+ recordsInPartition.size() == 1
+
+ // iterate over records to generate spans
for (record in recordsInPartition) {
assert record.value() == greeting
assert record.key() == null

MODIFY
instrumentation/kafka-streams-0.11/javaagent/src/test/groovy/KafkaStreamsTest.groovy
instrumentation/kafka-streams-0.11/javaagent/src/test/groovy/KafkaStreamsTest.groovy
@@ -13,6 +13,7 @@ import io.opentelemetry.semconv.trace.attributes.SemanticAttributes
import org.apache.kafka.clients.admin.AdminClient
import org.apache.kafka.clients.admin.NewTopic
import org.apache.kafka.clients.consumer.Consumer
+import org.apache.kafka.clients.consumer.ConsumerRebalanceListener
import org.apache.kafka.clients.consumer.KafkaConsumer
import org.apache.kafka.clients.producer.KafkaProducer
import org.apache.kafka.clients.producer.Producer
@@ -32,6 +33,7 @@ import org.testcontainers.containers.KafkaContainer
import spock.lang.Shared
import java.time.Duration
+import java.util.concurrent.CountDownLatch
import java.util.concurrent.TimeUnit
import static io.opentelemetry.api.trace.SpanKind.CONSUMER
@@ -48,7 +50,8 @@ class KafkaStreamsTest extends AgentInstrumentationSpecification {
static Producer<Integer, String> producer
@Shared
static Consumer<Integer, String> consumer
-
+ @Shared
+ static CountDownLatch consumerReady = new CountDownLatch(1)
def setupSpec() {
kafka = new KafkaContainer()
@@ -64,22 +67,27 @@ class KafkaStreamsTest extends AgentInstrumentationSpecification {
producer = new KafkaProducer<>(producerProps(kafka.bootstrapServers))
- // values copied from spring's KafkaTestUtils
def consumerProps = [
"bootstrap.servers" : kafka.bootstrapServers,
"group.id" : "test",
- "enable.auto.commit" : "false",
+ "enable.auto.commit" : "true",
"auto.commit.interval.ms": "10",
- "session.timeout.ms" : "60000",
+ "session.timeout.ms" : "30000",
"key.deserializer" : IntegerDeserializer,
"value.deserializer" : StringDeserializer
]
consumer = new KafkaConsumer<>(consumerProps)
- // assign topic partitions
- consumer.assign([
- new TopicPartition(STREAM_PROCESSED, 0)
- ])
+ consumer.subscribe([STREAM_PROCESSED], new ConsumerRebalanceListener() {
+ @Override
+ void onPartitionsRevoked(Collection<TopicPartition> collection) {
+ }
+
+ @Override
+ void onPartitionsAssigned(Collection<TopicPartition> collection) {
+ consumerReady.countDown()
+ }
+ })
}
def cleanupSpec() {
@@ -132,7 +140,7 @@ class KafkaStreamsTest extends AgentInstrumentationSpecification {
producer.send(new ProducerRecord<>(STREAM_PENDING, greeting))
then:
- // check that the message was received
+ awaitUntilConsumerIsReady()
def records = consumer.poll(Duration.ofSeconds(10).toMillis())
Headers receivedHeaders = null
for (record in records) {
@@ -283,4 +291,21 @@ class KafkaStreamsTest extends AgentInstrumentationSpecification {
"value.serializer" : StringSerializer
]
}
+
+ // Kafka's eventual consistency behavior forces us to do a couple of empty poll() calls until it gets properly assigned a topic partition
+ static void awaitUntilConsumerIsReady() {
+ if (consumerReady.await(0, TimeUnit.SECONDS)) {
+ return
+ }
+ for (i in 0..<10) {
+ consumer.poll(0)
+ if (consumerReady.await(1, TimeUnit.SECONDS)) {
+ break
+ }
+ }
+ if (consumerReady.getCount() != 0) {
+ throw new AssertionError("Consumer wasn't assigned any partitions!")
+ }
+ consumer.seekToBeginning([])
+ }
}

Fix kafka tests failing when consumer is not yet subscribed to topic (#4189)