Make Cloud Tasks Utils canary-aware (#2639)

Author: jianglai

core/src/main/java/google/registry/batch/CloudTasksUtils.java

@@ -17,6 +17,7 @@
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.collect.ImmutableList.toImmutableList;
+import static google.registry.config.RegistryConfig.CANARY_HEADER;
 import static java.util.concurrent.TimeUnit.SECONDS;
 
 import com.google.api.gax.rpc.ApiException;
@@ -190,6 +191,9 @@ protected Task createTask(
     requestBuilder.setOidcToken(oidcTokenBuilder.build());
     String totalPath = String.format("%s%s", service.getServiceUrl(), path);
     requestBuilder.setUrl(totalPath);
+    if (RegistryEnvironment.isCanary()) {
+      requestBuilder.putHeaders(CANARY_HEADER, "true");
+    }
     return Task.newBuilder().setHttpRequest(requestBuilder.build()).build();
   }
 
@@ -200,7 +204,7 @@ protected Task createTask(
    * default service account as the principal. That account must have permission to submit tasks to
    * Cloud Tasks.
    *
-   * <p>Prefer this overload over the one where the path and service are explicit defined, as this
+   * <p>Prefer this overload over the one where the path and service are explicitly defined, as this
    * class will automatically determine the service to use based on the action and the runtime.
    *
    * @param actionClazz the action class to run, must be annotated with {@link Action}.
@@ -269,7 +273,7 @@ public Task createTaskWithJitter(
   /**
    * Create a {@link Task} to be enqueued with a random delay up to {@code jitterSeconds}.
    *
-   * <p>Prefer this overload over the one where the path and service are explicit defined, as this
+   * <p>Prefer this overload over the one where the path and service are explicitly defined, as this
    * class will automatically determine the service to use based on the action and the runtime.
    *
    * @param actionClazz the action class to run, must be annotated with {@link Action}.
@@ -306,7 +310,7 @@ public Task createTaskWithJitter(
    * @param service the GAE/GKE service to route the request to.
    * @param params a multimap of URL query parameters. Duplicate keys are saved as is, and it is up
    *     to the server to process the duplicate keys.
-   * @param delay the amount of time that a task needs to delayed for.
+   * @param delay the amount of time that a task needs to be delayed for.
    * @return the enqueued task.
    * @see <a
    *     href=ttps://cloud.google.com/appengine/docs/standard/java/taskqueue/push/creating-tasks#target>Specifyinig
@@ -330,14 +334,14 @@ private Task createTaskWithDelay(
   /**
    * Create a {@link Task} to be enqueued with delay of {@code duration}.
    *
-   * <p>Prefer this overload over the one where the path and service are explicit defined, as this
+   * <p>Prefer this overload over the one where the path and service are explicitly defined, as this
    * class will automatically determine the service to use based on the action and the runtime.
    *
    * @param actionClazz the action class to run, must be annotated with {@link Action}.
    * @param method the HTTP method to be used for the request.
    * @param params a multimap of URL query parameters. Duplicate keys are saved as is, and it is up
    *     to the server to process the duplicate keys.
-   * @param delay the amount of time that a task needs to delayed for.
+   * @param delay the amount of time that a task needs to be delayed for.
    * @return the enqueued task.
    * @see <a
    *     href=ttps://cloud.google.com/appengine/docs/standard/java/taskqueue/push/creating-tasks#target>Specifyinig

core/src/main/java/google/registry/config/RegistryConfig.java

@@ -71,6 +71,7 @@
  */
 public final class RegistryConfig {
 
+  public static final String CANARY_HEADER = "canary";
   private static final String ENVIRONMENT_CONFIG_FORMAT = "files/nomulus-config-%s.yaml";
   private static final String YAML_CONFIG_PROD =
       readResourceUtf8(RegistryConfig.class, "files/default-config.yaml");

core/src/main/java/google/registry/tools/ServiceConnection.java

@@ -20,6 +20,7 @@
 import static com.google.common.net.HttpHeaders.X_REQUESTED_WITH;
 import static com.google.common.net.MediaType.JSON_UTF_8;
 import static google.registry.config.ConfigUtils.makeUrl;
+import static google.registry.config.RegistryConfig.CANARY_HEADER;
 import static google.registry.security.JsonHttp.JSON_SAFETY_PREFIX;
 import static java.nio.charset.StandardCharsets.UTF_8;
 
@@ -58,8 +59,6 @@ public class ServiceConnection {
   /** Pattern to heuristically extract title tag contents in HTML responses. */
   protected static final Pattern HTML_TITLE_TAG_PATTERN = Pattern.compile("<title>(.*?)</title>");
 
-  private static final String CANARY_HEADER = "canary";
-
   private final Service service;
   private final boolean useCanary;
   private final HttpRequestFactory requestFactory;

util/src/main/java/google/registry/util/RegistryEnvironment.java

@@ -58,9 +58,15 @@ public enum RegistryEnvironment {
    */
   private static final String JETTY_PROPERTY = "google.registry.jetty";
 
+  /** Name of the environmental variable of the container name. */
+  private static final String CONTAINER_ENV = "CONTAINER_NAME";
+
   private static final boolean ON_JETTY =
       Boolean.parseBoolean(System.getProperty(JETTY_PROPERTY, "false"));
 
+  private static final boolean IS_CANARY =
+      System.getenv().getOrDefault(CONTAINER_ENV, "").endsWith("-canary");
+
   /**
    * A thread local boolean that can be set in tests to indicate some code is running in a local
    * test server.
@@ -98,6 +104,10 @@ public static boolean isOnJetty() {
     return ON_JETTY;
   }
 
+  public static boolean isCanary() {
+    return IS_CANARY;
+  }
+
   public static void setIsInTestDriver(boolean value) {
     checkState(RegistryEnvironment.get() == RegistryEnvironment.UNITTEST);
     IN_TEST_SERVER.set(value);