clear waits

apache · Nov 17, 2021 · df61a96 · df61a96
1 parent 3a810a8
commit df61a96
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/computer-k8s/src/main/java/com/baidu/hugegraph/computer/k8s/driver/KubernetesDriver.java b/computer-k8s/src/main/java/com/baidu/hugegraph/computer/k8s/driver/KubernetesDriver.java
@@ -419,6 +419,11 @@ public String log(String jobId, int containerId, long offset, long length,
 
     @Override
     public void close() {
+        for (String jobId : this.waits.keySet()) {
+            this.waits.get(jobId).getKey().cancel(true);
+        }
+        this.waits.clear();
+
         if (this.watch != null) {
             this.watch.close();
             this.watchActive.setFalse();

diff --git a/computer-test/src/main/java/com/baidu/hugegraph/computer/k8s/KubernetesDriverTest.java b/computer-test/src/main/java/com/baidu/hugegraph/computer/k8s/KubernetesDriverTest.java
@@ -234,7 +234,11 @@ public void testWatchJobAndCancel() throws ExecutionException,
         this.driver.watchJob(jobId, params, jobObserver);
 
         this.driver.cancelJob(jobId, params);
-        this.driver.watchJob(jobId, params, jobObserver).get();
+        CompletableFuture<Void> watchJob = this.driver.watchJob(jobId, params,
+                                                                jobObserver);
+        if (watchJob != null) {
+            watchJob.get();
+        }
     }
 
     @Test