Skip to content

Commit

Permalink
Fix: register mulitple extensions.
Browse files Browse the repository at this point in the history 
Signed-off-by: dblock <dblock@amazon.com>
  • Loading branch information
@dblock
dblock committed Sep 29, 2023
1 parent e617da3 commit 35c6b61
Show file tree
Hide file tree
Showing 5 changed files with 144 additions and 4 deletions.
3 changes: 3 additions & 0 deletions server/src/main/java/org/opensearch/extensions/ExtensionsManager.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -340,7 +340,9 @@ private void validateExtension(Extension extension) throws IOException {
*/
public void initialize() {
for (DiscoveryExtensionNode extension : extensionIdMap.values()) {
// if (! initializedExtensions.containsKey(extension)) {
initializeExtension(extension);
// }
}
}

Expand Down Expand Up @@ -384,6 +386,7 @@ public String executor() {
transportService.getThreadPool().generic().execute(new AbstractRunnable() {
@Override
public void onFailure(Exception e) {
logger.warn(String.format("Error registering extension: %s", extension.getId()), e);
extensionIdMap.remove(extension.getId());
if (e.getCause() instanceof ConnectTransportException) {
logger.info("No response from extension to request.", e);
Expand Down
3 changes: 3 additions & 0 deletions server/src/main/java/org/opensearch/extensions/rest/RestActionsRequestHandler.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,9 @@ public TransportResponse handleRegisterRestActionsRequest(
DynamicActionRegistry dynamicActionRegistry
) throws Exception {
DiscoveryExtensionNode discoveryExtensionNode = extensionIdMap.get(restActionsRequest.getUniqueId());
if (discoveryExtensionNode == null) {
throw new IllegalStateException(String.format("Missing extension node for %s", restActionsRequest.getUniqueId()));
}
RestHandler handler = new RestSendToExtensionAction(
restActionsRequest,
discoveryExtensionNode,
Expand Down
2 changes: 1 addition & 1 deletion server/src/main/java/org/opensearch/extensions/rest/RestSendToExtensionAction.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -150,7 +150,7 @@ public RestSendToExtensionAction(

@Override
public String getName() {
return SEND_TO_EXTENSION_ACTION;
return this.discoveryExtensionNode.getId() + ":" + SEND_TO_EXTENSION_ACTION;
}

@Override
Expand Down
134 changes: 134 additions & 0 deletions server/src/test/java/org/opensearch/extensions/ExtensionsManagerTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,7 @@
import java.util.HashMap;
import java.util.List;
import java.util.Set;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

Expand Down Expand Up @@ -416,12 +417,81 @@ public void testInitialize() throws Exception {
}
}

public void testInitializeExtensionTwice() throws Exception {
ExtensionsManager extensionsManager = new ExtensionsManager(Set.of(), identityService);
initialize(extensionsManager);

ThreadPool mockThreadPool = spy(threadPool);
ExecutorService mockExecutorService = mock(ExecutorService.class);
when(mockThreadPool.generic()).thenReturn(mockExecutorService);

TransportService transportService = new TransportService(
Settings.EMPTY,
mock(Transport.class),
mockThreadPool,
TransportService.NOOP_TRANSPORT_INTERCEPTOR,
x -> null,
null,
Collections.emptySet(),
NoopTracer.INSTANCE
);

extensionsManager.initializeServicesAndRestHandler(
actionModule,
settingsModule,
transportService,
clusterService,
settings,
client,
identityService
);

Extension firstExtension = new Extension(
"firstExtension",
"uniqueid1",
"127.0.0.0",
"9301",
"0.0.7",
"2.0.0",
"2.0.0",
List.of(),
null
);

extensionsManager.loadExtension(firstExtension);
extensionsManager.initialize();

Extension secondExtension = new Extension(
"secondExtension",
"uniqueid2",
"127.0.0.0",
"9301",
"0.0.7",
"2.0.0",
"2.0.0",
List.of(),
null
);

extensionsManager.loadExtension(secondExtension);
extensionsManager.initialize();

// When execution is mocked, the successful registration callback is not called and the extension is never added to
// registered extensions.
// verify(mockExecutorService, times(2)).execute(any());
}

public void testHandleRegisterRestActionsRequest() throws Exception {

ExtensionsManager extensionsManager = new ExtensionsManager(Set.of(), identityService);
initialize(extensionsManager);

String uniqueIdStr = "uniqueid1";

extensionsManager.loadExtension(
new Extension("firstExtension", uniqueIdStr, "127.0.0.0", "9300", "0.0.7", "3.0.0", "3.0.0", List.of(), null)
);

List<String> actionsList = List.of("GET /foo foo", "PUT /bar bar", "POST /baz baz");
List<String> deprecatedActionsList = List.of("GET /deprecated/foo foo_deprecated", "It's deprecated!");
RegisterRestActionsRequest registerActionsRequest = new RegisterRestActionsRequest(uniqueIdStr, actionsList, deprecatedActionsList);
Expand All @@ -431,6 +501,58 @@ public void testHandleRegisterRestActionsRequest() throws Exception {
assertTrue(((AcknowledgedResponse) response).getStatus());
}

public void testHandleRegisterRestActionsRequestRequiresDiscoveryNode() throws Exception {

ExtensionsManager extensionsManager = new ExtensionsManager(Set.of(), identityService);
initialize(extensionsManager);

RegisterRestActionsRequest registerActionsRequest = new RegisterRestActionsRequest("uniqueId1", List.of(), List.of());

expectThrows(
IllegalStateException.class,
() -> extensionsManager.getRestActionsRequestHandler()
.handleRegisterRestActionsRequest(registerActionsRequest, actionModule.getDynamicActionRegistry())
);
}

public void testHandleRegisterTwoRestActionsRequest() throws Exception {

ExtensionsManager extensionsManager = new ExtensionsManager(Set.of(), identityService);
initialize(extensionsManager);

List<String> actionsList = List.of("GET /foo foo", "PUT /bar bar", "POST /baz baz");
List<String> deprecatedActionsList = List.of("GET /deprecated/foo foo_deprecated", "It's deprecated!");
for (int i = 0; i < 2; i++) {
String uniqueIdStr = String.format("uniqueid-%d", i);

Set<Setting<?>> additionalSettings = extAwarePlugin.getExtensionSettings().stream().collect(Collectors.toSet());
ExtensionScopedSettings extensionScopedSettings = new ExtensionScopedSettings(additionalSettings);
Extension firstExtension = new Extension(
String.format("Extension %s", i),
uniqueIdStr,
"127.0.0.0",
"9300",
"0.0.7",
"3.0.0",
"3.0.0",
List.of(),
extensionScopedSettings
);

extensionsManager.loadExtension(firstExtension);

RegisterRestActionsRequest registerActionsRequest = new RegisterRestActionsRequest(
uniqueIdStr,
actionsList,
deprecatedActionsList
);
TransportResponse response = extensionsManager.getRestActionsRequestHandler()
.handleRegisterRestActionsRequest(registerActionsRequest, actionModule.getDynamicActionRegistry());
assertEquals(AcknowledgedResponse.class, response.getClass());
assertTrue(((AcknowledgedResponse) response).getStatus());
}
}

public void testHandleRegisterSettingsRequest() throws Exception {
ExtensionsManager extensionsManager = new ExtensionsManager(Set.of(), identityService);
initialize(extensionsManager);
Expand All @@ -452,6 +574,9 @@ public void testHandleRegisterRestActionsRequestWithInvalidMethod() throws Excep
initialize(extensionsManager);

String uniqueIdStr = "uniqueid1";
extensionsManager.loadExtension(
new Extension("firstExtension", uniqueIdStr, "127.0.0.0", "9300", "0.0.7", "3.0.0", "3.0.0", List.of(), null)
);
List<String> actionsList = List.of("FOO /foo", "PUT /bar", "POST /baz");
List<String> deprecatedActionsList = List.of("GET /deprecated/foo", "It's deprecated!");
RegisterRestActionsRequest registerActionsRequest = new RegisterRestActionsRequest(uniqueIdStr, actionsList, deprecatedActionsList);
Expand All @@ -467,6 +592,9 @@ public void testHandleRegisterRestActionsRequestWithInvalidDeprecatedMethod() th
initialize(extensionsManager);

String uniqueIdStr = "uniqueid1";
extensionsManager.loadExtension(
new Extension("firstExtension", uniqueIdStr, "127.0.0.0", "9300", "0.0.7", "3.0.0", "3.0.0", List.of(), null)
);
List<String> actionsList = List.of("GET /foo", "PUT /bar", "POST /baz");
List<String> deprecatedActionsList = List.of("FOO /deprecated/foo", "It's deprecated!");
RegisterRestActionsRequest registerActionsRequest = new RegisterRestActionsRequest(uniqueIdStr, actionsList, deprecatedActionsList);
Expand All @@ -481,6 +609,9 @@ public void testHandleRegisterRestActionsRequestWithInvalidUri() throws Exceptio
ExtensionsManager extensionsManager = new ExtensionsManager(Set.of(), identityService);
initialize(extensionsManager);
String uniqueIdStr = "uniqueid1";
extensionsManager.loadExtension(
new Extension("firstExtension", uniqueIdStr, "127.0.0.0", "9300", "0.0.7", "3.0.0", "3.0.0", List.of(), null)
);
List<String> actionsList = List.of("GET", "PUT /bar", "POST /baz");
List<String> deprecatedActionsList = List.of("GET /deprecated/foo", "It's deprecated!");
RegisterRestActionsRequest registerActionsRequest = new RegisterRestActionsRequest(uniqueIdStr, actionsList, deprecatedActionsList);
Expand All @@ -495,6 +626,9 @@ public void testHandleRegisterRestActionsRequestWithInvalidDeprecatedUri() throw
ExtensionsManager extensionsManager = new ExtensionsManager(Set.of(), identityService);
initialize(extensionsManager);
String uniqueIdStr = "uniqueid1";
extensionsManager.loadExtension(
new Extension("firstExtension", uniqueIdStr, "127.0.0.0", "9300", "0.0.7", "3.0.0", "3.0.0", List.of(), null)
);
List<String> actionsList = List.of("GET /foo", "PUT /bar", "POST /baz");
List<String> deprecatedActionsList = List.of("GET", "It's deprecated!");
RegisterRestActionsRequest registerActionsRequest = new RegisterRestActionsRequest(uniqueIdStr, actionsList, deprecatedActionsList);
Expand Down
6 changes: 3 additions & 3 deletions server/src/test/java/org/opensearch/extensions/rest/RestSendToExtensionActionTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -150,7 +150,7 @@ public void testRestSendToExtensionAction() throws Exception {
identityService
);

assertEquals("send_to_extension_action", restSendToExtensionAction.getName());
assertEquals("uniqueid1:send_to_extension_action", restSendToExtensionAction.getName());
List<Route> expected = new ArrayList<>();
String uriPrefix = "/_extensions/_uniqueid1";
expected.add(new Route(Method.GET, uriPrefix + "/foo"));
Expand Down Expand Up @@ -183,7 +183,7 @@ public void testRestSendToExtensionActionWithNamedRoute() throws Exception {
identityService
);

assertEquals("send_to_extension_action", restSendToExtensionAction.getName());
assertEquals("uniqueid1:send_to_extension_action", restSendToExtensionAction.getName());
List<NamedRoute> expected = new ArrayList<>();
String uriPrefix = "/_extensions/_uniqueid1";
NamedRoute nr1 = new NamedRoute.Builder().method(Method.GET).path(uriPrefix + "/foo").uniqueName("foo").build();
Expand Down Expand Up @@ -229,7 +229,7 @@ public void testRestSendToExtensionActionWithNamedRouteAndLegacyActionName() thr
identityService
);

assertEquals("send_to_extension_action", restSendToExtensionAction.getName());
assertEquals("uniqueid1:send_to_extension_action", restSendToExtensionAction.getName());
List<NamedRoute> expected = new ArrayList<>();
String uriPrefix = "/_extensions/_uniqueid1";
NamedRoute nr1 = new NamedRoute.Builder().method(Method.GET)
Expand Down

0 comments on commit 35c6b61

Please sign in to comment.