Fix exception in GrpcRetryer.

temporal-serviceclient/build.gradle

@@ -24,6 +24,7 @@ dependencies {
     api "org.slf4j:slf4j-api:$slf4jVersion"
 
     testImplementation project(':temporal-testing')
+    testImplementation "io.grpc:grpc-testing:${grpcVersion}"
     testImplementation "junit:junit:${junitVersion}"
     testImplementation "org.mockito:mockito-core:${mockitoVersion}"
 

temporal-serviceclient/src/main/java/io/temporal/serviceclient/ChannelManager.java

@@ -27,8 +27,9 @@
 import io.grpc.health.v1.HealthGrpc;
 import io.grpc.netty.shaded.io.grpc.netty.NettyChannelBuilder;
 import io.grpc.stub.MetadataUtils;
-import io.temporal.api.workflowservice.v1.GetSystemInfoResponse;
+import io.temporal.api.workflowservice.v1.GetSystemInfoResponse.Capabilities;
 import io.temporal.internal.retryer.GrpcRetryer;
+import io.temporal.internal.retryer.GrpcRetryer.GrpcRetryerOptions;
 import java.time.Duration;
 import java.util.Collection;
 import java.util.List;
@@ -87,7 +88,7 @@ final class ChannelManager {
   private final Channel interceptedChannel;
   private final HealthGrpc.HealthBlockingStub healthBlockingStub;
 
-  private final CompletableFuture<GetSystemInfoResponse.Capabilities> serverCapabilitiesFuture =
+  private final CompletableFuture<Capabilities> serverCapabilitiesFuture =
       new CompletableFuture<>();
 
   public ChannelManager(
@@ -289,8 +290,8 @@ public void connect(String healthCheckServiceName, @Nullable Duration timeout) {
     if (timeout == null) {
       timeout = options.getRpcTimeout();
     }
-    GrpcRetryer.GrpcRetryerOptions grpcRetryerOptions =
-        new GrpcRetryer.GrpcRetryerOptions(
+    GrpcRetryerOptions grpcRetryerOptions =
+        new GrpcRetryerOptions(
             RpcRetryOptions.newBuilder().setExpiration(timeout).validateBuildWithDefaults(), null);
 
     new GrpcRetryer(getServerCapabilities())
@@ -310,30 +311,24 @@ public void connect(String healthCheckServiceName, @Nullable Duration timeout) {
    */
   public HealthCheckResponse healthCheck(
       String healthCheckServiceName, @Nullable Duration timeout) {
-    HealthGrpc.HealthBlockingStub stub;
-    if (timeout != null) {
-      stub =
-          this.healthBlockingStub.withDeadline(
-              Deadline.after(
-                  options.getHealthCheckAttemptTimeout().toMillis(), TimeUnit.MILLISECONDS));
-    } else {
-      stub = this.healthBlockingStub;
+    if (timeout == null) {
+      timeout = options.getHealthCheckAttemptTimeout();
     }
-    return stub.check(HealthCheckRequest.newBuilder().setService(healthCheckServiceName).build());
+    return this.healthBlockingStub
+        .withDeadline(deadlineFrom(timeout))
+        .check(HealthCheckRequest.newBuilder().setService(healthCheckServiceName).build());
   }
 
-  public Supplier<GetSystemInfoResponse.Capabilities> getServerCapabilities() {
-    return () -> {
-      synchronized (serverCapabilitiesFuture) {
-        GetSystemInfoResponse.Capabilities capabilities = serverCapabilitiesFuture.getNow(null);
-        if (capabilities == null) {
-          serverCapabilitiesFuture.complete(
-              SystemInfoInterceptor.getServerCapabilitiesOrThrow(interceptedChannel, null));
-          capabilities = serverCapabilitiesFuture.getNow(null);
-        }
-        return capabilities;
-      }
-    };
+  public Supplier<Capabilities> getServerCapabilities() {
+    return () ->
+        SystemInfoInterceptor.getServerCapabilitiesWithRetryOrThrow(
+            serverCapabilitiesFuture,
+            interceptedChannel,
+            deadlineFrom(options.getHealthCheckAttemptTimeout()));
+  }
+
+  private static Deadline deadlineFrom(Duration duration) {
+    return Deadline.after(duration.toMillis(), TimeUnit.MILLISECONDS);
   }
 
   public void shutdown() {

temporal-serviceclient/src/main/java/io/temporal/serviceclient/SystemInfoInterceptor.java

@@ -23,16 +23,22 @@
 import io.grpc.*;
 import io.temporal.api.workflowservice.v1.GetSystemInfoRequest;
 import io.temporal.api.workflowservice.v1.GetSystemInfoResponse;
+import io.temporal.api.workflowservice.v1.GetSystemInfoResponse.Capabilities;
 import io.temporal.api.workflowservice.v1.WorkflowServiceGrpc;
+import io.temporal.internal.retryer.GrpcRetryer;
+import io.temporal.internal.retryer.GrpcRetryer.GrpcRetryerOptions;
+import java.time.Duration;
+import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.TimeUnit;
+import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 
 public class SystemInfoInterceptor implements ClientInterceptor {
 
-  private final CompletableFuture<GetSystemInfoResponse.Capabilities> serverCapabilitiesFuture;
+  private final CompletableFuture<Capabilities> serverCapabilitiesFuture;
 
-  public SystemInfoInterceptor(
-      CompletableFuture<GetSystemInfoResponse.Capabilities> serverCapabilitiesFuture) {
+  public SystemInfoInterceptor(CompletableFuture<Capabilities> serverCapabilitiesFuture) {
     this.serverCapabilitiesFuture = serverCapabilitiesFuture;
   }
 
@@ -63,8 +69,7 @@ public void onMessage(RespT message) {
                   @Override
                   public void onClose(Status status, Metadata trailers) {
                     if (Status.UNIMPLEMENTED.getCode().equals(status.getCode())) {
-                      serverCapabilitiesFuture.complete(
-                          GetSystemInfoResponse.Capabilities.getDefaultInstance());
+                      serverCapabilitiesFuture.complete(Capabilities.getDefaultInstance());
                     }
                     super.onClose(status, trailers);
                   }
@@ -87,7 +92,39 @@ public void onClose(Status status, Metadata trailers) {
     };
   }
 
-  public static GetSystemInfoResponse.Capabilities getServerCapabilitiesOrThrow(
+  public static Capabilities getServerCapabilitiesWithRetryOrThrow(
+      @Nonnull CompletableFuture<Capabilities> future,
+      @Nonnull Channel channel,
+      @Nullable Deadline deadline) {
+    Capabilities capabilities = future.getNow(null);
+    if (capabilities == null) {
+      synchronized (Objects.requireNonNull(future)) {
+        capabilities = future.getNow(null);
+        if (capabilities == null) {
+          if (deadline == null) {
+            deadline = Deadline.after(30, TimeUnit.SECONDS);
+          }
+          Deadline computedDeadline = deadline;
+          RpcRetryOptions rpcRetryOptions =
+              RpcRetryOptions.newBuilder()
+                  .setExpiration(
+                      Duration.ofMillis(computedDeadline.timeRemaining(TimeUnit.MILLISECONDS)))
+                  .validateBuildWithDefaults();
+          GrpcRetryerOptions grpcRetryerOptions =
+              new GrpcRetryerOptions(rpcRetryOptions, computedDeadline);
+          capabilities =
+              new GrpcRetryer(Capabilities::getDefaultInstance)
+                  .retryWithResult(
+                      () -> getServerCapabilitiesOrThrow(channel, computedDeadline),
+                      grpcRetryerOptions);
+          future.complete(capabilities);
+        }
+      }
+    }
+    return capabilities;
+  }
+
+  public static Capabilities getServerCapabilitiesOrThrow(
       Channel channel, @Nullable Deadline deadline) {
     try {
       return WorkflowServiceGrpc.newBlockingStub(channel)
@@ -96,7 +133,7 @@ public static GetSystemInfoResponse.Capabilities getServerCapabilitiesOrThrow(
           .getCapabilities();
     } catch (StatusRuntimeException ex) {
       if (Status.Code.UNIMPLEMENTED.equals(ex.getStatus().getCode())) {
-        return GetSystemInfoResponse.Capabilities.getDefaultInstance();
+        return Capabilities.getDefaultInstance();
       }
       throw ex;
     }