From 996e6b0262ff90b2c292840c27823a839a9d544b Mon Sep 17 00:00:00 2001 From: Jing Zhang Date: Mon, 28 Oct 2024 14:32:53 -0700 Subject: [PATCH] increase the wait timeout to fetch the master key (#3151) * increase the wait timeout to fetch the master key Signed-off-by: Jing Zhang * fix UTs Signed-off-by: Jing Zhang --------- Signed-off-by: Jing Zhang (cherry picked from commit c7e564c04957feb2978f403c4279c09f006710b9) --- .../org/opensearch/ml/engine/encryptor/EncryptorImpl.java | 5 ++++- .../opensearch/ml/engine/encryptor/EncryptorImplTest.java | 8 ++++---- 2 files changed, 8 insertions(+), 5 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 index 6bf2adfa55..ced2d70fc5 100644 --- 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 @@ -181,7 +181,10 @@ private void initMasterKey() { })); try { - latch.await(1, SECONDS); + boolean completed = latch.await(3, SECONDS); + if (!completed) { + throw new MLException("Fetching master key timed out."); + } } catch (InterruptedException e) { throw new IllegalStateException(e); } diff --git a/ml-algorithms/src/test/java/org/opensearch/ml/engine/encryptor/EncryptorImplTest.java b/ml-algorithms/src/test/java/org/opensearch/ml/engine/encryptor/EncryptorImplTest.java index ba3b3a9167..3f1c8c3948 100644 --- a/ml-algorithms/src/test/java/org/opensearch/ml/engine/encryptor/EncryptorImplTest.java +++ b/ml-algorithms/src/test/java/org/opensearch/ml/engine/encryptor/EncryptorImplTest.java @@ -462,8 +462,8 @@ public void decrypt() { @Test public void encrypt_NullMasterKey_NullMasterKey_MasterKeyNotExistInIndex() { - exceptionRule.expect(ResourceNotFoundException.class); - exceptionRule.expectMessage(MASTER_KEY_NOT_READY_ERROR); + exceptionRule.expect(MLException.class); + exceptionRule.expectMessage("Fetching master key timed out."); doAnswer(invocation -> { ActionListener listener = invocation.getArgument(1); @@ -517,8 +517,8 @@ public void decrypt_NoResponseToInitConfigIndex() { @Test public void decrypt_MLConfigIndexNotFound() { - exceptionRule.expect(ResourceNotFoundException.class); - exceptionRule.expectMessage(MASTER_KEY_NOT_READY_ERROR); + exceptionRule.expect(MLException.class); + exceptionRule.expectMessage("Fetching master key timed out."); Metadata metadata = new Metadata.Builder().indices(ImmutableMap.of()).build(); when(clusterState.metadata()).thenReturn(metadata);