wip: adding a static SinkRecordHeadersGetter in KafkaConnectSingletons.

Author: vasantteja

instrumentation/kafka/kafka-connect-2.6/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/kafkaconnect/v2_6/KafkaConnectSingletons.java

@@ -6,15 +6,19 @@
 package io.opentelemetry.javaagent.instrumentation.kafkaconnect.v2_6;
 
 import io.opentelemetry.api.GlobalOpenTelemetry;
+import io.opentelemetry.context.propagation.TextMapGetter;
 import io.opentelemetry.context.propagation.TextMapPropagator;
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
+import org.apache.kafka.connect.sink.SinkRecord;
 
 public final class KafkaConnectSingletons {
 
   private static final String INSTRUMENTATION_NAME = "io.opentelemetry.kafka-connect-2.6";
   private static final TextMapPropagator PROPAGATOR = 
       GlobalOpenTelemetry.get().getPropagators().getTextMapPropagator();
 
+  private static final TextMapGetter<SinkRecord> SINK_RECORD_HEADER_GETTER = SinkRecordHeadersGetter.INSTANCE;
+  
   private static final Instrumenter<KafkaConnectTask, Void> INSTRUMENTER =
       Instrumenter.<KafkaConnectTask, Void>builder(
               GlobalOpenTelemetry.get(), INSTRUMENTATION_NAME, KafkaConnectTask::getSpanName)
@@ -29,5 +33,9 @@ public static TextMapPropagator propagator() {
     return PROPAGATOR;
   }
 
+  public static TextMapGetter<SinkRecord> sinkRecordHeaderGetter() {
+    return SINK_RECORD_HEADER_GETTER;
+  }
+
   private KafkaConnectSingletons() {}
 }

instrumentation/kafka/kafka-connect-2.6/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/kafkaconnect/v2_6/SinkTaskInstrumentation.java

@@ -12,7 +12,6 @@
 
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
-import io.opentelemetry.context.propagation.TextMapGetter;
 import io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
@@ -21,9 +20,6 @@
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 import org.apache.kafka.connect.sink.SinkRecord;
-import org.apache.kafka.connect.header.Header;
-import org.apache.kafka.connect.header.Headers;
-import java.nio.charset.StandardCharsets;
 import java.util.logging.Logger;
 
 public class SinkTaskInstrumentation implements TypeInstrumentation {
@@ -81,7 +77,8 @@ public static void onEnter(
       // Extract context from first record if available
       if (!records.isEmpty()) {
         SinkRecord firstRecord = records.iterator().next();
-        Context extractedContext = extractContextFromRecord(parentContext, firstRecord);
+        Context extractedContext = KafkaConnectSingletons.propagator()
+            .extract(parentContext, firstRecord, KafkaConnectSingletons.sinkRecordHeaderGetter());
         parentContext = extractedContext;
       }
 
@@ -112,36 +109,6 @@ public static void onExit(
       logger.info("KafkaConnect: Span ended");
     }
 
-    public static Context extractContextFromRecord(Context parentContext, SinkRecord record) {
-      // Create a simple TextMapGetter for SinkRecord headers
-      return KafkaConnectSingletons.propagator().extract(parentContext, record, new TextMapGetter<SinkRecord>() {
-        @Override
-        public Iterable<String> keys(SinkRecord sinkRecord) {
-          Headers headers = sinkRecord.headers();
-          java.util.List<String> keys = new java.util.ArrayList<>();
-          for (Header header : headers) {
-            keys.add(header.key());
-          }
-          return keys;
-        }
-
-        @Override
-        public String get(SinkRecord sinkRecord, String key) {
-          if (sinkRecord == null) {
-            return null;
-          }
-          Headers headers = sinkRecord.headers();
-          Header header = headers.lastWithName(key);
-          if (header == null) {
-            return null;
-          }
-          Object value = header.value();
-          if (value instanceof byte[]) {
-            return new String((byte[]) value, StandardCharsets.UTF_8);
-          }
-          return value != null ? value.toString() : null;
-        }
-      });
-    }
+
   }
 }

instrumentation/kafka/kafka-connect-2.6/testing/src/test/java/io/opentelemetry/instrumentation/kafkaconnect/v2_6/MongoKafkaConnectSinkTaskTest.java

@@ -135,8 +135,9 @@ public static void setup() throws IOException {
         .withExposedPorts(BACKEND_PORT)
         .withNetwork(network)
         .withNetworkAliases(BACKEND_ALIAS)
-        .waitingFor(Wait.forHttp("/health").forPort(BACKEND_PORT))
-        .withStartupTimeout(Duration.of(2, MINUTES));
+        .waitingFor(Wait.forHttp("/health").forPort(BACKEND_PORT).withStartupTimeout(Duration.of(5, MINUTES)))
+        .withStartupTimeout(Duration.of(5, MINUTES))
+        .withEnv("DOCKER_DEFAULT_PLATFORM", "linux/amd64"); // Force AMD64 for stability on ARM64 hosts
 
     logger.info("Starting backend container...");
     backend.start();
@@ -150,7 +151,7 @@ public static void setup() throws IOException {
             .withEnv("ZOOKEEPER_CLIENT_PORT", String.valueOf(ZOOKEEPER_INTERNAL_PORT))
             .withEnv("ZOOKEEPER_TICK_TIME", "2000")
             .withExposedPorts(ZOOKEEPER_INTERNAL_PORT)
-            .withStartupTimeout(Duration.of(3, MINUTES));
+            .withStartupTimeout(Duration.of(5, MINUTES));
     String zookeeperInternalUrl = ZOOKEEPER_NETWORK_ALIAS + ":" + ZOOKEEPER_INTERNAL_PORT;
 
     kafkaExposedPort = getRandomFreePort();
@@ -186,7 +187,7 @@ public static void setup() throws IOException {
             .withEnv("KAFKA_AUTO_CREATE_TOPICS_ENABLE", "true")
             .withEnv("KAFKA_OPTS", "-Djava.net.preferIPv4Stack=True")
             .withEnv("KAFKA_GROUP_INITIAL_REBALANCE_DELAY_MS", "100")
-            .withStartupTimeout(Duration.of(3, MINUTES));
+            .withStartupTimeout(Duration.of(5, MINUTES));
 
     // Get the agent path from system properties
     String agentPath = System.getProperty("otel.javaagent.testing.javaagent-jar-path");
@@ -237,7 +238,8 @@ public static void setup() throws IOException {
             .withEnv("CONNECT_CONFIG_STORAGE_REPLICATION_FACTOR", "1")
             .withEnv("CONNECT_OFFSET_STORAGE_REPLICATION_FACTOR", "1")
             .withEnv("CONNECT_STATUS_STORAGE_REPLICATION_FACTOR", "1")
-            .withStartupTimeout(Duration.of(3, MINUTES))
+            .waitingFor(Wait.forHttp("/").forPort(CONNECT_REST_PORT_INTERNAL).withStartupTimeout(Duration.of(5, MINUTES)))
+            .withStartupTimeout(Duration.of(5, MINUTES))
             .withCommand(
                 "bash",
                 "-c",
@@ -248,7 +250,7 @@ public static void setup() throws IOException {
         new MongoDBContainer(DockerImageName.parse("mongo:4.4"))
             .withNetwork(network)
             .withNetworkAliases(MONGO_NETWORK_ALIAS)
-            .withStartupTimeout(Duration.of(3, MINUTES));
+            .withStartupTimeout(Duration.of(5, MINUTES));
 
     // Start containers (backend already started)
     Startables.deepStart(Stream.of(zookeeper, kafka, kafkaConnect, mongoDB)).join();
@@ -314,12 +316,12 @@ public void testKafkaConnectMongoSinkTaskInstrumentation() throws IOException, I
       producer.flush();
     }
 
-    // Wait for message processing
-    await().atMost(Duration.ofSeconds(30)).until(() -> getRecordCountFromMongo() >= 1);
+    // Wait for message processing (increased timeout for ARM64 Docker emulation)
+    await().atMost(Duration.ofSeconds(60)).until(() -> getRecordCountFromMongo() >= 1);
 
-    // Wait for spans to arrive at backend
+    // Wait for spans to arrive at backend (increased timeout for ARM64 Docker emulation)
     String backendUrl = getBackendUrl();
-    await().atMost(Duration.ofSeconds(15)).until(() -> {
+    await().atMost(Duration.ofSeconds(30)).until(() -> {
       try {
         String traces = given()
             .when()
@@ -444,7 +446,7 @@ private static void createTopic(String topicName) {
   private static void awaitForTopicCreation(String topicName) {
     try (AdminClient adminClient = createAdminClient()) {
       await()
-          .atMost(Duration.ofSeconds(30))
+          .atMost(Duration.ofSeconds(60))
           .pollInterval(Duration.ofMillis(500))
           .until(() -> adminClient.listTopics().names().get().contains(topicName));
     }
@@ -503,14 +505,78 @@ private static void deleteConnectorIfExists() {
 
   @AfterAll
   public static void cleanup() {
+    // Close AdminClient first to release Kafka connections
+    if (adminClient != null) {
+      try {
+        adminClient.close();
+        logger.info("AdminClient closed");
+      } catch (RuntimeException e) {
+        logger.error("Error closing AdminClient: " + e.getMessage());
+      }
+    }
+    
+    // Stop all containers in reverse order of startup to ensure clean shutdown
+    if (kafkaConnect != null) {
+      try {
+        kafkaConnect.stop();
+        logger.info("Kafka Connect container stopped");
+      } catch (RuntimeException e) {
+        logger.error("Error stopping Kafka Connect: " + e.getMessage());
+      }
+    }
+    
+    if (mongoDB != null) {
+      try {
+        mongoDB.stop();
+        logger.info("MongoDB container stopped");
+      } catch (RuntimeException e) {
+        logger.error("Error stopping MongoDB: " + e.getMessage());
+      }
+    }
+    
+    if (kafka != null) {
+      try {
+        kafka.stop();
+        logger.info("Kafka container stopped");
+      } catch (RuntimeException e) {
+        logger.error("Error stopping Kafka: " + e.getMessage());
+      }
+    }
+    
+    if (zookeeper != null) {
+      try {
+        zookeeper.stop();
+        logger.info("Zookeeper container stopped");
+      } catch (RuntimeException e) {
+        logger.error("Error stopping Zookeeper: " + e.getMessage());
+      }
+    }
+    
     if (backend != null) {
-      backend.stop();
-      logger.info("Backend container stopped");
+      try {
+        backend.stop();
+        logger.info("Backend container stopped");
+      } catch (RuntimeException e) {
+        logger.error("Error stopping backend: " + e.getMessage());
+      }
     }
+    
     if (network != null) {
-      network.close();
-      logger.info("Network closed");
+      try {
+        network.close();
+        logger.info("Network closed");
+      } catch (RuntimeException e) {
+        logger.error("Error closing network: " + e.getMessage());
+      }
     }
+    
+    // Small delay to ensure containers are fully stopped before next test
+    try {
+      Thread.sleep(2000);
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+    }
+    
     logger.info("Test cleanup complete");
   }
 }