Skip to content

Commit 0321ede

Browse files
authored
Add some util methods for ClientPool (#1597)
Signed-off-by: yhmo <yihua.mo@zilliz.com>
1 parent a8d9c0d commit 0321ede

3 files changed

Lines changed: 35 additions & 2 deletions

File tree

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

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@
77
import org.slf4j.Logger;
88
import org.slf4j.LoggerFactory;
99

10+
import java.util.Set;
11+
1012
public class ClientPool<C, T> {
1113
protected static final Logger logger = LoggerFactory.getLogger(ClientPool.class);
1214
protected GenericKeyedObjectPool<String, T> clientPool;
@@ -42,6 +44,18 @@ public void configForKey(String key, C config) {
4244
this.clientFactory.configForKey(key, config);
4345
}
4446

47+
public C removeConfig(String key) {
48+
return this.clientFactory.removeConfig(key);
49+
}
50+
51+
public Set<String> configKeys() {
52+
return this.clientFactory.configKeys();
53+
}
54+
55+
public C getConfig(String key) {
56+
return this.clientFactory.getConfig(key);
57+
}
58+
4559
/**
4660
* Get a client object which is idle from the pool.
4761
* Once the client is hold by the caller, it will be marked as active state and cannot be fetched by other caller.

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

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010

1111
import java.lang.reflect.Constructor;
1212
import java.lang.reflect.Method;
13+
import java.util.Set;
1314
import java.util.concurrent.ConcurrentHashMap;
1415
import java.util.concurrent.ConcurrentMap;
1516

@@ -35,8 +36,20 @@ public PoolClientFactory(C configDefault, String clientClassName) throws ClassNo
3536
}
3637
}
3738

38-
public void configForKey(String key, C config) {
39-
configForKeys.put(key, config);
39+
public C configForKey(String key, C config) {
40+
return configForKeys.put(key, config);
41+
}
42+
43+
public C removeConfig(String key) {
44+
return configForKeys.remove(key);
45+
}
46+
47+
public Set<String> configKeys() {
48+
return configForKeys.keySet();
49+
}
50+
51+
public C getConfig(String key) {
52+
return configForKeys.get(key);
4053
}
4154

4255
@Override

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2134,6 +2134,10 @@ void testClientPool() {
21342134
.dbName(dummyDb)
21352135
.rpcDeadlineMs(100L)
21362136
.build());
2137+
Set<String> keys = pool.configKeys();
2138+
Assertions.assertTrue(keys.contains(dummyDb));
2139+
ConnectConfig dummyConfig = pool.getConfig(dummyDb);
2140+
Assertions.assertEquals(dummyConfig.getDbName(), dummyDb);
21372141

21382142
List<Thread> threadList = new ArrayList<>();
21392143
int threadCount = 10;
@@ -2165,6 +2169,8 @@ void testClientPool() {
21652169
// get client connect to the dummy db
21662170
MilvusClientV2 dummyClient = pool.getClient(dummyDb);
21672171
Assertions.assertEquals(dummyClient.currentUsedDatabase(), dummyDb);
2172+
pool.removeConfig(dummyDb);
2173+
Assertions.assertNull(pool.getConfig(dummyDb));
21682174
pool.close();
21692175
} catch (Exception e) {
21702176
System.out.println(e.getMessage());

0 commit comments

Comments
 (0)