diff --git a/server/src/main/java/org/opensearch/action/ActionModule.java b/server/src/main/java/org/opensearch/action/ActionModule.java index 1ddb988dc6508..bd4180ca07e70 100644 --- a/server/src/main/java/org/opensearch/action/ActionModule.java +++ b/server/src/main/java/org/opensearch/action/ActionModule.java @@ -500,6 +500,7 @@ import org.opensearch.rest.action.search.RestSearchAction; import org.opensearch.rest.action.search.RestSearchScrollAction; import org.opensearch.tasks.Task; +import org.opensearch.telemetry.tracing.TracerContextStorage; import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.client.node.NodeClient; import org.opensearch.usage.UsageService; @@ -594,9 +595,13 @@ public ActionModule( new RestHeaderDefinition(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, false) ) ).collect(Collectors.toSet()); + Set transients = Stream.concat( + actionPlugins.stream().flatMap(p -> p.getTransients().stream()), + Stream.of(TracerContextStorage.CURRENT_SPAN) + ).collect(Collectors.toSet()); UnaryOperator restWrapper = null; for (ActionPlugin plugin : actionPlugins) { - UnaryOperator newRestWrapper = plugin.getRestHandlerWrapper(threadPool.getThreadContext(), headers); + UnaryOperator newRestWrapper = plugin.getRestHandlerWrapper(threadPool.getThreadContext(), headers, transients); if (newRestWrapper != null) { logger.debug("Using REST wrapper from plugin " + plugin.getClass().getName()); if (restWrapper != null) { diff --git a/server/src/main/java/org/opensearch/plugins/ActionPlugin.java b/server/src/main/java/org/opensearch/plugins/ActionPlugin.java index a9abdb7b4bb77..832c508841b8a 100644 --- a/server/src/main/java/org/opensearch/plugins/ActionPlugin.java +++ b/server/src/main/java/org/opensearch/plugins/ActionPlugin.java @@ -121,6 +121,13 @@ default Collection getRestHeaders() { return Collections.emptyList(); } + /** + * Returns transients which should be copied through context propagation + */ + default Collection getTransients() { + return Collections.emptyList(); + } + /** * Returns headers which should be copied from internal requests into tasks. */ @@ -149,6 +156,15 @@ default Collection getTaskHeaders() { * * Note: Only one installed plugin may implement a rest wrapper. */ + default UnaryOperator getRestHandlerWrapper( + ThreadContext threadContext, + Set headersToCopy, + Set transients + ) { + return this.getRestHandlerWrapper(threadContext); + } + + @Deprecated(forRemoval = true) default UnaryOperator getRestHandlerWrapper(ThreadContext threadContext, Set headersToCopy) { return this.getRestHandlerWrapper(threadContext); } diff --git a/server/src/main/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorage.java b/server/src/main/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorage.java index 908164d1935a7..9f22d5c035a84 100644 --- a/server/src/main/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorage.java +++ b/server/src/main/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorage.java @@ -56,7 +56,7 @@ public Map transients(Map source) { final Map transients = new HashMap<>(); if (source.containsKey(CURRENT_SPAN)) { final SpanReference current = (SpanReference) source.get(CURRENT_SPAN); - if (current != null) { + if (current != null && current.getSpan() != null) { transients.put(CURRENT_SPAN, new SpanReference(current.getSpan())); } } diff --git a/server/src/test/java/org/opensearch/action/ActionModuleTests.java b/server/src/test/java/org/opensearch/action/ActionModuleTests.java index 0c1377cb0c6b2..cde93850501b1 100644 --- a/server/src/test/java/org/opensearch/action/ActionModuleTests.java +++ b/server/src/test/java/org/opensearch/action/ActionModuleTests.java @@ -53,11 +53,13 @@ import org.opensearch.rest.RestChannel; import org.opensearch.rest.RestController; import org.opensearch.rest.RestHandler; +import org.opensearch.rest.RestHeaderDefinition; import org.opensearch.rest.RestRequest; import org.opensearch.rest.RestRequest.Method; import org.opensearch.rest.action.RestMainAction; import org.opensearch.tasks.Task; import org.opensearch.tasks.TaskManager; +import org.opensearch.telemetry.tracing.TracerContextStorage; import org.opensearch.test.OpenSearchTestCase; import org.opensearch.threadpool.TestThreadPool; import org.opensearch.threadpool.ThreadPool; @@ -66,9 +68,11 @@ import java.io.IOException; import java.util.ArrayList; +import java.util.Collection; import java.util.List; import java.util.Set; import java.util.function.Supplier; +import java.util.function.UnaryOperator; import static java.util.Collections.emptyList; import static java.util.Collections.singletonList; @@ -274,4 +278,59 @@ public List routes() { threadPool.shutdown(); } } + + public void testTransientsCollectedFromPluginsAndCore() { + + ThreadPool threadPool = new TestThreadPool("test"); + try { + // Create a plugin that tracks what transients it receives + final Set[] receivedTransients = new Set[1]; + ActionPlugin testPlugin = new ActionPlugin() { + @Override + public Collection getTransients() { + return List.of("custom_transient"); + } + + @Override + public UnaryOperator getRestHandlerWrapper( + ThreadContext threadContext, + Set headersToCopy, + Set transients + ) { + // Capture the transients passed to this method + receivedTransients[0] = transients; + return handler -> handler; + } + }; + List plugins = List.of(testPlugin); + + SettingsModule settings = new SettingsModule(Settings.EMPTY); + UsageService usageService = new UsageService(); + ActionModule actionModule = new ActionModule( + settings.getSettings(), + new IndexNameExpressionResolver(new ThreadContext(Settings.EMPTY)), + settings.getIndexScopedSettings(), + settings.getClusterSettings(), + settings.getSettingsFilter(), + threadPool, + plugins, + null, + null, + usageService, + null, + new IdentityService(Settings.EMPTY, mock(ThreadPool.class), new ArrayList<>()), + new ExtensionsManager(Set.of(), new IdentityService(Settings.EMPTY, mock(ThreadPool.class), List.of())) + ); + + // Verify transients were passed to the plugin + assertNotNull("Plugin should have received transients", receivedTransients[0]); + assertTrue("Should contain custom transient from plugin", receivedTransients[0].contains("custom_transient")); + assertTrue("Should contain core CURRENT_SPAN", receivedTransients[0].contains(TracerContextStorage.CURRENT_SPAN)); + } catch (IOException e) { + throw new RuntimeException(e); + } finally { + threadPool.shutdown(); + } + } + } diff --git a/server/src/test/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorageTests.java b/server/src/test/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorageTests.java index 98dfc367c20f5..bc2975b97ec91 100644 --- a/server/src/test/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorageTests.java +++ b/server/src/test/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorageTests.java @@ -22,6 +22,8 @@ import org.junit.After; import org.junit.Before; +import java.util.HashMap; +import java.util.Map; import java.util.Optional; import java.util.Set; import java.util.concurrent.ExecutionException; @@ -269,4 +271,12 @@ public void testSpanNotPropagatedToChildSystemThreadContext() { assertThat(threadContext.getTransient(ThreadContextBasedTracerContextStorage.CURRENT_SPAN), is(not(nullValue()))); assertThat(threadContextStorage.get(ThreadContextBasedTracerContextStorage.CURRENT_SPAN), is(nullValue())); } + + public void testNullSpanWithinSpanReference() { + SpanReference spanReference = new SpanReference(null); + Map source = new HashMap<>(); + source.put(ThreadContextBasedTracerContextStorage.CURRENT_SPAN, spanReference); + ThreadContextBasedTracerContextStorage context = (ThreadContextBasedTracerContextStorage) threadContextStorage; + assertTrue(context.transients(source).isEmpty()); + } }