Skip to content

Commit

Permalink
Make Cloud Tasks Utils canary-aware (google#2639)
Browse files Browse the repository at this point in the history
  • Loading branch information
jianglai authored Jan 14, 2025
1 parent f8407c7 commit e4ee63b
Show file tree
Hide file tree
Showing 4 changed files with 21 additions and 7 deletions.
14 changes: 9 additions & 5 deletions core/src/main/java/google/registry/batch/CloudTasksUtils.java
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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();
}

Expand All @@ -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}.
Expand Down Expand Up @@ -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}.
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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");
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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;
Expand Down
10 changes: 10 additions & 0 deletions util/src/main/java/google/registry/util/RegistryEnvironment.java
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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);
Expand Down

0 comments on commit e4ee63b

Please sign in to comment.