Fix a bug of ClientPool that rpc channel is not properly closed

Signed-off-by: yhmo <yihua.mo@zilliz.com>

Author: yhmo

examples/src/main/java/io/milvus/v2/ConsistencyLevelExample.java

@@ -136,47 +136,51 @@ private static void testSessionLevel() throws Exception {
                 .build();
         MilvusClientV2Pool pool = new MilvusClientV2Pool(poolConfig, connectConfig);
 
-        // The same process, different MilvusClient object, insert and search with Session level.
-        // The Session level ensure that the newly inserted data instantaneously become searchable.
-        Gson gson = new Gson();
-        for (int i = 0; i < 100; i++) {
-            List<Float> vector = CommonUtils.generateFloatVector(VECTOR_DIM);
-            JsonObject row = new JsonObject();
-            row.addProperty("id", i);
-            row.add("vector", gson.toJsonTree(vector));
-
-            // insert by a MilvusClient
-            String clientName1 = String.format("client_%d", i % 10);
-            MilvusClientV2 client1 = pool.getClient(clientName1);
-            client1.insert(InsertReq.builder()
-                    .collectionName(collectionName)
-                    .data(Collections.singletonList(row))
-                    .build());
-            pool.returnClient(clientName1, client1); // don't forget to return the client to pool
-            System.out.println("insert");
-
-            // search by another MilvusClient, use the just inserted vector to search
-            // the returned item is expected to be the just inserted item
-            String clientName2 = String.format("client_%d", i % 10 + 1);
-            MilvusClientV2 client2 = pool.getClient(clientName2);
-            SearchResp searchR = client2.search(SearchReq.builder()
-                    .collectionName(collectionName)
-                    .data(Collections.singletonList(new FloatVec(vector)))
-                    .limit(1)
-                    .build());
-            pool.returnClient(clientName2, client2); // don't forget to return the client to pool
-            List<List<SearchResp.SearchResult>> searchResults = searchR.getSearchResults();
-            List<SearchResp.SearchResult> results = searchResults.get(0);
-            if (results.size() != 1) {
-                throw new RuntimeException("Search result is empty");
-            }
-            if (i != (Long) results.get(0).getId()) {
-                throw new RuntimeException("The just inserted entity is not found");
+        try {
+            // The same process, different MilvusClient object, insert and search with Session level.
+            // The Session level ensure that the newly inserted data instantaneously become searchable.
+            Gson gson = new Gson();
+            for (int i = 0; i < 100; i++) {
+                List<Float> vector = CommonUtils.generateFloatVector(VECTOR_DIM);
+                JsonObject row = new JsonObject();
+                row.addProperty("id", i);
+                row.add("vector", gson.toJsonTree(vector));
+
+                // insert by a MilvusClient
+                String clientName1 = String.format("client_%d", i % 10);
+                MilvusClientV2 client1 = pool.getClient(clientName1);
+                client1.insert(InsertReq.builder()
+                        .collectionName(collectionName)
+                        .data(Collections.singletonList(row))
+                        .build());
+                pool.returnClient(clientName1, client1); // don't forget to return the client to pool
+                System.out.println("insert");
+
+                // search by another MilvusClient, use the just inserted vector to search
+                // the returned item is expected to be the just inserted item
+                String clientName2 = String.format("client_%d", i % 10 + 1);
+                MilvusClientV2 client2 = pool.getClient(clientName2);
+                SearchResp searchR = client2.search(SearchReq.builder()
+                        .collectionName(collectionName)
+                        .data(Collections.singletonList(new FloatVec(vector)))
+                        .limit(1)
+                        .build());
+                pool.returnClient(clientName2, client2); // don't forget to return the client to pool
+                List<List<SearchResp.SearchResult>> searchResults = searchR.getSearchResults();
+                List<SearchResp.SearchResult> results = searchResults.get(0);
+                if (results.size() != 1) {
+                    throw new RuntimeException("Search result is empty");
+                }
+                if (i != (Long) results.get(0).getId()) {
+                    throw new RuntimeException("The just inserted entity is not found");
+                }
+                System.out.println("search");
             }
-            System.out.println("search");
-        }
 
-        System.out.println("Session level is working fine");
+            System.out.println("Session level is working fine");
+        } finally {
+            pool.close();
+        }
     }
 
     private static void testBoundedLevel() {
@@ -207,5 +211,7 @@ public static void main(String[] args) throws Exception {
         testBoundedLevel();
         System.out.println("==============================================================");
         testEventuallyLevel();
+
+        client.close();
     }
 }

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

@@ -38,6 +38,7 @@ protected ClientCache(String key, GenericKeyedObjectPool<String, T> pool) {
             @Override
             public Thread newThread(@NotNull Runnable r) {
                 Thread t = new Thread(r);
+                t.setDaemon(true);
                 t.setPriority(Thread.MAX_PRIORITY); // set the highest priority for the timer
                 return t;
             }
@@ -189,6 +190,16 @@ public void run() {
 
     public void stopTimer() {
         scheduler.shutdown();
+
+        // Return all active and retired clients to the pool so they can be properly destroyed
+        for (ClientWrapper<T> wrapper : activeClientList) {
+            returnToPool(wrapper.getRawClient());
+        }
+        activeClientList.clear();
+        for (ClientWrapper<T> wrapper : retireClientList) {
+            returnToPool(wrapper.getRawClient());
+        }
+        retireClientList.clear();
     }
 
     public T getClient() {
@@ -323,6 +334,10 @@ public T getClient() {
             return this.client;
         }
 
+        public T getRawClient() {
+            return this.client;
+        }
+
         public void returnClient() {
             this.refCount.decrementAndGet();
         }