Quellcode durchsuchen

Add some util methods for ClientPool (#1596)

Signed-off-by: yhmo <yihua.mo@zilliz.com>
groot vor 3 Wochen
Ursprung
Commit
27145d3059

+ 14 - 0
sdk-core/src/main/java/io/milvus/pool/ClientPool.java

@@ -7,6 +7,8 @@ import org.apache.commons.pool2.impl.GenericKeyedObjectPoolConfig;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.Set;
+
 public class ClientPool<C, T> {
     protected static final Logger logger = LoggerFactory.getLogger(ClientPool.class);
     protected GenericKeyedObjectPool<String, T> clientPool;
@@ -42,6 +44,18 @@ public class ClientPool<C, T> {
         this.clientFactory.configForKey(key, config);
     }
 
+    public C removeConfig(String key) {
+        return this.clientFactory.removeConfig(key);
+    }
+
+    public Set<String> 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.

+ 15 - 2
sdk-core/src/main/java/io/milvus/pool/PoolClientFactory.java

@@ -10,6 +10,7 @@ import org.slf4j.LoggerFactory;
 
 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 class PoolClientFactory<C, T> extends BaseKeyedPooledObjectFactory<String
         }
     }
 
-    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<String> configKeys() {
+        return configForKeys.keySet();
+    }
+
+    public C getConfig(String key) {
+        return configForKeys.get(key);
     }
 
     @Override

+ 6 - 0
sdk-core/src/test/java/io/milvus/v2/client/MilvusClientV2DockerTest.java

@@ -2134,6 +2134,10 @@ class MilvusClientV2DockerTest {
                     .dbName(dummyDb)
                     .rpcDeadlineMs(100L)
                     .build());
+            Set<String> keys = pool.configKeys();
+            Assertions.assertTrue(keys.contains(dummyDb));
+            ConnectConfig dummyConfig = pool.getConfig(dummyDb);
+            Assertions.assertEquals(dummyConfig.getDbName(), dummyDb);
 
             List<Thread> threadList = new ArrayList<>();
             int threadCount = 10;
@@ -2165,6 +2169,8 @@ class MilvusClientV2DockerTest {
             // 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());