graceful shutdown

grpc · Mar 21, 2023 · f653807 · f653807
1 parent f6af33f
commit f653807
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 13 deletions.
diff --git a/okhttp/src/main/java/io/grpc/okhttp/OkHttpServerTransport.java b/okhttp/src/main/java/io/grpc/okhttp/OkHttpServerTransport.java
@@ -136,6 +136,8 @@ final class OkHttpServerTransport implements ServerTransport,
   /** Non-{@code null} when waiting for forceful close GOAWAY to be sent. */
   @GuardedBy("lock")
   private ScheduledFuture<?> forcefulCloseTimer;
+  @GuardedBy("lock")
+  private Long gracefulShutdownPeriod = null;
 
   public OkHttpServerTransport(Config config, Socket bareSocket) {
     this.config = Preconditions.checkNotNull(config, "config");
@@ -233,8 +235,11 @@ public void data(boolean outFinished, int streamId, Buffer source, int byteCount
       if (config.maxConnectionAgeInNanos != MAX_CONNECTION_AGE_NANOS_DISABLED) {
         long maxConnectionAgeInNanos =
             (long) ((.9D + Math.random() * .2D) * config.maxConnectionAgeInNanos);
+        synchronized (lock) {
+          gracefulShutdownPeriod = config.maxConnectionAgeGraceInNanos;
+        }
         maxConnectionAgeMonitor = scheduledExecutorService.schedule(
-            new LogExceptionRunnable(() -> shutdown(config.maxConnectionAgeGraceInNanos)),
+            new LogExceptionRunnable(this::shutdown),
             maxConnectionAgeInNanos,
             TimeUnit.NANOSECONDS);
       }
@@ -254,10 +259,6 @@ public void data(boolean outFinished, int streamId, Buffer source, int byteCount
 
   @Override
   public void shutdown() {
-    shutdown(null);
-  }
-
-  private void shutdown(@Nullable Long graceTimeInNanos) {
     synchronized (lock) {
       if (gracefulShutdown || abruptShutdown) {
         return;
@@ -271,7 +272,7 @@ private void shutdown(@Nullable Long graceTimeInNanos) {
         // we also set a timer to limit the upper bound in case the PING is excessively stalled or
         // the client is malicious.
         secondGoawayTimer = scheduledExecutorService.schedule(
-            () -> triggerGracefulSecondGoaway(graceTimeInNanos),
+            this::triggerGracefulSecondGoaway,
             GRACEFUL_SHUTDOWN_PING_TIMEOUT_NANOS, TimeUnit.NANOSECONDS);
         frameWriter.goAway(Integer.MAX_VALUE, ErrorCode.NO_ERROR, new byte[0]);
         frameWriter.ping(false, 0, GRACEFUL_SHUTDOWN_PING);
@@ -280,7 +281,7 @@ private void shutdown(@Nullable Long graceTimeInNanos) {
     }
   }
 
-  private void triggerGracefulSecondGoaway(@Nullable Long gracePeriodNanos) {
+  private void triggerGracefulSecondGoaway() {
     synchronized (lock) {
       if (secondGoawayTimer == null) {
         return;
@@ -294,9 +295,9 @@ private void triggerGracefulSecondGoaway(@Nullable Long gracePeriodNanos) {
       } else {
         frameWriter.flush();
       }
-      if (gracePeriodNanos != null) {
+      if (gracefulShutdownPeriod != null) {
         forcefulCloseTimer = scheduledExecutorService.schedule(
-            this::triggerForcefulClose, gracePeriodNanos, TimeUnit.NANOSECONDS);
+            this::triggerForcefulClose, gracefulShutdownPeriod, TimeUnit.NANOSECONDS);
       }
     }
   }
@@ -935,7 +936,7 @@ public void ping(boolean ack, int payload1, int payload2) {
           return;
         }
         if (GRACEFUL_SHUTDOWN_PING == payload) {
-          triggerGracefulSecondGoaway(null);
+          triggerGracefulSecondGoaway();
           return;
         }
         log.log(Level.INFO, "Received unexpected ping ack: " + payload);

diff --git a/okhttp/src/test/java/io/grpc/okhttp/OkHttpServerTransportTest.java b/okhttp/src/test/java/io/grpc/okhttp/OkHttpServerTransportTest.java
@@ -156,7 +156,7 @@ public void startThenShutdown() throws Exception {
   @Test
   public void maxConnectionAge() throws Exception {
     serverBuilder.maxConnectionAge(5, TimeUnit.SECONDS)
-        .maxConnectionAgeGrace(1, TimeUnit.SECONDS);
+        .maxConnectionAgeGrace(3, TimeUnit.SECONDS);
     initTransport();
     handshake();
     clientFrameWriter.headers(1, Arrays.asList(
@@ -170,8 +170,7 @@ public void maxConnectionAge() throws Exception {
         new Header("some-client-sent-trailer", "trailer-value")));
     pingPong();
     fakeClock.forwardNanos(TimeUnit.SECONDS.toNanos(6)); // > 1.1 * 5
-    fakeClock.forwardNanos(TimeUnit.SECONDS.toNanos(1));
-    verifyGracefulShutdown(1, TimeUnit.SECONDS.toNanos(1));
+    verifyGracefulShutdown(1, TimeUnit.SECONDS.toNanos(3));
   }
 
   @Test
@@ -1231,6 +1230,7 @@ private void verifyGracefulShutdown(int lastStreamId, @Nullable Long gracePeriod
     assertThat(clientFrameReader.nextFrame(clientFramesRead)).isTrue();
     verify(clientFramesRead).goAway(lastStreamId, ErrorCode.NO_ERROR, ByteString.EMPTY);
     if (gracePeriodNanos != null) {
+      assertThat(fakeClock.getPendingTasks().size()).isEqualTo(1);
       assertThat(fakeClock.forwardNanos(gracePeriodNanos)).isEqualTo(1);
       assertThat(socket.isClosed()).isTrue();
     }