Promise response method exit handling logic updated

Author: kcsurapaneni

instrumentation/activej-http-6.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/activejhttp/ActivejHttpServerConnectionInstrumentation.java

@@ -71,7 +71,7 @@ public static void methodEnter(
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void methodExit(
         @Advice.This AsyncServlet asyncServlet,
-        @Advice.Return Promise<HttpResponse> responsePromise,
+        @Advice.Return(readOnly = false) Promise<HttpResponse> responsePromise,
         @Advice.Thrown Throwable throwable,
         @Advice.Local("otelContext") Context context,
         @Advice.Local("otelScope") Scope scope,
@@ -80,12 +80,11 @@ public static void methodExit(
         return;
       }
       scope.close();
-      instrumenter()
-          .end(
-              context,
-              httpRequest,
-              responsePromise == null ? null : responsePromise.getResult(),
-              throwable);
+      if (throwable != null) {
+        instrumenter().end(context, httpRequest, null, throwable);
+      } else {
+        responsePromise = PromiseWrapper.wrap(responsePromise, httpRequest, context);
+      }
     }
   }
 }

instrumentation/activej-http-6.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/activejhttp/PromiseWrapper.java

@@ -0,0 +1,51 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.activejhttp;
+
+import static io.opentelemetry.javaagent.instrumentation.activejhttp.ActivejHttpServerConnectionSingletons.instrumenter;
+
+import io.activej.http.HttpRequest;
+import io.activej.http.HttpResponse;
+import io.activej.promise.Promise;
+import io.activej.promise.SettablePromise;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+
+public class PromiseWrapper {
+
+  public static Promise<HttpResponse> wrap(
+      Promise<HttpResponse> promise, HttpRequest httpRequest, Context context) {
+    SettablePromise<HttpResponse> settablePromise = new SettablePromise<>();
+    promise
+        .whenResult(
+            result -> {
+              Exception error = null;
+              try (Scope ignored = context.makeCurrent()) {
+                settablePromise.set(result);
+              } catch (RuntimeException exception) {
+                error = exception;
+                settablePromise.setException(
+                    new RuntimeException("Context management failed", exception));
+              } finally {
+                instrumenter().end(context, httpRequest, result, error);
+              }
+            })
+        .whenException(
+            throwable -> {
+              try (Scope ignored = context.makeCurrent()) {
+                settablePromise.setException(throwable);
+              } catch (RuntimeException exception) {
+                settablePromise.setException(
+                    new RuntimeException("Context management failed", exception));
+              } finally {
+                instrumenter().end(context, httpRequest, null, throwable);
+              }
+            });
+    return settablePromise;
+  }
+
+  private PromiseWrapper() {}
+}