fix: add gccl-invocation-id interceptor (#1309)

frankyn · BenWhitehead · web-flow · commit 335c2679b70f · 2022-04-01T17:32:32.000-04:00
add gccl-invocation-id to all HTTP requests to allow identifying operation attempts across multiple rpcs. Excludes signed urls.

Co-authored-by: BenWhitehead &lt;BenWhitehead@users.noreply.github.com&gt;
diff --git a/google-cloud-storage/pom.xml b/google-cloud-storage/pom.xml
@@ -92,6 +92,10 @@
       <groupId>com.fasterxml.jackson.core</groupId>
       <artifactId>jackson-core</artifactId>
     </dependency>
+    <dependency>
+      <groupId>com.google.code.findbugs</groupId>
+      <artifactId>jsr305</artifactId>
+    </dependency>
 
     <!-- Test dependencies -->
     <dependency>
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/Retrying.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/Retrying.java
@@ -22,6 +22,7 @@
 import com.google.api.gax.retrying.ResultRetryAlgorithm;
 import com.google.api.gax.retrying.RetrySettings;
 import com.google.cloud.RetryHelper.RetryHelperException;
+import com.google.cloud.storage.spi.v1.HttpRpcContext;
 import java.util.concurrent.Callable;
 import java.util.function.Function;
 
@@ -47,11 +48,15 @@ final class Retrying {
    */
   static <T, U> U run(
       StorageOptions options, ResultRetryAlgorithm<?> algorithm, Callable<T> c, Function<T, U> f) {
+    HttpRpcContext httpRpcContext = HttpRpcContext.getInstance();
     try {
+      httpRpcContext.newInvocationId();
       T result = runWithRetries(c, options.getRetrySettings(), algorithm, options.getClock());
       return result == null ? null : f.apply(result);
     } catch (RetryHelperException e) {
       throw StorageException.coalesce(e);
+    } finally {
+      httpRpcContext.clearInvocationId();
     }
   }
 }
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/StorageOptions.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/StorageOptions.java
@@ -38,8 +38,9 @@ public class StorageOptions extends ServiceOptions<Storage, StorageOptions> {
   private static final String GCS_SCOPE = "https://www.googleapis.com/auth/devstorage.full_control";
   private static final Set<String> SCOPES = ImmutableSet.of(GCS_SCOPE);
   private static final String DEFAULT_HOST = "https://storage.googleapis.com";
-
+  private static final boolean DEFAULT_INCLUDE_INVOCATION_ID = true;
   private final RetryAlgorithmManager retryAlgorithmManager;
+  private final boolean includeInvocationId;
 
   public static class DefaultStorageFactory implements StorageFactory {
 
@@ -64,11 +65,13 @@ public ServiceRpc create(StorageOptions options) {
   public static class Builder extends ServiceOptions.Builder<Storage, StorageOptions, Builder> {
 
     private StorageRetryStrategy storageRetryStrategy;
+    private boolean includeInvocationId;
 
     private Builder() {}
 
     private Builder(StorageOptions options) {
       super(options);
+      this.includeInvocationId = options.includeInvocationId;
     }
 
     @Override
@@ -93,6 +96,17 @@ public Builder setStorageRetryStrategy(StorageRetryStrategy storageRetryStrategy
       return this;
     }
 
+    /**
+     * Override default enablement of invocation id added to x-goog-api-client header.
+     *
+     * @param includeInvocationId a boolean to change enablement of invocation id
+     * @return the builder
+     */
+    public Builder setIncludeInvocationId(boolean includeInvocationId) {
+      this.includeInvocationId = includeInvocationId;
+      return this;
+    }
+
     @Override
     public StorageOptions build() {
       return new StorageOptions(this, new StorageDefaults());
@@ -105,6 +119,7 @@ private StorageOptions(Builder builder, StorageDefaults serviceDefaults) {
         new RetryAlgorithmManager(
             MoreObjects.firstNonNull(
                 builder.storageRetryStrategy, serviceDefaults.getStorageRetryStrategy()));
+    this.includeInvocationId = builder.includeInvocationId;
   }
 
   private static class StorageDefaults implements ServiceDefaults<Storage, StorageOptions> {
@@ -127,6 +142,10 @@ public TransportOptions getDefaultTransportOptions() {
     public StorageRetryStrategy getStorageRetryStrategy() {
       return StorageRetryStrategy.getDefaultStorageRetryStrategy();
     }
+
+    public boolean isIncludeInvocationId() {
+      return DEFAULT_INCLUDE_INVOCATION_ID;
+    }
   }
 
   public static HttpTransportOptions getDefaultHttpTransportOptions() {
@@ -153,6 +172,11 @@ RetryAlgorithmManager getRetryAlgorithmManager() {
     return retryAlgorithmManager;
   }
 
+  /** Returns if Invocation ID is enabled and transmitted through x-goog-api-client header. */
+  public boolean isIncludeInvocationId() {
+    return includeInvocationId;
+  }
+
   /** Returns a default {@code StorageOptions} instance. */
   public static StorageOptions getDefaultInstance() {
     return newBuilder().build();
@@ -180,6 +204,8 @@ public boolean equals(Object obj) {
   }
 
   public static Builder newBuilder() {
-    return new Builder().setHost(DEFAULT_HOST);
+    return new Builder()
+        .setHost(DEFAULT_HOST)
+        .setIncludeInvocationId(DEFAULT_INCLUDE_INVOCATION_ID);
   }
 }
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/spi/v1/HttpRpcContext.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/spi/v1/HttpRpcContext.java
@@ -0,0 +1,72 @@
+/*
+ * Copyright 2022 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.storage.spi.v1;
+
+import com.google.api.core.InternalApi;
+import java.util.UUID;
+import java.util.function.Supplier;
+import javax.annotation.Nullable;
+
+@InternalApi
+public final class HttpRpcContext {
+
+  private static final Object GET_INSTANCE_LOCK = new Object();
+
+  private static volatile HttpRpcContext instance;
+
+  private final ThreadLocal<UUID> invocationId;
+  private final Supplier<UUID> supplier;
+
+  HttpRpcContext(Supplier<UUID> randomUUID) {
+    this.invocationId = new InheritableThreadLocal<>();
+    this.supplier = randomUUID;
+  }
+
+  @InternalApi
+  @Nullable
+  public UUID getInvocationId() {
+    return invocationId.get();
+  }
+
+  @InternalApi
+  public UUID newInvocationId() {
+    invocationId.set(supplier.get());
+    return getInvocationId();
+  }
+
+  @InternalApi
+  public void clearInvocationId() {
+    invocationId.remove();
+  }
+
+  @InternalApi
+  public static HttpRpcContext init() {
+    return new HttpRpcContext(UUID::randomUUID);
+  }
+
+  @InternalApi
+  public static HttpRpcContext getInstance() {
+    if (instance == null) {
+      synchronized (GET_INSTANCE_LOCK) {
+        if (instance == null) {
+          instance = init();
+        }
+      }
+    }
+    return instance;
+  }
+}
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/spi/v1/HttpStorageRpc.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/spi/v1/HttpStorageRpc.java
@@ -18,6 +18,7 @@
 
 import static com.google.common.base.MoreObjects.firstNonNull;
 import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Preconditions.checkNotNull;
 import static java.net.HttpURLConnection.HTTP_NOT_FOUND;
 
 import com.google.api.client.googleapis.batch.BatchRequest;
@@ -26,6 +27,7 @@
 import com.google.api.client.http.ByteArrayContent;
 import com.google.api.client.http.EmptyContent;
 import com.google.api.client.http.GenericUrl;
+import com.google.api.client.http.HttpExecuteInterceptor;
 import com.google.api.client.http.HttpHeaders;
 import com.google.api.client.http.HttpRequest;
 import com.google.api.client.http.HttpRequestFactory;
@@ -86,6 +88,8 @@
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
+import java.util.UUID;
+import javax.annotation.Nullable;
 
 public class HttpStorageRpc implements StorageRpc {
   public static final String DEFAULT_PROJECTION = "full";
@@ -114,6 +118,9 @@ public HttpStorageRpc(StorageOptions options) {
     // Open Census initialization
     censusHttpModule = new CensusHttpModule(tracer, true);
     initializer = censusHttpModule.getHttpRequestInitializer(initializer);
+    if (options.isIncludeInvocationId()) {
+      initializer = new InvocationIdInitializer(initializer);
+    }
     batchRequestInitializer = censusHttpModule.getHttpRequestInitializer(null);
     storage =
         new Storage.Builder(transport, new JacksonFactory(), initializer)
@@ -122,6 +129,54 @@ public HttpStorageRpc(StorageOptions options) {
             .build();
   }
 
+  private static final class InvocationIdInitializer implements HttpRequestInitializer {
+    @Nullable HttpRequestInitializer initializer;
+
+    private InvocationIdInitializer(@Nullable HttpRequestInitializer initializer) {
+      this.initializer = initializer;
+    }
+
+    @Override
+    public void initialize(HttpRequest request) throws IOException {
+      checkNotNull(request);
+      if (this.initializer != null) {
+        this.initializer.initialize(request);
+      }
+      request.setInterceptor(new InvocationIdInterceptor(request.getInterceptor()));
+    }
+  }
+
+  private static final class InvocationIdInterceptor implements HttpExecuteInterceptor {
+    @Nullable HttpExecuteInterceptor interceptor;
+
+    private InvocationIdInterceptor(@Nullable HttpExecuteInterceptor interceptor) {
+      this.interceptor = interceptor;
+    }
+
+    @Override
+    public void intercept(HttpRequest request) throws IOException {
+      checkNotNull(request);
+      if (this.interceptor != null) {
+        this.interceptor.intercept(request);
+      }
+      UUID invocationId = HttpRpcContext.getInstance().getInvocationId();
+      final String signatureKey = "Signature="; // For V2 and V4 signedURLs
+      final String builtURL = request.getUrl().build();
+      if (invocationId != null && !builtURL.contains(signatureKey)) {
+        HttpHeaders headers = request.getHeaders();
+        String existing = (String) headers.get("x-goog-api-client");
+        String invocationEntry = "gccl-invocation-id/" + invocationId;
+        final String newValue;
+        if (existing != null && !existing.isEmpty()) {
+          newValue = existing + " " + invocationEntry;
+        } else {
+          newValue = invocationEntry;
+        }
+        headers.set("x-goog-api-client", newValue);
+      }
+    }
+  }
+
   private class DefaultRpcBatch implements RpcBatch {
 
     // Batch size is limited as, due to some current service implementation details, the service
diff --git a/google-cloud-storage/src/test/java/com/google/cloud/storage/StorageOptionsTest.java b/google-cloud-storage/src/test/java/com/google/cloud/storage/StorageOptionsTest.java
@@ -17,6 +17,8 @@
 package com.google.cloud.storage;
 
 import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
 
 import com.google.cloud.TransportOptions;
 import org.easymock.EasyMock;
@@ -65,4 +67,18 @@ public void testDefaultInstanceSpecifiesCorrectHost() {
 
     assertThat(opts1.getHost()).isEqualTo("https://storage.googleapis.com");
   }
+
+  @Test
+  public void testDefaultInvocationId() {
+    StorageOptions opts1 = StorageOptions.getDefaultInstance();
+
+    assertTrue(opts1.isIncludeInvocationId());
+  }
+
+  @Test
+  public void testDisableInvocationId() {
+    StorageOptions opts1 = StorageOptions.newBuilder().setIncludeInvocationId(false).build();
+
+    assertFalse(opts1.isIncludeInvocationId());
+  }
 }
diff --git a/google-cloud-storage/src/test/java/com/google/cloud/storage/spi/v1/HttpRpcContextTest.java b/google-cloud-storage/src/test/java/com/google/cloud/storage/spi/v1/HttpRpcContextTest.java
