Skip to content

Commit ddf9ba7

Browse files
authored
Add some util methods for ClientPool (#1598)
Signed-off-by: yhmo <[email protected]>
1 parent 5d85287 commit ddf9ba7

File tree

3 files changed

+35
-2
lines changed

3 files changed

+35
-2
lines changed

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
@@ -1976,6 +1976,10 @@ void testClientPool() {
19761976
.dbName(dummyDb)
19771977
.rpcDeadlineMs(100L)
19781978
.build());
1979+
Set<String> keys = pool.configKeys();
1980+
Assertions.assertTrue(keys.contains(dummyDb));
1981+
ConnectConfig dummyConfig = pool.getConfig(dummyDb);
1982+
Assertions.assertEquals(dummyConfig.getDbName(), dummyDb);
19791983

19801984
List<Thread> threadList = new ArrayList<>();
19811985
int threadCount = 10;
@@ -2007,6 +2011,8 @@ void testClientPool() {
20072011
// get client connect to the dummy db
20082012
MilvusClientV2 dummyClient = pool.getClient(dummyDb);
20092013
Assertions.assertEquals(dummyClient.currentUsedDatabase(), dummyDb);
2014+
pool.removeConfig(dummyDb);
2015+
Assertions.assertNull(pool.getConfig(dummyDb));
20102016
pool.close();
20112017
} catch (Exception e) {
20122018
System.out.println(e.getMessage());

0 commit comments

Comments
 (0)