Add ServerCapabilities into RemoteExecutionClient (bazelbuild#18442)

In bazelbuild#18202, we discussed the possibility of conditionally using the new
field exclusively based on the Remote Execution Server's capabilities.

Capture Remote Execution Server's capabilities and store it in
RemoteExecutor implementations for furture usage.

Closes bazelbuild#18269.

PiperOrigin-RevId: 531999688
Change-Id: I370869a45c804af1ec499b9c1654c6977c7ab7d0

Co-authored-by: Son Luong Ngoc <[email protected]>

src/main/java/com/google/devtools/build/lib/remote/ExperimentalGrpcRemoteExecutor.java

@@ -20,6 +20,7 @@
 import build.bazel.remote.execution.v2.ExecutionGrpc;
 import build.bazel.remote.execution.v2.ExecutionGrpc.ExecutionBlockingStub;
 import build.bazel.remote.execution.v2.RequestMetadata;
+import build.bazel.remote.execution.v2.ServerCapabilities;
 import build.bazel.remote.execution.v2.WaitExecutionRequest;
 import com.google.common.base.Preconditions;
 import com.google.devtools.build.lib.authandtls.CallCredentialsProvider;
@@ -56,6 +57,7 @@
 @ThreadSafe
 public class ExperimentalGrpcRemoteExecutor implements RemoteExecutionClient {
 
+  private final ServerCapabilities serverCapabilities;
   private final RemoteOptions remoteOptions;
   private final ReferenceCountedChannel channel;
   private final CallCredentialsProvider callCredentialsProvider;
@@ -64,10 +66,12 @@ public class ExperimentalGrpcRemoteExecutor implements RemoteExecutionClient {
   private final AtomicBoolean closed = new AtomicBoolean();
 
   public ExperimentalGrpcRemoteExecutor(
+      ServerCapabilities serverCapabilities,
       RemoteOptions remoteOptions,
       ReferenceCountedChannel channel,
       CallCredentialsProvider callCredentialsProvider,
       RemoteRetrier retrier) {
+    this.serverCapabilities = serverCapabilities;
     this.remoteOptions = remoteOptions;
     this.channel = channel;
     this.callCredentialsProvider = callCredentialsProvider;
@@ -318,6 +322,11 @@ static ExecuteResponse extractResponseOrThrowIfError(Operation operation) throws
     }
   }
 
+  @Override
+  public ServerCapabilities getServerCapabilities() {
+    return this.serverCapabilities;
+  }
+
   @Override
   public ExecuteResponse executeRemotely(
       RemoteActionExecutionContext context, ExecuteRequest request, OperationObserver observer)

src/main/java/com/google/devtools/build/lib/remote/GrpcRemoteExecutor.java

@@ -19,6 +19,7 @@
 import build.bazel.remote.execution.v2.ExecutionGrpc;
 import build.bazel.remote.execution.v2.ExecutionGrpc.ExecutionBlockingStub;
 import build.bazel.remote.execution.v2.RequestMetadata;
+import build.bazel.remote.execution.v2.ServerCapabilities;
 import build.bazel.remote.execution.v2.WaitExecutionRequest;
 import com.google.common.base.Preconditions;
 import com.google.devtools.build.lib.authandtls.CallCredentialsProvider;
@@ -43,16 +44,19 @@
 @ThreadSafe
 class GrpcRemoteExecutor implements RemoteExecutionClient {
 
+  private final ServerCapabilities serverCapabilities;
   private final ReferenceCountedChannel channel;
   private final CallCredentialsProvider callCredentialsProvider;
   private final RemoteRetrier retrier;
 
   private final AtomicBoolean closed = new AtomicBoolean();
 
   public GrpcRemoteExecutor(
+      ServerCapabilities serverCapabilities,
       ReferenceCountedChannel channel,
       CallCredentialsProvider callCredentialsProvider,
       RemoteRetrier retrier) {
+    this.serverCapabilities = serverCapabilities;
     this.channel = channel;
     this.callCredentialsProvider = callCredentialsProvider;
     this.retrier = retrier;
@@ -89,6 +93,11 @@ private ExecuteResponse getOperationResponse(Operation op) throws IOException {
     return null;
   }
 
+  @Override
+  public ServerCapabilities getServerCapabilities() {
+    return this.serverCapabilities;
+  }
+
   /* Execute has two components: the Execute call and (optionally) the WaitExecution call.
    * This is the simple flow without any errors:
    *

src/main/java/com/google/devtools/build/lib/remote/RemoteModule.java

@@ -489,11 +489,12 @@ public void beforeCommand(CommandEnvironment env) throws AbruptExitException {
     //
     // If they point to different endpoints, we check the endpoint with execution or cache
     // capabilities respectively.
+    ServerCapabilities executionCapabilities = null;
     ServerCapabilities cacheCapabilities = null;
     try {
       if (execChannel != null) {
         if (cacheChannel != execChannel) {
-          var unused =
+          executionCapabilities =
               getAndVerifyServerCapabilities(
                   remoteOptions,
                   execChannel,
@@ -521,6 +522,7 @@ public void beforeCommand(CommandEnvironment env) throws AbruptExitException {
                   env,
                   digestUtil,
                   ServerCapabilitiesRequirement.EXECUTION_AND_CACHE);
+          executionCapabilities = cacheCapabilities;
         }
       } else {
         cacheCapabilities =
@@ -601,7 +603,11 @@ public void beforeCommand(CommandEnvironment env) throws AbruptExitException {
                 circuitBreaker);
         remoteExecutor =
             new ExperimentalGrpcRemoteExecutor(
-                remoteOptions, execChannel.retain(), callCredentialsProvider, execRetrier);
+                executionCapabilities,
+                remoteOptions,
+                execChannel.retain(),
+                callCredentialsProvider,
+                execRetrier);
       } else {
         RemoteRetrier execRetrier =
             new RemoteRetrier(
@@ -610,7 +616,8 @@ public void beforeCommand(CommandEnvironment env) throws AbruptExitException {
                 retryScheduler,
                 circuitBreaker);
         remoteExecutor =
-            new GrpcRemoteExecutor(execChannel.retain(), callCredentialsProvider, execRetrier);
+            new GrpcRemoteExecutor(
+                executionCapabilities, execChannel.retain(), callCredentialsProvider, execRetrier);
       }
       execChannel.release();
       RemoteExecutionCache remoteCache =

src/main/java/com/google/devtools/build/lib/remote/common/RemoteExecutionClient.java

@@ -15,6 +15,7 @@
 
 import build.bazel.remote.execution.v2.ExecuteRequest;
 import build.bazel.remote.execution.v2.ExecuteResponse;
+import build.bazel.remote.execution.v2.ServerCapabilities;
 import java.io.IOException;
 
 /**
@@ -24,6 +25,9 @@
  */
 public interface RemoteExecutionClient {
 
+  /** Returns the cache capabilities of the remote execution server */
+  ServerCapabilities getServerCapabilities();
+
   /** Execute an action remotely using Remote Execution API. */
   ExecuteResponse executeRemotely(
       RemoteActionExecutionContext context, ExecuteRequest request, OperationObserver observer)

src/test/java/com/google/devtools/build/lib/remote/ExperimentalGrpcRemoteExecutorTest.java

@@ -21,8 +21,10 @@
 import build.bazel.remote.execution.v2.Digest;
 import build.bazel.remote.execution.v2.ExecuteRequest;
 import build.bazel.remote.execution.v2.ExecuteResponse;
+import build.bazel.remote.execution.v2.ExecutionCapabilities;
 import build.bazel.remote.execution.v2.OutputFile;
 import build.bazel.remote.execution.v2.RequestMetadata;
+import build.bazel.remote.execution.v2.ServerCapabilities;
 import com.google.common.util.concurrent.ListeningScheduledExecutorService;
 import com.google.common.util.concurrent.MoreExecutors;
 import com.google.devtools.build.lib.authandtls.CallCredentialsProvider;
@@ -133,9 +135,15 @@ public int maxConcurrency() {
               }
             });
 
+    ServerCapabilities caps =
+        ServerCapabilities.newBuilder()
+            .setExecutionCapabilities(
+                ExecutionCapabilities.newBuilder().setExecEnabled(true).build())
+            .build();
+
     executor =
         new ExperimentalGrpcRemoteExecutor(
-            remoteOptions, channel, CallCredentialsProvider.NO_CREDENTIALS, retrier);
+            caps, remoteOptions, channel, CallCredentialsProvider.NO_CREDENTIALS, retrier);
   }
 
   @After

src/test/java/com/google/devtools/build/lib/remote/RemoteSpawnRunnerWithGrpcRemoteExecutorTest.java

@@ -34,6 +34,7 @@
 import build.bazel.remote.execution.v2.Directory;
 import build.bazel.remote.execution.v2.ExecuteRequest;
 import build.bazel.remote.execution.v2.ExecuteResponse;
+import build.bazel.remote.execution.v2.ExecutionCapabilities;
 import build.bazel.remote.execution.v2.ExecutionGrpc.ExecutionImplBase;
 import build.bazel.remote.execution.v2.FileNode;
 import build.bazel.remote.execution.v2.FindMissingBlobsRequest;
@@ -42,6 +43,7 @@
 import build.bazel.remote.execution.v2.OutputDirectory;
 import build.bazel.remote.execution.v2.OutputFile;
 import build.bazel.remote.execution.v2.RequestMetadata;
+import build.bazel.remote.execution.v2.ServerCapabilities;
 import build.bazel.remote.execution.v2.Tree;
 import build.bazel.remote.execution.v2.WaitExecutionRequest;
 import com.google.bytestream.ByteStreamGrpc.ByteStreamImplBase;
@@ -198,13 +200,13 @@ public final void setUp() throws Exception {
             new FakeOwner("Mnemonic", "Progress Message", "//dummy:label"),
             ImmutableList.of("/bin/echo", "Hi!"),
             ImmutableMap.of("VARIABLE", "value"),
-            /*executionInfo=*/ ImmutableMap.<String, String>of(),
-            /*runfilesSupplier=*/ null,
-            /*filesetMappings=*/ ImmutableMap.of(),
-            /*inputs=*/ NestedSetBuilder.create(
+            /* executionInfo= */ ImmutableMap.<String, String>of(),
+            /* runfilesSupplier= */ null,
+            /* filesetMappings= */ ImmutableMap.of(),
+            /* inputs= */ NestedSetBuilder.create(
                 Order.STABLE_ORDER, ActionInputHelper.fromPath("input")),
-            /*tools=*/ NestedSetBuilder.emptySet(Order.STABLE_ORDER),
-            /*outputs=*/ ImmutableSet.of(
+            /* tools= */ NestedSetBuilder.emptySet(Order.STABLE_ORDER),
+            /* outputs= */ ImmutableSet.of(
                 new ActionInput() {
                   @Override
                   public String getExecPathString() {
@@ -247,7 +249,7 @@ public PathFragment getExecPath() {
                     return PathFragment.create("bar");
                   }
                 }),
-            /*mandatoryOutputs=*/ ImmutableSet.of(),
+            /* mandatoryOutputs= */ ImmutableSet.of(),
             ResourceSet.ZERO);
 
     Path stdout = fs.getPath("/tmp/stdout");
@@ -295,8 +297,14 @@ public int maxConcurrency() {
                 return 100;
               }
             });
+    ServerCapabilities caps =
+        ServerCapabilities.newBuilder()
+            .setExecutionCapabilities(
+                ExecutionCapabilities.newBuilder().setExecEnabled(true).build())
+            .build();
     GrpcRemoteExecutor executor =
-        new GrpcRemoteExecutor(channel.retain(), CallCredentialsProvider.NO_CREDENTIALS, retrier);
+        new GrpcRemoteExecutor(
+            caps, channel.retain(), CallCredentialsProvider.NO_CREDENTIALS, retrier);
     CallCredentialsProvider callCredentialsProvider =
         GoogleAuthUtils.newCallCredentialsProvider(null);
     GrpcCacheClient cacheProtocol =
@@ -326,7 +334,7 @@ public int maxConcurrency() {
             remoteOptions,
             Options.getDefaults(ExecutionOptions.class),
             /* verboseFailures= */ true,
-            /*cmdlineReporter=*/ null,
+            /* cmdlineReporter= */ null,
             retryService,
             logDir,
             remoteExecutionService);