Skip to content

Commit eb3ae1d

Browse files
akolarkunnudhrubo-oszane-neo
authored
[FEATURE] Improve EncryptorImpl with Asynchronous Handling for Scalability (#3919)
* [FEATURE] Improve EncryptorImpl with Asynchronous Handling for Scalability Removed the usage of ContDownLatch. Every requets will be submitted and returns the Future. Added a list to track the ongoing master key generation. If any tenant id is in the list, then it's key generation is on going and it will wait until other thread completes the key genearion. Same time system will accept other requests, if key is already avaialble in the map that will procced otherwise key generation for new tenant will start in different thread. So, multiple tenants key generation can happen simulatneuosly. Resolves #3510 Signed-off-by: Abdul Muneer Kolarkunnu <muneer.kolarkunnu@netapp.com> * Fixed review comments Resolves #3510 Signed-off-by: Abdul Muneer Kolarkunnu <muneer.kolarkunnu@netapp.com> * [FEATURE] Improve EncryptorImpl with Asynchronous Handling for Scalability Removed the usage of ContDownLatch. Every requets will be submitted and returns the Future. Added a list to track the ongoing master key generation. If any tenant id is in the list, then it's key generation is on going and it will wait until other thread completes the key genearion. Same time system will accept other requests, if key is already avaialble in the map that will procced otherwise key generation for new tenant will start in different thread. So, multiple tenants key generation can happen simulatneuosly. Resolves #3510 Signed-off-by: Abdul Muneer Kolarkunnu <muneer.kolarkunnu@netapp.com> * Removed wait and notify Resolves #3510 Signed-off-by: Abdul Muneer Kolarkunnu <muneer.kolarkunnu@netapp.com> * [FEATURE] Support MCP for Flow and Conversational Flow Agent Resolves #3807 Signed-off-by: Abdul Muneer Kolarkunnu <muneer.kolarkunnu@netapp.com> * [FEATURE] Support MCP for Flow and Conversational Flow Agent Resolves #3807 Signed-off-by: Abdul Muneer Kolarkunnu <muneer.kolarkunnu@netapp.com> * Improved the test coverage and handled the edge case Resolves #3510 Signed-off-by: Abdul Muneer Kolarkunnu <muneer.kolarkunnu@netapp.com> * Fixing coderabbitai reported issues Resolves #3510 Signed-off-by: Abdul Muneer Kolarkunnu <muneer.kolarkunnu@netapp.com> * Fixing flaky tests Resolves #3510 Signed-off-by: Abdul Muneer Kolarkunnu <muneer.kolarkunnu@netapp.com> * Fixing test issues and merging conflicts * fixing coderabbitai comments Resolves #3510 Signed-off-by: Abdul Muneer Kolarkunnu <muneer.kolarkunnu@netapp.com> * fixing coderabbitai comments Resolves #3510 Signed-off-by: Abdul Muneer Kolarkunnu <muneer.kolarkunnu@netapp.com> * Merged latest changes Resolves #3510 Signed-off-by: Abdul Muneer Kolarkunnu <muneer.kolarkunnu@netapp.com> * Fixed spotless Resolves #3510 Signed-off-by: Abdul Muneer Kolarkunnu <muneer.kolarkunnu@netapp.com> * Syned with PR #4678 Resolves #3510 Signed-off-by: Abdul Muneer Kolarkunnu <muneer.kolarkunnu@netapp.com> * Improvements in tests Resolves #3510 Signed-off-by: Abdul Muneer Kolarkunnu <muneer.kolarkunnu@netapp.com> * Added 1 more test with multiple texts in encrypt and decrypt Resolves #3510 Signed-off-by: Abdul Muneer Kolarkunnu <muneer.kolarkunnu@netapp.com> * Fixed PR Code Analyzer review comments Resolves #3510 Signed-off-by: Abdul Muneer Kolarkunnu <muneer.kolarkunnu@netapp.com> * Fixed Spotless Resolves #3510 Signed-off-by: Abdul Muneer Kolarkunnu <muneer.kolarkunnu@netapp.com> * Removing extra onResponse() invocation Resolves #3510 Signed-off-by: Abdul Muneer Kolarkunnu <muneer.kolarkunnu@netapp.com> * Removing extra onResponse() invocation Resolves #3510 Signed-off-by: Abdul Muneer Kolarkunnu <muneer.kolarkunnu@netapp.com> * Fixed review comments Resolves #3510 Signed-off-by: Abdul Muneer Kolarkunnu <muneer.kolarkunnu@netapp.com> * Reverting null tenantid related code Resolves #3510 Signed-off-by: Abdul Muneer Kolarkunnu <muneer.kolarkunnu@netapp.com> * Reverting null tenantid related code Resolves #3510 Signed-off-by: Abdul Muneer Kolarkunnu <muneer.kolarkunnu@netapp.com> --------- Signed-off-by: Abdul Muneer Kolarkunnu <muneer.kolarkunnu@netapp.com> Signed-off-by: Muneer Kolarkunnu <33829651+akolarkunnu@users.noreply.github.com> Co-authored-by: Dhrubo Saha <dhrubo@amazon.com> Co-authored-by: zane-neo <zaniu@amazon.com>
1 parent 72b888d commit eb3ae1d

37 files changed

Lines changed: 1508 additions & 970 deletions

‎common/src/main/java/org/opensearch/ml/common/connector/AbstractConnector.java‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,14 +10,17 @@
1010

1111
import java.io.IOException;
1212
import java.time.Instant;
13+
import java.util.ArrayList;
1314
import java.util.HashMap;
1415
import java.util.List;
1516
import java.util.Map;
1617
import java.util.Optional;
1718

1819
import org.apache.commons.text.StringSubstitutor;
20+
import org.opensearch.common.TriConsumer;
1921
import org.opensearch.common.xcontent.XContentType;
2022
import org.opensearch.commons.authuser.User;
23+
import org.opensearch.core.action.ActionListener;
2124
import org.opensearch.core.xcontent.NamedXContentRegistry;
2225
import org.opensearch.core.xcontent.XContentParser;
2326
import org.opensearch.ml.common.AccessMode;
@@ -26,7 +29,9 @@
2629

2730
import lombok.Getter;
2831
import lombok.Setter;
32+
import lombok.extern.log4j.Log4j2;
2933

34+
@Log4j2
3035
@Getter
3136
public abstract class AbstractConnector implements Connector {
3237
public static final String ACCESS_KEY_FIELD = "access_key";
@@ -157,4 +162,68 @@ public String getActionEndpoint(String action, Map<String, String> parameters) {
157162
return predictEndpoint;
158163
}
159164

165+
@Override
166+
public void encrypt(
167+
TriConsumer<List<String>, String, ActionListener<List<String>>> function,
168+
String tenantId,
169+
ActionListener<Boolean> listener
170+
) {
171+
if (credential == null || credential.isEmpty()) {
172+
listener.onResponse(true);
173+
return;
174+
}
175+
List<String> orderedEncryptKeys = new ArrayList<>();
176+
List<String> orderedToEncrypt = new ArrayList<>();
177+
for (String key : credential.keySet()) {
178+
orderedEncryptKeys.add(key);
179+
orderedToEncrypt.add(credential.get(key));
180+
}
181+
ActionListener<List<String>> updateEncryptedCredentialsListener = ActionListener.wrap(r -> {
182+
Map<String, String> encryptedCredentials = new HashMap<>();
183+
for (int i = 0; i < r.size(); i++) {
184+
encryptedCredentials.put(orderedEncryptKeys.get(i), r.get(i));
185+
}
186+
credential = encryptedCredentials;
187+
listener.onResponse(true);
188+
}, e -> {
189+
log.error("Failed to encrypt credentials in connector", e);
190+
listener.onFailure(e);
191+
});
192+
function.apply(orderedToEncrypt, tenantId, updateEncryptedCredentialsListener);
193+
}
194+
195+
@Override
196+
public void decrypt(
197+
String action,
198+
TriConsumer<List<String>, String, ActionListener<List<String>>> function,
199+
String tenantId,
200+
ActionListener<Boolean> listener
201+
) {
202+
if (credential == null || credential.isEmpty()) {
203+
this.decryptedHeaders = createDecryptedHeaders(getAllHeaders(action));
204+
listener.onResponse(true);
205+
return;
206+
}
207+
List<String> orderedDecryptKeys = new ArrayList<>();
208+
List<String> orderedToDecrypt = new ArrayList<>();
209+
for (Map.Entry<String, String> entry : credential.entrySet()) {
210+
orderedDecryptKeys.add(entry.getKey());
211+
orderedToDecrypt.add(entry.getValue());
212+
}
213+
ActionListener<List<String>> updateDecryptedCredentialsListener = ActionListener.wrap(r -> {
214+
decryptedCredential = new HashMap<>();
215+
for (int i = 0; i < r.size(); i++) {
216+
decryptedCredential.put(orderedDecryptKeys.get(i), r.get(i));
217+
}
218+
this.decryptedHeaders = createDecryptedHeaders(getAllHeaders(action));
219+
listener.onResponse(true);
220+
}, e -> {
221+
log.error("Failed to decrypt credentials in connector", e);
222+
listener.onFailure(e);
223+
});
224+
function.apply(orderedToDecrypt, tenantId, updateDecryptedCredentialsListener);
225+
}
226+
227+
protected abstract Map<String, String> getAllHeaders(String action);
228+
160229
}

‎common/src/main/java/org/opensearch/ml/common/connector/Connector.java‎

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,14 +16,15 @@
1616
import java.util.List;
1717
import java.util.Map;
1818
import java.util.Optional;
19-
import java.util.function.BiFunction;
2019
import java.util.regex.Matcher;
2120
import java.util.regex.Pattern;
2221

2322
import org.apache.commons.text.StringSubstitutor;
23+
import org.opensearch.common.TriConsumer;
2424
import org.opensearch.common.xcontent.LoggingDeprecationHandler;
2525
import org.opensearch.common.xcontent.XContentType;
2626
import org.opensearch.commons.authuser.User;
27+
import org.opensearch.core.action.ActionListener;
2728
import org.opensearch.core.common.io.stream.StreamInput;
2829
import org.opensearch.core.common.io.stream.StreamOutput;
2930
import org.opensearch.core.common.io.stream.Writeable;
@@ -79,9 +80,18 @@ public interface Connector extends ToXContentObject, Writeable {
7980

8081
<T> T createPayload(String action, Map<String, String> parameters);
8182

82-
void decrypt(String action, BiFunction<String, String, String> function, String tenantId);
83+
void decrypt(
84+
String action,
85+
TriConsumer<List<String>, String, ActionListener<List<String>>> function,
86+
String tenantId,
87+
ActionListener<Boolean> listener
88+
);
8389

84-
void encrypt(BiFunction<String, String, String> function, String tenantId);
90+
void encrypt(
91+
TriConsumer<List<String>, String, ActionListener<List<String>>> function,
92+
String tenantId,
93+
ActionListener<Boolean> listener
94+
);
8595

8696
Connector cloneConnector();
8797

@@ -91,7 +101,7 @@ public interface Connector extends ToXContentObject, Writeable {
91101

92102
void writeTo(StreamOutput out) throws IOException;
93103

94-
void update(MLCreateConnectorInput updateContent, BiFunction<String, String, String> function);
104+
void update(MLCreateConnectorInput updateContent);
95105

96106
<T> void parseResponse(T orElse, List<ModelTensor> modelTensors, boolean b) throws IOException;
97107

‎common/src/main/java/org/opensearch/ml/common/connector/HttpConnector.java‎

Lines changed: 6 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,6 @@
2424
import java.util.Locale;
2525
import java.util.Map;
2626
import java.util.Optional;
27-
import java.util.function.BiFunction;
2827
import java.util.regex.Matcher;
2928
import java.util.regex.Pattern;
3029

@@ -315,7 +314,7 @@ public void writeTo(StreamOutput out) throws IOException {
315314
}
316315

317316
@Override
318-
public void update(MLCreateConnectorInput updateContent, BiFunction<String, String, String> function) {
317+
public void update(MLCreateConnectorInput updateContent) {
319318
if (updateContent.getName() != null) {
320319
this.name = updateContent.getName();
321320
}
@@ -333,7 +332,6 @@ public void update(MLCreateConnectorInput updateContent, BiFunction<String, Stri
333332
}
334333
if (updateContent.getCredential() != null && !updateContent.getCredential().isEmpty()) {
335334
this.credential = updateContent.getCredential();
336-
encrypt(function, this.tenantId);
337335
}
338336
if (updateContent.getActions() != null) {
339337
this.actions = updateContent.getActions();
@@ -422,18 +420,6 @@ private List<String> findStringParametersWithNullDefaultValue(String input) {
422420
return paramList;
423421
}
424422

425-
@Override
426-
public void decrypt(String action, BiFunction<String, String, String> function, String tenantId) {
427-
Map<String, String> decrypted = new HashMap<>();
428-
for (String key : credential.keySet()) {
429-
decrypted.put(key, function.apply(credential.get(key), tenantId));
430-
}
431-
this.decryptedCredential = decrypted;
432-
Optional<ConnectorAction> connectorAction = findAction(action);
433-
Map<String, String> headers = connectorAction.map(ConnectorAction::getHeaders).orElse(null);
434-
this.decryptedHeaders = createDecryptedHeaders(headers);
435-
}
436-
437423
@Override
438424
public Connector cloneConnector() {
439425
try (BytesStreamOutput bytesStreamOutput = new BytesStreamOutput()) {
@@ -445,17 +431,14 @@ public Connector cloneConnector() {
445431
}
446432
}
447433

448-
@Override
449-
public void encrypt(BiFunction<String, String, String> function, String tenantId) {
450-
for (String key : credential.keySet()) {
451-
String encrypted = function.apply(credential.get(key), tenantId);
452-
credential.put(key, encrypted);
453-
}
454-
}
455-
456434
@Override
457435
public String getActionHttpMethod(String action) {
458436
return findAction(action).get().getMethod();
459437
}
460438

439+
@Override
440+
protected Map<String, String> getAllHeaders(String action) {
441+
Optional<ConnectorAction> connectorAction = findAction(action);
442+
return connectorAction.map(ConnectorAction::getHeaders).orElse(null);
443+
}
461444
}

‎common/src/main/java/org/opensearch/ml/common/connector/McpConnector.java‎

Lines changed: 4 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -6,22 +6,12 @@
66
package org.opensearch.ml.common.connector;
77

88
import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
9-
import static org.opensearch.ml.common.CommonValue.ACCESS_FIELD;
109
import static org.opensearch.ml.common.CommonValue.BACKEND_ROLES_FIELD;
11-
import static org.opensearch.ml.common.CommonValue.CLIENT_CONFIG_FIELD;
12-
import static org.opensearch.ml.common.CommonValue.CREATED_TIME_FIELD;
1310
import static org.opensearch.ml.common.CommonValue.CREDENTIAL_FIELD;
14-
import static org.opensearch.ml.common.CommonValue.DESCRIPTION_FIELD;
1511
import static org.opensearch.ml.common.CommonValue.HEADERS_FIELD;
16-
import static org.opensearch.ml.common.CommonValue.LAST_UPDATED_TIME_FIELD;
17-
import static org.opensearch.ml.common.CommonValue.NAME_FIELD;
18-
import static org.opensearch.ml.common.CommonValue.OWNER_FIELD;
19-
import static org.opensearch.ml.common.CommonValue.PARAMETERS_FIELD;
20-
import static org.opensearch.ml.common.CommonValue.PROTOCOL_FIELD;
2112
import static org.opensearch.ml.common.CommonValue.TENANT_ID_FIELD;
2213
import static org.opensearch.ml.common.CommonValue.URL_FIELD;
2314
import static org.opensearch.ml.common.CommonValue.VERSION_3_1_0;
24-
import static org.opensearch.ml.common.CommonValue.VERSION_FIELD;
2515
import static org.opensearch.ml.common.connector.ConnectorProtocols.MCP_SSE;
2616
import static org.opensearch.ml.common.connector.ConnectorProtocols.validateProtocol;
2717

@@ -32,7 +22,6 @@
3222
import java.util.List;
3323
import java.util.Map;
3424
import java.util.Optional;
35-
import java.util.function.BiFunction;
3625
import java.util.regex.Matcher;
3726
import java.util.regex.Pattern;
3827

@@ -60,32 +49,8 @@
6049
@EqualsAndHashCode
6150
@Getter
6251
@org.opensearch.ml.common.annotation.Connector(MCP_SSE)
63-
public class McpConnector implements Connector {
52+
public class McpConnector extends AbstractConnector {
6453

65-
protected String name;
66-
protected String description;
67-
protected String version;
68-
protected String protocol;
69-
70-
protected Map<String, String> credential;
71-
protected Map<String, String> decryptedHeaders;
72-
protected Map<String, String> parameters;
73-
@Setter
74-
protected Map<String, String> decryptedCredential;
75-
@Setter
76-
protected List<String> backendRoles;
77-
@Setter
78-
protected User owner;
79-
@Setter
80-
protected AccessMode access;
81-
@Setter
82-
protected Instant createdTime;
83-
@Setter
84-
protected Instant lastUpdateTime;
85-
@Setter
86-
protected ConnectorClientConfig connectorClientConfig;
87-
@Setter
88-
protected String tenantId;
8954
@Setter
9055
@Getter
9156
protected String url;
@@ -204,28 +169,8 @@ protected Map<String, String> createDecryptedHeaders(Map<String, String> headers
204169
return decryptedHeaders;
205170
}
206171

207-
@Override
208-
public void decrypt(String action, BiFunction<String, String, String> function, String tenantId) {
209-
if (credential != null) {
210-
Map<String, String> decrypted = new HashMap<>();
211-
for (String key : credential.keySet()) {
212-
decrypted.put(key, function.apply(credential.get(key), tenantId));
213-
}
214-
this.decryptedCredential = decrypted;
215-
} else {
216-
this.decryptedCredential = new HashMap<>();
217-
}
218-
this.decryptedHeaders = createDecryptedHeaders(headers);
219-
}
220-
221-
@Override
222-
public void encrypt(BiFunction<String, String, String> function, String tenantId) {
223-
if (credential != null) {
224-
for (String key : credential.keySet()) {
225-
String encrypted = function.apply(credential.get(key), tenantId);
226-
credential.put(key, encrypted);
227-
}
228-
}
172+
protected Map<String, String> getAllHeaders(String action) {
173+
return headers;
229174
}
230175

231176
@Override
@@ -338,7 +283,7 @@ public void writeTo(StreamOutput out) throws IOException {
338283
}
339284

340285
@Override
341-
public void update(MLCreateConnectorInput updateContent, BiFunction<String, String, String> function) {
286+
public void update(MLCreateConnectorInput updateContent) {
342287
if (updateContent.getName() != null) {
343288
this.name = updateContent.getName();
344289
}
@@ -353,7 +298,6 @@ public void update(MLCreateConnectorInput updateContent, BiFunction<String, Stri
353298
}
354299
if (updateContent.getCredential() != null && !updateContent.getCredential().isEmpty()) {
355300
this.credential = updateContent.getCredential();
356-
encrypt(function, this.tenantId);
357301
}
358302
if (updateContent.getBackendRoles() != null) {
359303
this.backendRoles = updateContent.getBackendRoles();

0 commit comments

Comments
 (0)