diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/client/AvroGenericDaVinciClient.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/client/AvroGenericDaVinciClient.java index 34b1edbae8..348b48128f 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/client/AvroGenericDaVinciClient.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/client/AvroGenericDaVinciClient.java @@ -842,7 +842,21 @@ public synchronized void start() { @Override public synchronized void close() { - throwIfNotReady(); + if (isReady()) { + closeInner(); + } else { + getClientLogger() + .warn("Client is not ready or already closed, will ignore close request, storeName=" + getStoreName()); + } + } + + @Override + public String toString() { + return this.getClass().getSimpleName(); + } + + // Visible for testing + void closeInner() { try { logger.info("Closing client, storeName=" + getStoreName()); ready.set(false); @@ -850,7 +864,7 @@ public synchronized void close() { cacheBackend.close(); } daVinciBackend.release(); - logger.info("Client is closed successfully, storeName=" + getStoreName()); + logger.info("Client is closed successfully, storeName={}", getStoreName()); } catch (Throwable e) { String msg = "Unable to close Da Vinci client, storeName=" + getStoreName(); logger.error(msg, e); @@ -858,8 +872,8 @@ public synchronized void close() { } } - @Override - public String toString() { - return this.getClass().getSimpleName(); + Logger getClientLogger() { + return logger; } + } diff --git a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/client/AvroGenericDaVinciClientTest.java b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/client/AvroGenericDaVinciClientTest.java index 9712c73d21..cf84fd4a7a 100644 --- a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/client/AvroGenericDaVinciClientTest.java +++ b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/client/AvroGenericDaVinciClientTest.java @@ -4,10 +4,14 @@ import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyInt; import static org.mockito.Mockito.anyString; +import static org.mockito.Mockito.doCallRealMethod; import static org.mockito.Mockito.doNothing; import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertThrows; @@ -42,6 +46,7 @@ import java.util.concurrent.ExecutionException; import java.util.concurrent.Executor; import org.apache.avro.Schema; +import org.apache.logging.log4j.LogManager; import org.testng.Assert; import org.testng.annotations.Test; @@ -230,5 +235,20 @@ public void constructorTest() { readChunkExecutor); assertEquals(daVinciClient.getReadChunkExecutorForLargeRequest(), readChunkExecutor); + // Close a not-ready client won't throw exception. + daVinciClient.close(); + } + + @Test + public void closeTest() { + AvroGenericDaVinciClient client = mock(AvroGenericDaVinciClient.class); + doCallRealMethod().when(client).close(); + doReturn(LogManager.getLogger(AvroGenericDaVinciClient.class)).when(client).getClientLogger(); + doReturn(false).when(client).isReady(); + client.close(); + verify(client, never()).closeInner(); + doReturn(true).when(client).isReady(); + client.close(); + verify(client, times(1)).closeInner(); } } diff --git a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/DaVinciClientTest.java b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/DaVinciClientTest.java index 6d413399ef..5714633ec9 100644 --- a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/DaVinciClientTest.java +++ b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/DaVinciClientTest.java @@ -1106,6 +1106,8 @@ public void testHybridStore() throws Exception { // Verify that closed cached client can be restarted. client.close(); + // Verify that 2nd close call on the same store won't throw exception. + client.close(); DaVinciClient client1 = factory.getAndStartGenericAvroClient(storeName, new DaVinciConfig()); assertEquals((int) client1.get(1).get(), 1);