fix init master key bug

Signed-off-by: Yaliang Wu <[email protected]>
opensearch-project · Jul 12, 2023 · 3d2e924 · 3d2e924
1 parent d9d1190
commit 3d2e924
Showing 1 changed file with 13 additions and 2 deletions.
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/encryptor/EncryptorImpl.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/encryptor/EncryptorImpl.java
@@ -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;
 
@@ -98,7 +99,7 @@ public String generateMasterKey() {
         return base64Key;
     }
 
-    private void initMasterKey() {
+    private synchronized void initMasterKey() {
         if (masterKey != null) {
             return;
         }
@@ -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"));
                     }
@@ -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) {
@@ -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");
+        }
     }
 }