WIP on plugin aware thread context

Signed-off-by: Craig Perkins <[email protected]>

qa/smoke-test-http/src/test/java/org/opensearch/http/ExecutionContextPluginIT.java

@@ -0,0 +1,79 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+/*
+ * Modifications Copyright OpenSearch Contributors. See
+ * GitHub history for details.
+ */
+
+package org.opensearch.http;
+
+import org.opensearch.client.Request;
+import org.opensearch.client.RequestOptions;
+import org.opensearch.client.Response;
+import org.opensearch.client.ResponseException;
+import org.opensearch.plugins.Plugin;
+import org.opensearch.test.OpenSearchIntegTestCase.ClusterScope;
+import org.opensearch.test.OpenSearchIntegTestCase.Scope;
+
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.Collection;
+
+import static org.hamcrest.Matchers.equalTo;
+
+/**
+ * Test a rest action that sets special response headers
+ */
+@ClusterScope(scope = Scope.SUITE, supportsDedicatedMasters = false, numDataNodes = 1)
+public class ExecutionContextPluginIT extends HttpSmokeTestCase {
+
+    @Override
+    protected boolean addMockHttpTransport() {
+        return false; // enable http
+    }
+
+    @Override
+    protected Collection<Class<? extends Plugin>> nodePlugins() {
+        ArrayList<Class<? extends Plugin>> plugins = new ArrayList<>(super.nodePlugins());
+        plugins.add(TestExecutionContextPlugin.class);
+        return plugins;
+    }
+
+    public void testThatSettingHeadersWorks() throws IOException {
+        ensureGreen();
+        try {
+            Response response = getRestClient().performRequest(new Request("GET", "/_execution_context"));
+            System.out.println("Response body: " + new String(response.getEntity().getContent().readAllBytes(), StandardCharsets.UTF_8));
+//            fail("request should have failed");
+        } catch(ResponseException e) {
+            Response response = e.getResponse();
+            assertThat(response.getStatusLine().getStatusCode(), equalTo(401));
+            assertThat(response.getHeader("Secret"), equalTo("required"));
+        }
+    }
+}

qa/smoke-test-http/src/test/java/org/opensearch/http/TestExecutionContextPlugin.java

@@ -0,0 +1,91 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*
+ * Modifications Copyright OpenSearch Contributors. See
+ * GitHub history for details.
+ */
+
+package org.opensearch.http;
+
+import org.opensearch.client.Client;
+import org.opensearch.cluster.metadata.IndexNameExpressionResolver;
+import org.opensearch.cluster.node.DiscoveryNodes;
+import org.opensearch.cluster.service.ClusterService;
+import org.opensearch.common.settings.ClusterSettings;
+import org.opensearch.common.settings.IndexScopedSettings;
+import org.opensearch.common.settings.Settings;
+import org.opensearch.common.settings.SettingsFilter;
+import org.opensearch.core.common.io.stream.NamedWriteableRegistry;
+import org.opensearch.core.xcontent.NamedXContentRegistry;
+import org.opensearch.env.Environment;
+import org.opensearch.env.NodeEnvironment;
+import org.opensearch.plugins.ActionPlugin;
+import org.opensearch.plugins.Plugin;
+import org.opensearch.repositories.RepositoriesService;
+import org.opensearch.rest.RestController;
+import org.opensearch.rest.RestHandler;
+import org.opensearch.script.ScriptService;
+import org.opensearch.threadpool.ThreadPool;
+import org.opensearch.watcher.ResourceWatcherService;
+
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import java.util.function.Supplier;
+
+import static java.util.Collections.singletonList;
+
+public class TestExecutionContextPlugin extends Plugin implements ActionPlugin {
+
+    private ThreadPool threadPool;
+
+    @Override
+    public Collection<Object> createComponents(
+        Client client,
+        ClusterService clusterService,
+        ThreadPool threadPool,
+        ResourceWatcherService resourceWatcherService,
+        ScriptService scriptService,
+        NamedXContentRegistry xContentRegistry,
+        Environment environment,
+        NodeEnvironment nodeEnvironment,
+        NamedWriteableRegistry namedWriteableRegistry,
+        IndexNameExpressionResolver expressionResolver,
+        Supplier<RepositoriesService> repositoriesServiceSupplier
+    ) {
+        this.threadPool = threadPool;
+        return Collections.emptyList();
+    }
+
+    @Override
+    public List<RestHandler> getRestHandlers(Settings settings, RestController restController, ClusterSettings clusterSettings,
+            IndexScopedSettings indexScopedSettings, SettingsFilter settingsFilter, IndexNameExpressionResolver indexNameExpressionResolver,
+            Supplier<DiscoveryNodes> nodesInCluster) {
+        return singletonList(new TestExecutionContextRestAction(threadPool));
+    }
+}

qa/smoke-test-http/src/test/java/org/opensearch/http/TestExecutionContextRestAction.java

@@ -0,0 +1,72 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+/*
+ * Modifications Copyright OpenSearch Contributors. See
+ * GitHub history for details.
+ */
+
+package org.opensearch.http;
+
+import org.opensearch.client.node.NodeClient;
+import org.opensearch.core.rest.RestStatus;
+import org.opensearch.rest.BaseRestHandler;
+import org.opensearch.rest.BytesRestResponse;
+import org.opensearch.rest.RestRequest;
+import org.opensearch.rest.RestResponse;
+import org.opensearch.threadpool.ThreadPool;
+
+import java.util.List;
+
+import static java.util.Collections.singletonList;
+import static org.opensearch.rest.RestRequest.Method.GET;
+
+public class TestExecutionContextRestAction extends BaseRestHandler {
+
+    private final ThreadPool threadPool;
+
+    public TestExecutionContextRestAction(ThreadPool threadPool) {
+        this.threadPool = threadPool;
+    }
+
+    @Override
+    public List<Route> routes() {
+        return singletonList(new Route(GET, "/_execution_context"));
+    }
+
+    @Override
+    public String getName() {
+        return "test_execution_context_action";
+    }
+
+    @Override
+    public RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) {
+        System.out.println("Plugin execution context: " + threadPool.getThreadContext().getExecutionContext());
+        threadPool.getThreadContext().setExecutionContext("should-not-allow-plugin-to-set-execution-context");
+        RestResponse response = new BytesRestResponse(RestStatus.OK, "Should not happen");
+        return channel -> channel.sendResponse(response);
+    }
+}

server/src/main/java/org/opensearch/action/ActionModule.java

@@ -322,6 +322,7 @@
 import org.opensearch.rest.NamedRoute;
 import org.opensearch.rest.RestController;
 import org.opensearch.rest.RestHandler;
+import org.opensearch.rest.RestHandlerProxy;
 import org.opensearch.rest.RestHeaderDefinition;
 import org.opensearch.rest.action.RestFieldCapabilitiesAction;
 import org.opensearch.rest.action.RestMainAction;
@@ -993,7 +994,8 @@ public void initRestHandlers(Supplier<DiscoveryNodes> nodesInCluster) {
                 indexNameExpressionResolver,
                 nodesInCluster
             )) {
-                registerHandler.accept(handler);
+                RestHandler handlerProxy = RestHandlerProxy.newInstance(handler, threadPool, plugin);
+                registerHandler.accept(handlerProxy);
             }
         }
         registerHandler.accept(new RestCatAction(catActions));

server/src/main/java/org/opensearch/rest/RestHandlerProxy.java

@@ -0,0 +1,53 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.rest;
+
+import org.opensearch.plugins.ActionPlugin;
+import org.opensearch.plugins.Plugin;
+import org.opensearch.threadpool.ThreadPool;
+
+import java.lang.reflect.InvocationHandler;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import java.lang.reflect.Proxy;
+
+public class RestHandlerProxy implements InvocationHandler {
+    private final RestHandler restHandler;
+    private final ThreadPool threadPool;
+    private final ActionPlugin plugin;
+
+    public static RestHandler newInstance(RestHandler obj, ThreadPool threadPool, ActionPlugin plugin) {
+        return (RestHandler) Proxy.newProxyInstance(
+            obj.getClass().getClassLoader(),
+            new Class<?>[] { RestHandler.class },
+            new RestHandlerProxy(obj, threadPool, plugin)
+        );
+    }
+
+    private RestHandlerProxy(RestHandler restHandler, ThreadPool threadPool, ActionPlugin plugin) {
+        this.restHandler = restHandler;
+        this.threadPool = threadPool;
+        this.plugin = plugin;
+    }
+
+    @Override
+    public Object invoke(Object proxy, Method m, Object[] args) throws Throwable {
+        Object result;
+        try {
+            threadPool.getThreadContext().setExecutionContext(Proxy.getInvocationHandler(plugin).getClass().getName());
+            result = m.invoke(restHandler, args);
+            threadPool.getThreadContext().clearExecutionContext();
+        } catch (InvocationTargetException e) {
+            throw e.getTargetException();
+        } catch (Exception e) {
+            throw new RuntimeException("unexpected invocation exception: " + e.getMessage());
+        }
+        return result;
+    }
+}