diff --git a/sdk-core/src/main/java/io/milvus/pool/ClientPool.java b/sdk-core/src/main/java/io/milvus/pool/ClientPool.java index d8d9b580a..358b5237e 100644 --- a/sdk-core/src/main/java/io/milvus/pool/ClientPool.java +++ b/sdk-core/src/main/java/io/milvus/pool/ClientPool.java @@ -7,6 +7,8 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.util.Set; + public class ClientPool { protected static final Logger logger = LoggerFactory.getLogger(ClientPool.class); protected GenericKeyedObjectPool clientPool; @@ -42,6 +44,18 @@ public void configForKey(String key, C config) { this.clientFactory.configForKey(key, config); } + public C removeConfig(String key) { + return this.clientFactory.removeConfig(key); + } + + public Set configKeys() { + return this.clientFactory.configKeys(); + } + + public C getConfig(String key) { + return this.clientFactory.getConfig(key); + } + /** * Get a client object which is idle from the pool. * Once the client is hold by the caller, it will be marked as active state and cannot be fetched by other caller. diff --git a/sdk-core/src/main/java/io/milvus/pool/PoolClientFactory.java b/sdk-core/src/main/java/io/milvus/pool/PoolClientFactory.java index 5dabef75d..6f736507c 100644 --- a/sdk-core/src/main/java/io/milvus/pool/PoolClientFactory.java +++ b/sdk-core/src/main/java/io/milvus/pool/PoolClientFactory.java @@ -10,6 +10,7 @@ import java.lang.reflect.Constructor; import java.lang.reflect.Method; +import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; @@ -35,8 +36,20 @@ public PoolClientFactory(C configDefault, String clientClassName) throws ClassNo } } - public void configForKey(String key, C config) { - configForKeys.put(key, config); + public C configForKey(String key, C config) { + return configForKeys.put(key, config); + } + + public C removeConfig(String key) { + return configForKeys.remove(key); + } + + public Set configKeys() { + return configForKeys.keySet(); + } + + public C getConfig(String key) { + return configForKeys.get(key); } @Override diff --git a/sdk-core/src/test/java/io/milvus/v2/client/MilvusClientV2DockerTest.java b/sdk-core/src/test/java/io/milvus/v2/client/MilvusClientV2DockerTest.java index c353e39f8..3e5c3397e 100644 --- a/sdk-core/src/test/java/io/milvus/v2/client/MilvusClientV2DockerTest.java +++ b/sdk-core/src/test/java/io/milvus/v2/client/MilvusClientV2DockerTest.java @@ -2134,6 +2134,10 @@ void testClientPool() { .dbName(dummyDb) .rpcDeadlineMs(100L) .build()); + Set keys = pool.configKeys(); + Assertions.assertTrue(keys.contains(dummyDb)); + ConnectConfig dummyConfig = pool.getConfig(dummyDb); + Assertions.assertEquals(dummyConfig.getDbName(), dummyDb); List threadList = new ArrayList<>(); int threadCount = 10; @@ -2165,6 +2169,8 @@ void testClientPool() { // get client connect to the dummy db MilvusClientV2 dummyClient = pool.getClient(dummyDb); Assertions.assertEquals(dummyClient.currentUsedDatabase(), dummyDb); + pool.removeConfig(dummyDb); + Assertions.assertNull(pool.getConfig(dummyDb)); pool.close(); } catch (Exception e) { System.out.println(e.getMessage());