Add channel shutdown grace period property

Author: ST-DDT

grpc-client-spring-boot-autoconfigure/src/main/java/net/devh/boot/grpc/client/channelfactory/AbstractChannelFactory.java

@@ -17,12 +17,16 @@
 
 package net.devh.boot.grpc.client.channelfactory;
 
+import static java.util.Comparator.comparingLong;
 import static java.util.Objects.requireNonNull;
+import static java.util.concurrent.TimeUnit.MILLISECONDS;
 
 import java.time.Duration;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Map.Entry;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
@@ -302,16 +306,30 @@ public synchronized void close() {
             return;
         }
         this.shutdown = true;
-        for (final ManagedChannel channel : this.channels.values()) {
+        final List<ShutdownRecord> shutdownEntries = new ArrayList<>();
+        for (final Entry<String, ManagedChannel> entry : this.channels.entrySet()) {
+            final ManagedChannel channel = entry.getValue();
             channel.shutdown();
+            final long gracePeriod = this.properties.getChannel(entry.getKey()).getShutdownGracePeriod().toMillis();
+            shutdownEntries.add(new ShutdownRecord(entry.getKey(), channel, gracePeriod));
         }
         try {
-            final long waitLimit = System.currentTimeMillis() + 60_000; // wait 60 seconds at max
-            for (final ManagedChannel channel : this.channels.values()) {
-                int i = 0;
-                do {
-                    log.debug("Awaiting channel shutdown: {} ({}s)", channel, i++);
-                } while (System.currentTimeMillis() < waitLimit && !channel.awaitTermination(1, TimeUnit.SECONDS));
+            final long start = System.currentTimeMillis();
+            shutdownEntries.sort(comparingLong(ShutdownRecord::getGracePeriod));
+
+            for (final ShutdownRecord entry : shutdownEntries) {
+                if (!entry.channel.isTerminated()) {
+                    log.debug("Awaiting channel termination: {}", entry.name);
+
+                    final long waitedTime = System.currentTimeMillis() - start;
+                    final long waitTime = entry.gracePeriod - waitedTime;
+
+                    if (waitTime > 0) {
+                        entry.channel.awaitTermination(waitTime, MILLISECONDS);
+                    }
+                    entry.channel.shutdownNow();
+                }
+                log.debug("Completed channel termination: {}", entry.name);
             }
         } catch (final InterruptedException e) {
             Thread.currentThread().interrupt();
@@ -330,4 +348,22 @@ public synchronized void close() {
         log.debug("GrpcCannelFactory closed (including {} channels)", channelCount);
     }
 
+    private static class ShutdownRecord {
+
+        private final String name;
+        private final ManagedChannel channel;
+        private final long gracePeriod;
+
+        public ShutdownRecord(final String name, final ManagedChannel channel, final long gracePeriod) {
+            this.name = name;
+            this.channel = channel;
+            this.gracePeriod = gracePeriod;
+        }
+
+        long getGracePeriod() {
+            return this.gracePeriod;
+        }
+
+    }
+
 }

grpc-client-spring-boot-autoconfigure/src/main/java/net/devh/boot/grpc/client/config/GrpcChannelProperties.java

@@ -263,6 +263,32 @@ public void setKeepAliveWithoutCalls(final Boolean keepAliveWithoutCalls) {
         this.keepAliveWithoutCalls = keepAliveWithoutCalls;
     }
 
+    // --------------------------------------------------
+
+    @DurationUnit(ChronoUnit.SECONDS)
+    private Duration shutdownGracePeriod;
+    private static final Duration DEFAULT_SHUTDOWN_GRACE_PERIOD = Duration.ofSeconds(30);
+
+    /**
+     * Gets the time to wait for the channel to gracefully shutdown. If set to {@code -1} the server waits forever. If
+     * set to {@code 0} the server will force shutdown immediately. Defaults to {@code 30s}.
+     *
+     * @return The time to wait for a graceful shutdown.
+     */
+    public Duration getShutdownGracePeriod() {
+        return this.shutdownGracePeriod == null ? DEFAULT_SHUTDOWN_GRACE_PERIOD : this.shutdownGracePeriod;
+    }
+
+    /**
+     * Sets the time to wait for the channel to gracefully shutdown (completing all requests). If set to {@code -1} the
+     * channel waits forever. If set to {@code 0} the channel will force shutdown immediately. Defaults to {@code 30s}.
+     *
+     * @param shutdownGracePeriod The time to wait for a graceful shutdown.
+     */
+    public void setShutdownGracePeriod(final Duration shutdownGracePeriod) {
+        this.shutdownGracePeriod = shutdownGracePeriod;
+    }
+
     // --------------------------------------------------
     // Message Transfer
     // --------------------------------------------------