From 8f7da5296c79ac4c5a71fc634590d0a6c19a1c3b Mon Sep 17 00:00:00 2001
From: Saket Chaudhary <saketchaudharynit@gmail.com>
Date: Wed, 29 Oct 2025 22:08:14 +0530
Subject: [PATCH] new request handlers

---
 .../java/com/newrelic/java/HandlerSetup.java  | 52 +++++++++++++++++++
 .../newrelic/java/RequestHandlerWrapper.java  | 18 +++++++
 .../newrelic/java/StreamHandlerWrapper.java   | 22 ++++++++
 3 files changed, 92 insertions(+)
 create mode 100644 java/src/main/java/com/newrelic/java/HandlerSetup.java
 create mode 100644 java/src/main/java/com/newrelic/java/RequestHandlerWrapper.java
 create mode 100644 java/src/main/java/com/newrelic/java/StreamHandlerWrapper.java

diff --git a/java/src/main/java/com/newrelic/java/HandlerSetup.java b/java/src/main/java/com/newrelic/java/HandlerSetup.java
new file mode 100644
index 00000000..e59b40ab
--- /dev/null
+++ b/java/src/main/java/com/newrelic/java/HandlerSetup.java
@@ -0,0 +1,52 @@
+package com.newrelic.java;
+
+import com.amazonaws.services.lambda.runtime.RequestHandler;
+import com.amazonaws.services.lambda.runtime.RequestStreamHandler;
+import com.newrelic.opentracing.LambdaTracer;
+import io.opentracing.Tracer;
+import io.opentracing.util.GlobalTracer;
+
+public class HandlerSetup {
+
+    public static final String HANDLER_ENV_VAR = "NEW_RELIC_LAMBDA_HANDLER";
+
+    public static RequestHandler<Object, Object> requestHandler;
+    public static RequestStreamHandler requestStreamHandler;
+
+    static {
+        // Obtain an instance of the OpenTracing Tracer of your choice
+        Tracer tracer = LambdaTracer.INSTANCE;
+        // Register your tracer as the Global Tracer
+        GlobalTracer.registerIfAbsent(tracer);
+
+        // Set up handlers
+        setupHandlers();
+    }
+
+    static void setupHandlers() {
+        String handler = System.getenv(HANDLER_ENV_VAR);
+        String[] parts = handler.split("::");
+        String handlerClass = parts[0];
+        String handlerMethod = parts.length == 2 ? parts[1] : "handleRequest";
+
+        try {
+            Class<?> loadedClass = JavaClassLoader.class.getClassLoader().loadClass(handlerClass);
+
+            boolean isRequestStreamHandler = RequestStreamHandler.class.isAssignableFrom(loadedClass);
+            requestHandler = isRequestStreamHandler
+                    ? (input, context) -> {
+                        throw new IllegalStateException("" + handlerClass + " is RequestStreamHandler, use handleStreamsRequest instead");
+                    }
+                    : JavaClassLoader.initializeRequestHandler(loadedClass, handlerMethod);
+            requestStreamHandler = isRequestStreamHandler
+                    ? (RequestStreamHandler) loadedClass.getDeclaredConstructor().newInstance()
+                    : (input, output, context) -> {
+                        throw new IllegalStateException("" + handlerClass + " is not RequestStreamHandler, use handleRequest instead");
+                    };
+
+        } catch (ReflectiveOperationException e) {
+            throw new RuntimeException("Error occurred during initialization of javaClassLoader:", e);
+        }
+    }
+
+}
diff --git a/java/src/main/java/com/newrelic/java/RequestHandlerWrapper.java b/java/src/main/java/com/newrelic/java/RequestHandlerWrapper.java
new file mode 100644
index 00000000..73c7624c
--- /dev/null
+++ b/java/src/main/java/com/newrelic/java/RequestHandlerWrapper.java
@@ -0,0 +1,18 @@
+package com.newrelic.java;
+
+import com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestHandler;
+import com.newrelic.opentracing.aws.LambdaTracing;
+
+
+public class RequestHandlerWrapper implements RequestHandler<Object, Object> {
+
+    @Override
+    public Object handleRequest(Object input, Context context) {
+        return LambdaTracing.instrument(
+                input,
+                context,
+                HandlerSetup.requestHandler::handleRequest
+        );
+    }
+}
diff --git a/java/src/main/java/com/newrelic/java/StreamHandlerWrapper.java b/java/src/main/java/com/newrelic/java/StreamHandlerWrapper.java
new file mode 100644
index 00000000..85633f6d
--- /dev/null
+++ b/java/src/main/java/com/newrelic/java/StreamHandlerWrapper.java
@@ -0,0 +1,22 @@
+package com.newrelic.java;
+
+import com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestStreamHandler;
+import com.newrelic.opentracing.aws.StreamLambdaTracing;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+
+public class StreamHandlerWrapper implements RequestStreamHandler {
+
+    @Override
+    public void handleRequest(InputStream input, OutputStream output, Context context) throws IOException {
+        StreamLambdaTracing.instrument(
+                input,
+                output,
+                context,
+                HandlerSetup.requestStreamHandler
+        );
+    }
+}