From e3ed9d133d95a374673331b6e9f0e9324b10fd09 Mon Sep 17 00:00:00 2001
From: Lauri Tulmin <ltulmin@splunk.com>
Date: Wed, 17 Jun 2026 13:20:09 +0300
Subject: [PATCH] Run openserch asyn callbacks with parent context

---
 .../OpenSearchTransportInstrumentation.java   | 27 ++++++++++++++++---
 .../v3_0/AbstractOpenSearchTest.java          |  2 +-
 .../v3_0/OpenSearchAwsSdk2TransportTest.java  |  2 +-
 3 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/instrumentation/opensearch/opensearch-java-3.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opensearch/v3_0/OpenSearchTransportInstrumentation.java b/instrumentation/opensearch/opensearch-java-3.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opensearch/v3_0/OpenSearchTransportInstrumentation.java
index a9bd40ead735..3ad024fba05a 100644
--- a/instrumentation/opensearch/opensearch-java-3.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opensearch/v3_0/OpenSearchTransportInstrumentation.java
+++ b/instrumentation/opensearch/opensearch-java-3.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opensearch/v3_0/OpenSearchTransportInstrumentation.java
@@ -58,11 +58,14 @@ public void transform(TypeTransformer transformer) {
 
   public static class AdviceScope {
     private final OpenSearchRequest otelRequest;
+    private final Context parentContext;
     private final Context context;
     private final Scope scope;
 
-    private AdviceScope(OpenSearchRequest otelRequest, Context context, Scope scope) {
+    private AdviceScope(
+        OpenSearchRequest otelRequest, Context parentContext, Context context, Scope scope) {
       this.otelRequest = otelRequest;
+      this.parentContext = parentContext;
       this.context = context;
       this.scope = scope;
     }
@@ -87,11 +90,29 @@ public static AdviceScope start(
         return null;
       }
       Context context = instrumenter().start(parentContext, otelRequest);
-      return new AdviceScope(otelRequest, context, context.makeCurrent());
+      return new AdviceScope(otelRequest, parentContext, context, context.makeCurrent());
     }
 
     public CompletableFuture<Object> wrapFuture(CompletableFuture<Object> future) {
-      return future.whenComplete(new OpenSearchResponseHandler(context, otelRequest));
+      return wrapFuture(
+          future.whenComplete(new OpenSearchResponseHandler(context, otelRequest)), parentContext);
+    }
+
+    private static <T> CompletableFuture<T> wrapFuture(
+        CompletableFuture<T> future, Context context) {
+      CompletableFuture<T> result = new CompletableFuture<>();
+      future.whenComplete(
+          (T value, Throwable throwable) -> {
+            try (Scope ignored = context.makeCurrent()) {
+              if (throwable != null) {
+                result.completeExceptionally(throwable);
+              } else {
+                result.complete(value);
+              }
+            }
+          });
+
+      return result;
     }
 
     public void end(@Nullable Throwable throwable) {
diff --git a/instrumentation/opensearch/opensearch-java-3.0/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/opensearch/v3_0/AbstractOpenSearchTest.java b/instrumentation/opensearch/opensearch-java-3.0/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/opensearch/v3_0/AbstractOpenSearchTest.java
index ed7966a3c58d..ea4fa4d6a7c2 100644
--- a/instrumentation/opensearch/opensearch-java-3.0/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/opensearch/v3_0/AbstractOpenSearchTest.java
+++ b/instrumentation/opensearch/opensearch-java-3.0/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/opensearch/v3_0/AbstractOpenSearchTest.java
@@ -159,7 +159,7 @@ void shouldGetStatusAsyncWithTraces() throws Exception {
                     span ->
                         span.hasName("callback")
                             .hasKind(SpanKind.INTERNAL)
-                            .hasParent(trace.getSpan(1))));
+                            .hasParent(trace.getSpan(0))));
   }
 
   @Test
diff --git a/instrumentation/opensearch/opensearch-java-3.0/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/opensearch/v3_0/OpenSearchAwsSdk2TransportTest.java b/instrumentation/opensearch/opensearch-java-3.0/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/opensearch/v3_0/OpenSearchAwsSdk2TransportTest.java
index 8205dfba7de1..1b68baaea632 100644
--- a/instrumentation/opensearch/opensearch-java-3.0/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/opensearch/v3_0/OpenSearchAwsSdk2TransportTest.java
+++ b/instrumentation/opensearch/opensearch-java-3.0/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/opensearch/v3_0/OpenSearchAwsSdk2TransportTest.java
@@ -203,6 +203,6 @@ void shouldGetStatusAsyncWithTraces() throws Exception {
                     span ->
                         span.hasName("callback")
                             .hasKind(SpanKind.INTERNAL)
-                            .hasParent(trace.getSpan(1))));
+                            .hasParent(trace.getSpan(0))));
   }
 }