refactor: work in deadlines (#21)

BREAKING CHANGE: changed new apis added in last release before they're adopted

Author: aaron-steinfeld

grpc-client-utils/build.gradle.kts

@@ -6,19 +6,11 @@ plugins {
 }
 
 dependencies {
-  implementation(project(":grpc-context-utils"))
+  api("io.grpc:grpc-context:1.40.0")
+  api("io.grpc:grpc-api:1.40.0")
 
-  // Logging
+  implementation(project(":grpc-context-utils"))
   implementation("org.slf4j:slf4j-api:1.7.30")
-  // End Logging
-
-  // grpc
-  implementation("io.grpc:grpc-core:1.40.0")
-  constraints {
-    implementation("com.google.guava:guava:30.0-jre") {
-      because("https://snyk.io/vuln/SNYK-JAVA-COMGOOGLEGUAVA-1015415")
-    }
-  }
 
   testImplementation("org.junit.jupiter:junit-jupiter:5.7.0")
   testImplementation("org.mockito:mockito-core:3.12.1")

grpc-client-utils/src/main/java/org/hypertrace/core/grpcutils/client/GrpcChannelRegistry.java

@@ -1,12 +1,10 @@
 package org.hypertrace.core.grpcutils.client;
 
-import static java.time.temporal.ChronoUnit.MINUTES;
-import static java.time.temporal.ChronoUnit.SECONDS;
+import static java.util.concurrent.TimeUnit.MILLISECONDS;
 
+import io.grpc.Deadline;
 import io.grpc.ManagedChannel;
 import io.grpc.ManagedChannelBuilder;
-import java.time.Clock;
-import java.time.Instant;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.TimeUnit;
@@ -16,18 +14,8 @@
 public class GrpcChannelRegistry {
   private static final Logger LOG = LoggerFactory.getLogger(GrpcChannelRegistry.class);
   private final Map<String, ManagedChannel> channelMap = new ConcurrentHashMap<>();
-  private final Clock clock;
   private volatile boolean isShutdown = false;
 
-  @Deprecated
-  public GrpcChannelRegistry() {
-    this(Clock.systemUTC());
-  }
-
-  public GrpcChannelRegistry(Clock clock) {
-    this.clock = clock;
-  }
-
   /**
    * Use either {@link #forSecureAddress(String, int)} or {@link #forPlaintextAddress(String, int)}
    */
@@ -68,10 +56,10 @@ private String getChannelId(String host, int port, boolean isPlaintext) {
   /**
    * Shuts down channels using a default deadline of 1 minute.
    *
-   * @see #shutdown(Instant)
+   * @see #shutdown(Deadline)
    */
   public void shutdown() {
-    this.shutdown(this.clock.instant().plus(1, MINUTES));
+    this.shutdown(Deadline.after(1, TimeUnit.MINUTES));
   }
 
   /**
@@ -96,11 +84,11 @@ public void shutdown() {
    *
    * @param deadline Deadline for all channels to complete graceful shutdown.
    */
-  public void shutdown(Instant deadline) {
+  public void shutdown(Deadline deadline) {
     channelMap.forEach(this::initiateChannelShutdown);
     channelMap.keySet().stream()
         .filter(channelId -> !this.waitForGracefulShutdown(channelId, deadline))
-        .forEach(this::forceShutdown);
+        .forEach(channelId -> this.forceShutdown(channelId, deadline.offset(5, TimeUnit.SECONDS)));
 
     this.isShutdown = true;
     this.channelMap.clear();
@@ -111,30 +99,28 @@ private void initiateChannelShutdown(String channelId, ManagedChannel managedCha
     managedChannel.shutdown();
   }
 
-  private boolean waitForGracefulShutdown(String channelId, Instant deadline) {
+  private boolean waitForGracefulShutdown(String channelId, Deadline deadline) {
     boolean successfullyShutdown = this.waitForTermination(channelId, deadline);
     if (successfullyShutdown) {
       LOG.info("Shutdown channel successfully [{}]", channelId);
     }
     return successfullyShutdown;
   }
 
-  private void forceShutdown(String channelId) {
+  private void forceShutdown(String channelId, Deadline deadline) {
     LOG.error("Shutting down channel [{}] forcefully", channelId);
     this.channelMap.get(channelId).shutdownNow();
-    Instant forceShutdownDeadline = this.clock.instant().plus(5, SECONDS);
-    if (this.waitForTermination(channelId, forceShutdownDeadline)) {
+    if (this.waitForTermination(channelId, deadline)) {
       LOG.error("Forced channel [{}] shutdown successful", channelId);
     } else {
       LOG.error("Unable to force channel [{}] shutdown in 5s - giving up!", channelId);
     }
   }
 
-  private boolean waitForTermination(String channelId, Instant deadline) {
+  private boolean waitForTermination(String channelId, Deadline deadline) {
     ManagedChannel managedChannel = this.channelMap.get(channelId);
-    long millisRemaining = Math.max(0, deadline.toEpochMilli() - this.clock.millis());
     try {
-      if (!managedChannel.awaitTermination(millisRemaining, TimeUnit.MILLISECONDS)) {
+      if (!managedChannel.awaitTermination(deadline.timeRemaining(MILLISECONDS), MILLISECONDS)) {
         LOG.error("Channel [{}] did not shut down after waiting", channelId);
       }
     } catch (InterruptedException ex) {

grpc-client-utils/src/test/java/org/hypertrace/core/grpcutils/client/GrpcChannelRegistryTest.java

@@ -12,11 +12,10 @@
 import static org.mockito.Mockito.when;
 
 import io.grpc.Channel;
+import io.grpc.Deadline;
+import io.grpc.Deadline.Ticker;
 import io.grpc.ManagedChannel;
 import io.grpc.ManagedChannelBuilder;
-import java.time.Clock;
-import java.time.Instant;
-import java.time.ZoneOffset;
 import java.util.concurrent.TimeUnit;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -30,7 +29,7 @@ class GrpcChannelRegistryTest {
 
   @BeforeEach
   void beforeEach() {
-    this.channelRegistry = new GrpcChannelRegistry(Clock.systemUTC());
+    this.channelRegistry = new GrpcChannelRegistry();
   }
 
   @Test
@@ -51,10 +50,10 @@ void reusesChannelsForDuplicateRequests() {
     assertNotSame(firstChannelSecure, this.channelRegistry.forSecureAddress("bar", 1000));
   }
 
+  @SuppressWarnings("rawtypes")
   @Test
   void shutdownAllChannelsOnShutdown() throws InterruptedException {
-    this.channelRegistry =
-        new GrpcChannelRegistry(Clock.fixed(Instant.ofEpochMilli(0), ZoneOffset.UTC));
+    this.channelRegistry = new GrpcChannelRegistry();
     try (MockedStatic<ManagedChannelBuilder> mockedBuilderStatic =
         Mockito.mockStatic(ManagedChannelBuilder.class)) {
 
@@ -78,8 +77,12 @@ void shutdownAllChannelsOnShutdown() throws InterruptedException {
       // Second does not
       when(secondChannel.isTerminated()).thenReturn(false);
 
+      Ticker mockTicker = mock(Ticker.class);
+
+      when(mockTicker.nanoTime()).thenReturn(0L);
+
       // Wait for 10ms (test clock fixed at 0)
-      this.channelRegistry.shutdown(Instant.ofEpochMilli(10));
+      this.channelRegistry.shutdown(Deadline.after(10, TimeUnit.MILLISECONDS, mockTicker));
 
       // First channel requests shutdown, waits, succeeds and checks result
       InOrder firstChannelVerifier = inOrder(firstChannel);
@@ -95,8 +98,8 @@ void shutdownAllChannelsOnShutdown() throws InterruptedException {
       secondChannelVerifier.verify(secondChannel).awaitTermination(10, TimeUnit.MILLISECONDS);
       secondChannelVerifier.verify(secondChannel).isTerminated();
       secondChannelVerifier.verify(secondChannel).shutdownNow();
-      // hardcoded 5s for force shutdown
-      secondChannelVerifier.verify(secondChannel).awaitTermination(5000, TimeUnit.MILLISECONDS);
+      // hardcoded additional 5s for force shutdown
+      secondChannelVerifier.verify(secondChannel).awaitTermination(5010, TimeUnit.MILLISECONDS);
       secondChannelVerifier.verify(secondChannel).isTerminated();
     }
   }

grpc-server-utils/build.gradle.kts

@@ -10,19 +10,11 @@ tasks.test {
 }
 
 dependencies {
-  implementation(project(":grpc-context-utils"))
+  api("io.grpc:grpc-context:1.40.0")
+  api("io.grpc:grpc-api:1.40.0")
 
-  // Logging
+  implementation(project(":grpc-context-utils"))
   implementation("org.slf4j:slf4j-api:1.7.30")
-  // End Logging
-
-  // grpc
-  implementation("io.grpc:grpc-core:1.40.0")
-  constraints {
-    implementation("com.google.guava:guava:30.0-jre") {
-      because("https://snyk.io/vuln/SNYK-JAVA-COMGOOGLEGUAVA-1015415")
-    }
-  }
 
   annotationProcessor("org.projectlombok:lombok:1.18.20")
   compileOnly("org.projectlombok:lombok:1.18.20")

grpc-server-utils/src/main/java/org/hypertrace/core/grpcutils/server/ServerManagementUtil.java

@@ -1,43 +1,56 @@
 package org.hypertrace.core.grpcutils.server;
 
+import static java.util.concurrent.TimeUnit.MILLISECONDS;
+import static java.util.concurrent.TimeUnit.SECONDS;
+
+import io.grpc.Deadline;
 import io.grpc.Server;
-import java.time.Duration;
-import java.util.concurrent.TimeUnit;
 import lombok.extern.slf4j.Slf4j;
 
 @Slf4j
 public class ServerManagementUtil {
 
-  public static void shutdownServer(Server grpcServer, String name, Duration timeout) {
+  /**
+   * Attempts to gracefully shutdown the server by the provided deadline. If unsuccessful, it then
+   * uses up to 5 additional seconds to forcefully shutdown the server, returning with the final
+   * success status.
+   *
+   * @param grpcServer
+   * @param name
+   * @param deadline
+   * @return boolean - true if the server is terminated successfully, false otherwise
+   */
+  public static boolean shutdownServer(Server grpcServer, String name, Deadline deadline) {
     log.info("Starting shutdown for service [{}]", name);
     grpcServer.shutdown();
-    boolean gracefullyShutdown = waitForGracefulShutdown(grpcServer, name, timeout);
-    if (!gracefullyShutdown) {
-      forceShutdown(grpcServer, name);
-    }
+    return waitForGracefulShutdown(grpcServer, name, deadline)
+        || forceShutdown(grpcServer, name, deadline.offset(5, SECONDS));
   }
 
-  private static boolean waitForGracefulShutdown(Server grpcServer, String name, Duration timeout) {
-    boolean successfullyShutdown = waitForTermination(grpcServer, name, timeout);
+  private static boolean waitForGracefulShutdown(
+      Server grpcServer, String name, Deadline deadline) {
+    boolean successfullyShutdown = waitForTermination(grpcServer, name, deadline);
     if (successfullyShutdown) {
       log.info("Shutdown service successfully [{}]", name);
     }
     return successfullyShutdown;
   }
 
-  private static void forceShutdown(Server grpcServer, String name) {
+  private static boolean forceShutdown(Server grpcServer, String name, Deadline deadline) {
     log.error("Shutting down service [{}] forcefully", name);
     grpcServer.shutdownNow();
-    if (waitForTermination(grpcServer, name, Duration.ofSeconds(5))) {
+    boolean successfullyShutdown = waitForTermination(grpcServer, name, deadline);
+    if (successfullyShutdown) {
       log.error("Forced service [{}] shutdown successful", name);
     } else {
       log.error("Unable to force service [{}] shutdown in 5s - giving up!", name);
     }
+    return successfullyShutdown;
   }
 
-  private static boolean waitForTermination(Server grpcServer, String name, Duration deadline) {
+  private static boolean waitForTermination(Server grpcServer, String name, Deadline deadline) {
     try {
-      if (!grpcServer.awaitTermination(deadline.toMillis(), TimeUnit.MILLISECONDS)) {
+      if (!grpcServer.awaitTermination(deadline.timeRemaining(MILLISECONDS), MILLISECONDS)) {
         log.error("Service [{}] did not shut down after waiting", name);
       }
     } catch (InterruptedException ex) {

grpc-server-utils/src/test/java/org/hypertrace/core/grpcutils/server/ServerManagementUtilTest.java

@@ -1,13 +1,13 @@
 package org.hypertrace.core.grpcutils.server;
 
+import static java.util.concurrent.TimeUnit.MILLISECONDS;
 import static org.mockito.Mockito.inOrder;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
+import io.grpc.Deadline;
+import io.grpc.Deadline.Ticker;
 import io.grpc.Server;
-import java.time.Duration;
-import java.time.temporal.ChronoUnit;
-import java.util.concurrent.TimeUnit;
 import org.junit.jupiter.api.Test;
 import org.mockito.InOrder;
 
@@ -17,12 +17,14 @@ class ServerManagementUtilTest {
   void canShutdownGracefully() throws InterruptedException {
     Server mockServer = mock(Server.class);
     when(mockServer.isTerminated()).thenReturn(true);
+    Ticker mockTicker = mock(Ticker.class);
+    when(mockTicker.nanoTime()).thenReturn(0L);
     ServerManagementUtil.shutdownServer(
-        mockServer, "mockServer", Duration.of(10, ChronoUnit.MILLIS));
+        mockServer, "mockServer", Deadline.after(10, MILLISECONDS, mockTicker));
 
     InOrder serverVerifier = inOrder(mockServer);
     serverVerifier.verify(mockServer).shutdown();
-    serverVerifier.verify(mockServer).awaitTermination(10, TimeUnit.MILLISECONDS);
+    serverVerifier.verify(mockServer).awaitTermination(10, MILLISECONDS);
     serverVerifier.verify(mockServer).isTerminated();
     serverVerifier.verifyNoMoreInteractions();
   }
@@ -31,15 +33,17 @@ void canShutdownGracefully() throws InterruptedException {
   void canShutdownForcefully() throws InterruptedException {
     Server mockServer = mock(Server.class);
     when(mockServer.isTerminated()).thenReturn(false);
+    Ticker mockTicker = mock(Ticker.class);
+    when(mockTicker.nanoTime()).thenReturn(0L);
     ServerManagementUtil.shutdownServer(
-        mockServer, "mockServer", Duration.of(10, ChronoUnit.MILLIS));
+        mockServer, "mockServer", Deadline.after(10, MILLISECONDS, mockTicker));
 
     InOrder serverVerifier = inOrder(mockServer);
     serverVerifier.verify(mockServer).shutdown();
-    serverVerifier.verify(mockServer).awaitTermination(10, TimeUnit.MILLISECONDS);
+    serverVerifier.verify(mockServer).awaitTermination(10, MILLISECONDS);
     serverVerifier.verify(mockServer).isTerminated();
     serverVerifier.verify(mockServer).shutdownNow();
-    serverVerifier.verify(mockServer).awaitTermination(5000, TimeUnit.MILLISECONDS);
+    serverVerifier.verify(mockServer).awaitTermination(5010, MILLISECONDS);
     serverVerifier.verify(mockServer).isTerminated();
     serverVerifier.verifyNoMoreInteractions();
   }