Skip to content

Commit

Permalink
fix init master key bug (#1094)
Browse files Browse the repository at this point in the history
Signed-off-by: Yaliang Wu <[email protected]>
(cherry picked from commit b782f29)
  • Loading branch information
ylwu-amzn authored and github-actions[bot] committed Jul 12, 2023
1 parent a966bf1 commit 870aa0e
Showing 1 changed file with 12 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.atomic.AtomicReference;

import static java.util.concurrent.TimeUnit.SECONDS;
import static org.opensearch.ml.common.CommonValue.MASTER_KEY;
import static org.opensearch.ml.common.CommonValue.ML_CONFIG_INDEX;

Expand Down Expand Up @@ -111,7 +112,7 @@ private void initMasterKey() {
client.get(getRequest, new LatchedActionListener(ActionListener.<GetResponse>wrap(r -> {
if (r.isExists()) {
String masterKey = (String) r.getSourceAsMap().get(MASTER_KEY);
setMasterKey(masterKey);
this.masterKey = masterKey;
} else {
exceptionRef.set(new ResourceNotFoundException("ML encryption master key not initialized yet"));
}
Expand All @@ -122,6 +123,13 @@ private void initMasterKey() {
}
} else {
exceptionRef.set(new ResourceNotFoundException("ML encryption master key not initialized yet"));
latch.countDown();
}

try {
latch.await(5, SECONDS);
} catch (InterruptedException e) {
throw new IllegalStateException(e);
}

if (exceptionRef.get() != null) {
Expand All @@ -132,5 +140,8 @@ private void initMasterKey() {
throw new MLException(exceptionRef.get());
}
}
if (masterKey == null) {
throw new ResourceNotFoundException("ML encryption master key not initialized yet");
}
}
}

0 comments on commit 870aa0e

Please sign in to comment.