From 196a69485764e8cf8c1cc82d68d2f3cd245d4800 Mon Sep 17 00:00:00 2001
From: Daniel Mai <daniel@dgraph.io>
Date: Fri, 24 Dec 2021 13:17:08 -0800
Subject: [PATCH] feat: Add shutdown method to client. (#173)

Convenience method to shutdown all ManagedChannels in the client.
---
 src/main/java/io/dgraph/DgraphAsyncClient.java | 18 ++++++++++++++++++
 src/main/java/io/dgraph/DgraphClient.java      |  5 +++++
 .../java/io/dgraph/DgraphIntegrationTest.java  |  3 +++
 3 files changed, 26 insertions(+)

diff --git a/src/main/java/io/dgraph/DgraphAsyncClient.java b/src/main/java/io/dgraph/DgraphAsyncClient.java
index 173a1375..85bafa6c 100644
--- a/src/main/java/io/dgraph/DgraphAsyncClient.java
+++ b/src/main/java/io/dgraph/DgraphAsyncClient.java
@@ -21,7 +21,9 @@
 import io.dgraph.DgraphProto.Payload;
 import io.dgraph.DgraphProto.TxnContext;
 import io.dgraph.DgraphProto.Version;
+import io.grpc.Channel;
 import io.grpc.Context;
+import io.grpc.ManagedChannel;
 import io.grpc.Metadata;
 import io.grpc.Status;
 import io.grpc.StatusRuntimeException;
@@ -371,4 +373,20 @@ public AsyncTransaction newReadOnlyTransaction() {
   public AsyncTransaction newReadOnlyTransaction(TxnContext context) {
     return new AsyncTransaction(this, this.anyClient(), context, true);
   }
+
+  /** Calls %{@link io.grpc.ManagedChannel#shutdown} on all connections for this client */
+  public CompletableFuture<Void> shutdown() {
+    CompletableFuture<Void> future =
+        CompletableFuture.runAsync(
+            () -> {
+              for (DgraphGrpc.DgraphStub stub : this.stubs) {
+                Channel chan = stub.getChannel();
+                if (chan instanceof ManagedChannel) {
+                  ((ManagedChannel) chan).shutdown();
+                }
+              }
+            },
+            this.executor);
+    return future;
+  }
 }
diff --git a/src/main/java/io/dgraph/DgraphClient.java b/src/main/java/io/dgraph/DgraphClient.java
index 6d27e518..7008032d 100644
--- a/src/main/java/io/dgraph/DgraphClient.java
+++ b/src/main/java/io/dgraph/DgraphClient.java
@@ -228,4 +228,9 @@ public void login(String userid, String password) {
   public void loginIntoNamespace(String userid, String password, long namespace) {
     asyncClient.loginIntoNamespace(userid, password, namespace).join();
   }
+
+  /** Calls %{@link io.grpc.ManagedChannel#shutdown} on all connections for this client */
+  public void shutdown() {
+    asyncClient.shutdown().join();
+  }
 }
diff --git a/src/test/java/io/dgraph/DgraphIntegrationTest.java b/src/test/java/io/dgraph/DgraphIntegrationTest.java
index 2b516ae2..92393f7c 100644
--- a/src/test/java/io/dgraph/DgraphIntegrationTest.java
+++ b/src/test/java/io/dgraph/DgraphIntegrationTest.java
@@ -85,8 +85,11 @@ public static void beforeClass() throws InterruptedException {
 
   @AfterClass
   public static void afterClass() throws InterruptedException {
+    // Shutdown channel connections
     channel1.shutdown().awaitTermination(5, TimeUnit.SECONDS);
     channel2.shutdown().awaitTermination(5, TimeUnit.SECONDS);
     channel3.shutdown().awaitTermination(5, TimeUnit.SECONDS);
+    // Or, alternatively, shutdown channels from the client
+    dgraphClient.shutdown();
   }
 }