Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions sdk-core/src/main/java/io/milvus/pool/ClientPool.java
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,8 @@
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;
Expand Down Expand Up @@ -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<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.
Expand Down
17 changes: 15 additions & 2 deletions sdk-core/src/main/java/io/milvus/pool/PoolClientFactory.java
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -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<String> configKeys() {
return configForKeys.keySet();
}

public C getConfig(String key) {
return configForKeys.get(key);
}

@Override
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2134,6 +2134,10 @@ void testClientPool() {
.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;
Expand Down Expand Up @@ -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());
Expand Down
Loading