mirrors/nomulus
1
0
Fork 0
mirror of https://github.com/google/nomulus synced 2026-09-19 06:32:00 +00:00
Code Issues Releases Wiki Activity

Finish RDE pipeline implementation in SQL mode (#1330)

Browse Source 
This PR adds the final step in RDE pipeline (enqueueing the next action
  to Cloud Tasks) and makes some necessary changes, namely by making all
  CloudTasksUtils related classes serializable, so that they can be used
  on Beam.

<!-- Reviewable:start -->
This change is [<img src="https://reviewable.io/review_button.svg" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/google/nomulus/1330)
<!-- Reviewable:end -->
This commit is contained in:
Lai Jiang
2021-10-04 21:02:44 -04:00
committed by GitHub
parent f6920454f6
commit 62eb8801c5
27 changed files with 903 additions and 206 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/src/main/java/google/registry/beam/rde/RdeIO.java
+58 -4
View File
@@ -19,10 +19,13 @@ import static com.google.common.base.Verify.verify;
import static google.registry.model.common.Cursor.getCursorTimeOrStartOfTime;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import static google.registry.rde.RdeModule.BRDA_QUEUE;
import static google.registry.rde.RdeModule.RDE_UPLOAD_QUEUE;
import static java.nio.charset.StandardCharsets.UTF_8;
import com.google.auto.value.AutoValue;
import com.google.cloud.storage.BlobId;
import com.google.common.collect.ImmutableMultimap;
import com.google.common.flogger.FluentLogger;
import google.registry.gcs.GcsUtils;
import google.registry.keyring.api.PgpHelper;
@@ -31,14 +34,20 @@ import google.registry.model.rde.RdeMode;
import google.registry.model.rde.RdeNamingUtils;
import google.registry.model.rde.RdeRevision;
import google.registry.model.tld.Registry;
import google.registry.rde.BrdaCopyAction;
import google.registry.rde.DepositFragment;
import google.registry.rde.Ghostryde;
import google.registry.rde.PendingDeposit;
import google.registry.rde.RdeCounter;
import google.registry.rde.RdeMarshaller;
import google.registry.rde.RdeModule;
import google.registry.rde.RdeResourceType;
import google.registry.rde.RdeUploadAction;
import google.registry.rde.RdeUtil;
import google.registry.request.Action.Service;
import google.registry.request.RequestParameters;
import google.registry.tldconfig.idn.IdnTableEnum;
import google.registry.util.CloudTasksUtils;
import google.registry.xjc.rdeheader.XjcRdeHeader;
import google.registry.xjc.rdeheader.XjcRdeHeaderElement;
import google.registry.xml.ValidationMode;
@@ -68,6 +77,8 @@ public class RdeIO {
abstract GcsUtils gcsUtils();
abstract CloudTasksUtils cloudTasksUtils();
abstract String rdeBucket();
// It's OK to return a primitive array because we are only using it to construct the
@@ -83,7 +94,9 @@ public class RdeIO {
@AutoValue.Builder
abstract static class Builder {
abstract Builder setGcsUtils(GcsUtils gcsUtils);
abstract Builder setGcsUtils(GcsUtils value);
abstract Builder setCloudTasksUtils(CloudTasksUtils value);
abstract Builder setRdeBucket(String value);
@@ -100,7 +113,8 @@ public class RdeIO {
.apply(
"Write to GCS",
ParDo.of(new RdeWriter(gcsUtils(), rdeBucket(), stagingKeyBytes(), validationMode())))
.apply("Update cursors", ParDo.of(new CursorUpdater()));
.apply("Update cursors", ParDo.of(new CursorUpdater()))
.apply("Enqueue upload action", ParDo.of(new UploadEnqueuer(cloudTasksUtils())));
return PDone.in(input.getPipeline());
}
}
@@ -236,11 +250,12 @@ public class RdeIO {
}
}
private static class CursorUpdater extends DoFn<KV<PendingDeposit, Integer>, Void> {
private static class CursorUpdater extends DoFn<KV<PendingDeposit, Integer>, PendingDeposit> {
private static final FluentLogger logger = FluentLogger.forEnclosingClass();
@ProcessElement
public void processElement(@Element KV<PendingDeposit, Integer> input) {
public void processElement(
@Element KV<PendingDeposit, Integer> input, OutputReceiver<PendingDeposit> outputReceiver) {
tm().transact(
() -> {
PendingDeposit key = input.getKey();
@@ -268,6 +283,45 @@ public class RdeIO {
"Rolled forward %s on %s cursor to %s.", key.cursor(), key.tld(), newPosition);
RdeRevision.saveRevision(key.tld(), key.watermark(), key.mode(), revision);
});
outputReceiver.output(input.getKey());
}
}
private static class UploadEnqueuer extends DoFn<PendingDeposit, Void> {
private final CloudTasksUtils cloudTasksUtils;
private UploadEnqueuer(CloudTasksUtils cloudTasksUtils) {
this.cloudTasksUtils = cloudTasksUtils;
}
@ProcessElement
public void processElement(@Element PendingDeposit input, PipelineOptions options) {
if (input.mode() == RdeMode.FULL) {
cloudTasksUtils.enqueue(
RDE_UPLOAD_QUEUE,
CloudTasksUtils.createPostTask(
RdeUploadAction.PATH,
Service.BACKEND.getServiceId(),
ImmutableMultimap.of(
RequestParameters.PARAM_TLD,
input.tld(),
RdeModule.PARAM_PREFIX,
options.getJobName() + '/')));
} else {
cloudTasksUtils.enqueue(
BRDA_QUEUE,
CloudTasksUtils.createPostTask(
BrdaCopyAction.PATH,
Service.BACKEND.getServiceId(),
ImmutableMultimap.of(
RequestParameters.PARAM_TLD,
input.tld(),
RdeModule.PARAM_WATERMARK,
input.watermark().toString(),
RdeModule.PARAM_PREFIX,
options.getJobName() + '/')));
}
}
}
}
core/src/main/java/google/registry/beam/rde/RdePipeline.java
+26 -15
View File
@@ -27,6 +27,7 @@ import com.google.common.io.BaseEncoding;
import dagger.BindsInstance;
import dagger.Component;
import google.registry.beam.common.RegistryJpaIO;
import google.registry.config.CloudTasksUtilsModule;
import google.registry.config.CredentialModule;
import google.registry.config.RegistryConfig.ConfigModule;
import google.registry.gcs.GcsUtils;
@@ -44,6 +45,8 @@ import google.registry.rde.PendingDeposit;
import google.registry.rde.PendingDeposit.PendingDepositCoder;
import google.registry.rde.RdeFragmenter;
import google.registry.rde.RdeMarshaller;
import google.registry.util.CloudTasksUtils;
import google.registry.util.UtilsModule;
import google.registry.xml.ValidationMode;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
@@ -66,7 +69,6 @@ import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.transforms.FlatMapElements;
import org.apache.beam.sdk.transforms.Flatten;
import org.apache.beam.sdk.transforms.GroupByKey;
import org.apache.beam.sdk.transforms.Reshuffle;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionList;
@@ -93,6 +95,7 @@ public class RdePipeline implements Serializable {
private final String rdeBucket;
private final byte[] stagingKeyBytes;
private final GcsUtils gcsUtils;
private final CloudTasksUtils cloudTasksUtils;
// Registrars to be excluded from data escrow. Not including the sandbox-only OTE type so that
// if sneaks into production we would get an extra signal.
@@ -111,13 +114,14 @@ public class RdePipeline implements Serializable {
}
@Inject
RdePipeline(RdePipelineOptions options, GcsUtils gcsUtils) {
RdePipeline(RdePipelineOptions options, GcsUtils gcsUtils, CloudTasksUtils cloudTasksUtils) {
this.options = options;
this.mode = ValidationMode.valueOf(options.getValidationMode());
this.pendings = decodePendings(options.getPendings());
this.rdeBucket = options.getGcsBucket();
this.rdeBucket = options.getRdeStagingBucket();
this.stagingKeyBytes = BaseEncoding.base64Url().decode(options.getStagingKey());
this.gcsUtils = gcsUtils;
this.cloudTasksUtils = cloudTasksUtils;
}
PipelineResult run() {
@@ -140,10 +144,11 @@ public class RdePipeline implements Serializable {
void persistData(PCollection<KV<PendingDeposit, Iterable<DepositFragment>>> input) {
input.apply(
"Write to GCS and update cursors",
"Write to GCS, update cursors, and enqueue upload tasks",
RdeIO.Write.builder()
.setRdeBucket(rdeBucket)
.setGcsUtils(gcsUtils)
.setCloudTasksUtils(cloudTasksUtils)
.setValidationMode(mode)
.setStagingKeyBytes(stagingKeyBytes)
.build());
@@ -177,18 +182,13 @@ public class RdePipeline implements Serializable {
}));
}
@SuppressWarnings("deprecation") // Reshuffle is still recommended by Dataflow.
<T extends EppResource>
PCollection<KV<PendingDeposit, DepositFragment>> processNonRegistrarEntities(
Pipeline pipeline, Class<T> clazz) {
return createInputs(pipeline, clazz)
.apply("Marshal " + clazz.getSimpleName() + " into DepositFragment", mapToFragments(clazz))
.setCoder(KvCoder.of(PendingDepositCoder.of(), SerializableCoder.of(DepositFragment.class)))
.apply(
"Reshuffle KV<PendingDeposit, DepositFragment> of "
+ clazz.getSimpleName()
+ " to prevent fusion",
Reshuffle.of());
.setCoder(
KvCoder.of(PendingDepositCoder.of(), SerializableCoder.of(DepositFragment.class)));
}
<T extends EppResource> PCollection<VKey<T>> createInputs(Pipeline pipeline, Class<T> clazz) {
@@ -202,7 +202,7 @@ public class RdePipeline implements Serializable {
String.class,
// TODO: consider adding coders for entities and pass them directly instead of using
// VKeys.
x -> VKey.create(clazz, x)));
x -> VKey.createSql(clazz, x)));
}
<T extends EppResource>
@@ -270,7 +270,7 @@ public class RdePipeline implements Serializable {
* Encodes the TLD to pending deposit map in an URL safe string that is sent to the pipeline
* worker by the pipeline launcher as a pipeline option.
*/
static String encodePendings(ImmutableSetMultimap<String, PendingDeposit> pendings)
public static String encodePendings(ImmutableSetMultimap<String, PendingDeposit> pendings)
throws IOException {
try (ByteArrayOutputStream baos = new ByteArrayOutputStream()) {
ObjectOutputStream oos = new ObjectOutputStream(baos);
@@ -282,13 +282,24 @@ public class RdePipeline implements Serializable {
public static void main(String[] args) throws IOException, ClassNotFoundException {
PipelineOptionsFactory.register(RdePipelineOptions.class);
RdePipelineOptions options = PipelineOptionsFactory.fromArgs(args).as(RdePipelineOptions.class);
RdePipelineOptions options =
PipelineOptionsFactory.fromArgs(args).withValidation().as(RdePipelineOptions.class);
// RegistryPipelineWorkerInitializer only initializes before pipeline executions, after the
// main() function constructed the graph. We need the registry environment set up so that we
// can create a CloudTasksUtils which uses the environment-dependent config file.
options.getRegistryEnvironment().setup();
options.setIsolationOverride(TransactionIsolationLevel.TRANSACTION_READ_COMMITTED);
DaggerRdePipeline_RdePipelineComponent.builder().options(options).build().rdePipeline().run();
}
@Singleton
@Component(modules = {CredentialModule.class, ConfigModule.class})
@Component(
modules = {
CredentialModule.class,
ConfigModule.class,
CloudTasksUtilsModule.class,
UtilsModule.class
})
interface RdePipelineComponent {
RdePipeline rdePipeline();
core/src/main/java/google/registry/beam/rde/RdePipelineOptions.java
+2 -2
View File
@@ -31,9 +31,9 @@ public interface RdePipelineOptions extends RegistryPipelineOptions {
void setValidationMode(String value);
@Description("The GCS bucket where the encrypted RDE deposits will be uploaded to.")
String getGcsBucket();
String getRdeStagingBucket();
void setGcsBucket(String value);
void setRdeStagingBucket(String value);
@Description("The Base64-encoded PGP public key to encrypt the deposits.")
String getStagingKey();
core/src/main/java/google/registry/config/CloudTasksUtilsModule.java
+30 -16
View File
@@ -22,10 +22,13 @@ import dagger.Provides;
import google.registry.config.CredentialModule.DefaultCredential;
import google.registry.config.RegistryConfig.Config;
import google.registry.util.CloudTasksUtils;
import google.registry.util.CloudTasksUtils.GcpCloudTasksClient;
import google.registry.util.CloudTasksUtils.SerializableCloudTasksClient;
import google.registry.util.GoogleCredentialsBundle;
import google.registry.util.Retrier;
import java.io.IOException;
import javax.inject.Provider;
import java.io.Serializable;
import java.util.function.Supplier;
import javax.inject.Singleton;
/**
@@ -42,24 +45,35 @@ public abstract class CloudTasksUtilsModule {
public static CloudTasksUtils provideCloudTasksUtils(
@Config("projectId") String projectId,
@Config("locationId") String locationId,
// Use a provider so that we can use try-with-resources with the client, which implements
// Autocloseable.
Provider<CloudTasksClient> clientProvider,
SerializableCloudTasksClient client,
Retrier retrier) {
return new CloudTasksUtils(retrier, projectId, locationId, clientProvider);
return new CloudTasksUtils(retrier, projectId, locationId, client);
}
// Provides a supplier instead of using a Dagger @Provider because the latter is not serializable.
@Provides
public static Supplier<CloudTasksClient> provideCloudTasksClientSupplier(
@DefaultCredential GoogleCredentialsBundle credentials) {
return (Supplier<CloudTasksClient> & Serializable)
() -> {
CloudTasksClient client;
try {
client =
CloudTasksClient.create(
CloudTasksSettings.newBuilder()
.setCredentialsProvider(
FixedCredentialsProvider.create(credentials.getGoogleCredentials()))
.build());
} catch (IOException e) {
throw new RuntimeException(e);
}
return client;
};
}
@Provides
public static CloudTasksClient provideCloudTasksClient(
@DefaultCredential GoogleCredentialsBundle credentials) {
try {
return CloudTasksClient.create(
CloudTasksSettings.newBuilder()
.setCredentialsProvider(
FixedCredentialsProvider.create(credentials.getGoogleCredentials()))
.build());
} catch (IOException e) {
throw new RuntimeException(e);
}
public static SerializableCloudTasksClient provideSerializableCloudTasksClient(
final Supplier<CloudTasksClient> clientSupplier) {
return new GcpCloudTasksClient(clientSupplier);
}
}
core/src/main/java/google/registry/rde/BrdaCopyAction.java
+15 -11
View File
@@ -31,6 +31,7 @@ import google.registry.request.auth.Auth;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.Optional;
import javax.inject.Inject;
import org.bouncycastle.openpgp.PGPKeyPair;
import org.bouncycastle.openpgp.PGPPrivateKey;
@@ -60,7 +61,7 @@ import org.joda.time.DateTime;
auth = Auth.AUTH_INTERNAL_OR_ADMIN)
public final class BrdaCopyAction implements Runnable {
static final String PATH = "/_dr/task/brdaCopy";
public static final String PATH = "/_dr/task/brdaCopy";
private static final FluentLogger logger = FluentLogger.forEnclosingClass();
@@ -69,6 +70,7 @@ public final class BrdaCopyAction implements Runnable {
@Inject @Config("rdeBucket") String stagingBucket;
@Inject @Parameter(RequestParameters.PARAM_TLD) String tld;
@Inject @Parameter(RdeModule.PARAM_WATERMARK) DateTime watermark;
@Inject @Parameter(RdeModule.PARAM_PREFIX) Optional<String> prefix;
@Inject @Key("brdaReceiverKey") PGPPublicKey receiverKey;
@Inject @Key("brdaSigningKey") PGPKeyPair signingKey;
@Inject @Key("rdeStagingDecryptionKey") PGPPrivateKey stagingDecryptionKey;
@@ -84,11 +86,12 @@ public final class BrdaCopyAction implements Runnable {
}
private void copyAsRyde() throws IOException {
String prefix = RdeNamingUtils.makeRydeFilename(tld, watermark, THIN, 1, 0);
BlobId xmlFilename = BlobId.of(stagingBucket, prefix + ".xml.ghostryde");
BlobId xmlLengthFilename = BlobId.of(stagingBucket, prefix + ".xml.length");
BlobId rydeFile = BlobId.of(brdaBucket, prefix + ".ryde");
BlobId sigFile = BlobId.of(brdaBucket, prefix + ".sig");
String nameWithoutPrefix = RdeNamingUtils.makeRydeFilename(tld, watermark, THIN, 1, 0);
String name = prefix.orElse("") + nameWithoutPrefix;
BlobId xmlFilename = BlobId.of(stagingBucket, name + ".xml.ghostryde");
BlobId xmlLengthFilename = BlobId.of(stagingBucket, name + ".xml.length");
BlobId rydeFile = BlobId.of(brdaBucket, nameWithoutPrefix + ".ryde");
BlobId sigFile = BlobId.of(brdaBucket, nameWithoutPrefix + ".sig");
long xmlLength = readXmlLength(xmlLengthFilename);
@@ -97,11 +100,12 @@ public final class BrdaCopyAction implements Runnable {
InputStream ghostrydeDecoder = Ghostryde.decoder(gcsInput, stagingDecryptionKey);
OutputStream rydeOut = gcsUtils.openOutputStream(rydeFile);
OutputStream sigOut = gcsUtils.openOutputStream(sigFile);
RydeEncoder rydeEncoder = new RydeEncoder.Builder()
.setRydeOutput(rydeOut, receiverKey)
.setSignatureOutput(sigOut, signingKey)
.setFileMetadata(prefix, xmlLength, watermark)
.build()) {
RydeEncoder rydeEncoder =
new RydeEncoder.Builder()
.setRydeOutput(rydeOut, receiverKey)
.setSignatureOutput(sigOut, signingKey)
.setFileMetadata(nameWithoutPrefix, xmlLength, watermark)
.build()) {
ByteStreams.copy(ghostrydeDecoder, rydeEncoder);
}
}
core/src/main/java/google/registry/rde/RdeModule.java
+8 -3
View File
@@ -49,6 +49,10 @@ public abstract class RdeModule {
public static final String PARAM_MODE = "mode";
public static final String PARAM_REVISION = "revision";
public static final String PARAM_LENIENT = "lenient";
public static final String PARAM_PREFIX = "prefix";
public static final String RDE_UPLOAD_QUEUE = "rde-upload";
public static final String RDE_REPORT_QUEUE = "rde-report";
public static final String BRDA_QUEUE = "brda";
@Provides
@Parameter(PARAM_WATERMARK)
@@ -92,10 +96,11 @@ public abstract class RdeModule {
return extractBooleanParameter(req, PARAM_LENIENT);
}
// TODO (jianglai): Make it a required parameter once we migrate to Cloud SQL.
@Provides
@Named("brda")
static Queue provideQueueBrda() {
return getQueue("brda");
@Parameter(PARAM_PREFIX)
static Optional<String> providePrefix(HttpServletRequest req) {
return extractOptionalParameter(req, PARAM_PREFIX);
}
@Provides
core/src/main/java/google/registry/rde/RdeReportAction.java
+4 -2
View File
@@ -68,6 +68,7 @@ public final class RdeReportAction implements Runnable, EscrowTask {
@Inject Response response;
@Inject RdeReporter reporter;
@Inject @Parameter(RequestParameters.PARAM_TLD) String tld;
@Inject @Parameter(RdeModule.PARAM_PREFIX) Optional<String> prefix;
@Inject @Config("rdeBucket") String bucket;
@Inject @Config("rdeInterval") Duration interval;
@Inject @Config("rdeReportLockTimeout") Duration timeout;
@@ -96,8 +97,9 @@ public final class RdeReportAction implements Runnable, EscrowTask {
RdeRevision.getCurrentRevision(tld, watermark, FULL)
.orElseThrow(
() -> new IllegalStateException("RdeRevision was not set on generated deposit"));
String prefix = RdeNamingUtils.makeRydeFilename(tld, watermark, FULL, 1, revision);
BlobId reportFilename = BlobId.of(bucket, prefix + "-report.xml.ghostryde");
String name =
prefix.orElse("") + RdeNamingUtils.makeRydeFilename(tld, watermark, FULL, 1, revision);
BlobId reportFilename = BlobId.of(bucket, name + "-report.xml.ghostryde");
verify(gcsUtils.existsAndNotEmpty(reportFilename), "Missing file: %s", reportFilename);
reporter.send(readReportFromGcs(reportFilename));
response.setContentType(PLAIN_TEXT_UTF_8);
core/src/main/java/google/registry/rde/RdeStagingAction.java
+75 -16
View File
@@ -14,20 +14,33 @@
package google.registry.rde;
import static google.registry.beam.BeamUtils.createJobName;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.request.Action.Method.GET;
import static google.registry.request.Action.Method.POST;
import static google.registry.xml.ValidationMode.LENIENT;
import static google.registry.xml.ValidationMode.STRICT;
import static javax.servlet.http.HttpServletResponse.SC_INTERNAL_SERVER_ERROR;
import static javax.servlet.http.HttpServletResponse.SC_NO_CONTENT;
import static javax.servlet.http.HttpServletResponse.SC_OK;
import com.google.api.services.dataflow.Dataflow;
import com.google.api.services.dataflow.model.LaunchFlexTemplateParameter;
import com.google.api.services.dataflow.model.LaunchFlexTemplateRequest;
import com.google.api.services.dataflow.model.LaunchFlexTemplateResponse;
import com.google.common.base.Ascii;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.ImmutableSetMultimap;
import com.google.common.collect.Multimaps;
import com.google.common.flogger.FluentLogger;
import com.google.common.io.BaseEncoding;
import google.registry.beam.rde.RdePipeline;
import google.registry.config.RegistryConfig.Config;
import google.registry.config.RegistryEnvironment;
import google.registry.gcs.GcsUtils;
import google.registry.keyring.api.KeyModule.Key;
import google.registry.mapreduce.MapreduceRunner;
import google.registry.mapreduce.inputs.EppResourceInputs;
import google.registry.mapreduce.inputs.NullInput;
@@ -47,6 +60,7 @@ import google.registry.request.Response;
import google.registry.request.auth.Auth;
import google.registry.util.Clock;
import google.registry.xml.ValidationMode;
import java.io.IOException;
import java.util.Optional;
import javax.inject.Inject;
import org.joda.time.DateTime;
@@ -201,6 +215,7 @@ public final class RdeStagingAction implements Runnable {
public static final String PATH = "/_dr/task/rdeStaging";
private static final String PIPELINE_NAME = "rde_pipeline";
private static final FluentLogger logger = FluentLogger.forEnclosingClass();
@Inject Clock clock;
@@ -209,7 +224,11 @@ public final class RdeStagingAction implements Runnable {
@Inject Response response;
@Inject GcsUtils gcsUtils;
@Inject MapreduceRunner mrRunner;
@Inject @Config("projectId") String projectId;
@Inject @Config("defaultJobRegion") String jobRegion;
@Inject @Config("transactionCooldown") Duration transactionCooldown;
@Inject @Config("beamStagingBucketUrl") String stagingBucketUrl;
@Inject @Config("rdeBucket") String rdeBucket;
@Inject @Parameter(RdeModule.PARAM_MANUAL) boolean manual;
@Inject @Parameter(RdeModule.PARAM_DIRECTORY) Optional<String> directory;
@Inject @Parameter(RdeModule.PARAM_MODE) ImmutableSet<String> modeStrings;
@@ -217,7 +236,8 @@ public final class RdeStagingAction implements Runnable {
@Inject @Parameter(RdeModule.PARAM_WATERMARKS) ImmutableSet<DateTime> watermarks;
@Inject @Parameter(RdeModule.PARAM_REVISION) Optional<Integer> revision;
@Inject @Parameter(RdeModule.PARAM_LENIENT) boolean lenient;
@Inject @Key("rdeStagingEncryptionKey") byte[] stagingKeyBytes;
@Inject Dataflow dataflow;
@Inject RdeStagingAction() {}
@Override
@@ -228,27 +248,66 @@ public final class RdeStagingAction implements Runnable {
String message = "Nothing needs to be deposited.";
logger.atInfo().log(message);
response.setStatus(SC_NO_CONTENT);
response.setPayload(message);
// No need to set payload as HTTP 204 response status code does not allow a payload.
return;
}
for (PendingDeposit pending : pendings.values()) {
logger.atInfo().log("Pending deposit: %s", pending);
}
ValidationMode validationMode = lenient ? LENIENT : STRICT;
RdeStagingMapper mapper = new RdeStagingMapper(validationMode, pendings);
RdeStagingReducer reducer = reducerFactory.create(validationMode, gcsUtils);
mrRunner
.setJobName("Stage escrow deposits for all TLDs")
.setModuleName("backend")
.setDefaultReduceShards(pendings.size())
.runMapreduce(
mapper,
reducer,
ImmutableList.of(
// Add an extra shard that maps over a null resource. See the mapper code for why.
new NullInput<>(), EppResourceInputs.createEntityInput(EppResource.class)))
.sendLinkToMapreduceConsole(response);
if (tm().isOfy()) {
RdeStagingMapper mapper = new RdeStagingMapper(validationMode, pendings);
RdeStagingReducer reducer = reducerFactory.create(validationMode, gcsUtils);
mrRunner
.setJobName("Stage escrow deposits for all TLDs")
.setModuleName("backend")
.setDefaultReduceShards(pendings.size())
.runMapreduce(
mapper,
reducer,
ImmutableList.of(
// Add an extra shard that maps over a null resource. See the mapper code for why.
new NullInput<>(), EppResourceInputs.createEntityInput(EppResource.class)))
.sendLinkToMapreduceConsole(response);
} else {
try {
LaunchFlexTemplateParameter parameter =
new LaunchFlexTemplateParameter()
.setJobName(createJobName("rde", clock))
.setContainerSpecGcsPath(
String.format("%s/%s_metadata.json", stagingBucketUrl, PIPELINE_NAME))
.setParameters(
ImmutableMap.of(
"pendings",
RdePipeline.encodePendings(pendings),
"validationMode",
validationMode.name(),
"rdeStagingBucket",
rdeBucket,
"stagingKey",
BaseEncoding.base64Url().omitPadding().encode(stagingKeyBytes),
"registryEnvironment",
RegistryEnvironment.get().name()));
LaunchFlexTemplateResponse launchResponse =
dataflow
.projects()
.locations()
.flexTemplates()
.launch(
projectId,
jobRegion,
new LaunchFlexTemplateRequest().setLaunchParameter(parameter))
.execute();
logger.atInfo().log("Got response: %s", launchResponse.getJob().toPrettyString());
response.setStatus(SC_OK);
response.setPayload(
String.format("Launched RDE pipeline: %s", launchResponse.getJob().getId()));
} catch (IOException e) {
logger.atWarning().withCause(e).log("Pipeline Launch failed");
response.setStatus(SC_INTERNAL_SERVER_ERROR);
response.setPayload(String.format("Pipeline launch failed: %s", e.getMessage()));
}
}
}
private ImmutableSetMultimap<String, PendingDeposit> getStandardPendingDeposits() {
core/src/main/java/google/registry/rde/RdeUploadAction.java
+29 -17
View File
@@ -14,7 +14,6 @@
package google.registry.rde;
import static com.google.appengine.api.taskqueue.TaskOptions.Builder.withUrl;
import static com.google.common.base.Verify.verify;
import static com.google.common.net.MediaType.PLAIN_TEXT_UTF_8;
import static com.jcraft.jsch.ChannelSftp.OVERWRITE;
@@ -24,14 +23,15 @@ import static google.registry.model.common.Cursor.getCursorTimeOrStartOfTime;
import static google.registry.model.rde.RdeMode.FULL;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import static google.registry.rde.RdeModule.RDE_REPORT_QUEUE;
import static google.registry.request.Action.Method.POST;
import static google.registry.util.DateTimeUtils.START_OF_TIME;
import static google.registry.util.DateTimeUtils.isBeforeOrAt;
import static java.util.Arrays.asList;
import com.google.appengine.api.taskqueue.Queue;
import com.google.cloud.storage.BlobId;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.HashMultimap;
import com.google.common.flogger.FluentLogger;
import com.google.common.io.ByteStreams;
import com.jcraft.jsch.JSch;
@@ -49,14 +49,15 @@ import google.registry.model.tld.Registry;
import google.registry.rde.EscrowTaskRunner.EscrowTask;
import google.registry.rde.JSchSshSession.JSchSshSessionFactory;
import google.registry.request.Action;
import google.registry.request.Action.Service;
import google.registry.request.HttpException.NoContentException;
import google.registry.request.Parameter;
import google.registry.request.RequestParameters;
import google.registry.request.Response;
import google.registry.request.auth.Auth;
import google.registry.util.Clock;
import google.registry.util.CloudTasksUtils;
import google.registry.util.Retrier;
import google.registry.util.TaskQueueUtils;
import google.registry.util.TeeOutputStream;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
@@ -66,7 +67,6 @@ import java.io.OutputStream;
import java.net.URI;
import java.util.Optional;
import javax.inject.Inject;
import javax.inject.Named;
import org.bouncycastle.openpgp.PGPKeyPair;
import org.bouncycastle.openpgp.PGPPrivateKey;
import org.bouncycastle.openpgp.PGPPublicKey;
@@ -90,7 +90,7 @@ import org.joda.time.Duration;
auth = Auth.AUTH_INTERNAL_OR_ADMIN)
public final class RdeUploadAction implements Runnable, EscrowTask {
static final String PATH = "/_dr/task/rdeUpload";
public static final String PATH = "/_dr/task/rdeUpload";
private static final FluentLogger logger = FluentLogger.forEnclosingClass();
@@ -109,9 +109,10 @@ public final class RdeUploadAction implements Runnable, EscrowTask {
@Inject JSchSshSessionFactory jschSshSessionFactory;
@Inject Response response;
@Inject SftpProgressMonitor sftpProgressMonitor;
@Inject TaskQueueUtils taskQueueUtils;
@Inject CloudTasksUtils cloudTasksUtils;
@Inject Retrier retrier;
@Inject @Parameter(RequestParameters.PARAM_TLD) String tld;
@Inject @Parameter(RdeModule.PARAM_PREFIX) Optional<String> prefix;
@Inject @Config("rdeBucket") String bucket;
@Inject @Config("rdeInterval") Duration interval;
@Inject @Config("rdeUploadLockTimeout") Duration timeout;
@@ -120,15 +121,21 @@ public final class RdeUploadAction implements Runnable, EscrowTask {
@Inject @Key("rdeReceiverKey") PGPPublicKey receiverKey;
@Inject @Key("rdeSigningKey") PGPKeyPair signingKey;
@Inject @Key("rdeStagingDecryptionKey") PGPPrivateKey stagingDecryptionKey;
@Inject @Named("rde-report") Queue reportQueue;
@Inject RdeUploadAction() {}
@Override
public void run() {
logger.atInfo().log("Attempting to acquire RDE upload lock for TLD '%s'.", tld);
runner.lockRunAndRollForward(this, Registry.get(tld), timeout, CursorType.RDE_UPLOAD, interval);
taskQueueUtils.enqueue(
reportQueue, withUrl(RdeReportAction.PATH).param(RequestParameters.PARAM_TLD, tld));
HashMultimap<String, String> params = HashMultimap.create();
params.put(RequestParameters.PARAM_TLD, tld);
if (prefix.isPresent()) {
params.put(RdeModule.PARAM_PREFIX, prefix.get());
}
cloudTasksUtils.enqueue(
RDE_REPORT_QUEUE,
CloudTasksUtils.createPostTask(
RdeReportAction.PATH, Service.BACKEND.getServiceId(), params));
}
@Override
@@ -164,7 +171,9 @@ public final class RdeUploadAction implements Runnable, EscrowTask {
RdeRevision.getCurrentRevision(tld, watermark, FULL)
.orElseThrow(
() -> new IllegalStateException("RdeRevision was not set on generated deposit"));
final String name = RdeNamingUtils.makeRydeFilename(tld, watermark, FULL, 1, revision);
final String nameWithoutPrefix =
RdeNamingUtils.makeRydeFilename(tld, watermark, FULL, 1, revision);
final String name = prefix.orElse("") + nameWithoutPrefix;
final BlobId xmlFilename = BlobId.of(bucket, name + ".xml.ghostryde");
final BlobId xmlLengthFilename = BlobId.of(bucket, name + ".xml.length");
BlobId reportFilename = BlobId.of(bucket, name + "-report.xml.ghostryde");
@@ -174,7 +183,8 @@ public final class RdeUploadAction implements Runnable, EscrowTask {
logger.atInfo().log("Commencing RDE upload for TLD '%s' to '%s'.", tld, uploadUrl);
final long xmlLength = readXmlLength(xmlLengthFilename);
retrier.callWithRetry(
() -> upload(xmlFilename, xmlLength, watermark, name), JSchException.class);
() -> upload(xmlFilename, xmlLength, watermark, name, nameWithoutPrefix),
JSchException.class);
logger.atInfo().log(
"Updating RDE cursor '%s' for TLD '%s' following successful upload.", RDE_UPLOAD_SFTP, tld);
tm().transact(
@@ -210,7 +220,8 @@ public final class RdeUploadAction implements Runnable, EscrowTask {
* }</pre>
*/
@VisibleForTesting
protected void upload(BlobId xmlFile, long xmlLength, DateTime watermark, String name)
protected void upload(
BlobId xmlFile, long xmlLength, DateTime watermark, String name, String nameWithoutPrefix)
throws Exception {
logger.atInfo().log("Uploading XML file '%s' to remote path '%s'.", xmlFile, uploadUrl);
try (InputStream gcsInput = gcsUtils.openInputStream(xmlFile);
@@ -218,8 +229,8 @@ public final class RdeUploadAction implements Runnable, EscrowTask {
try (JSchSshSession session = jschSshSessionFactory.create(lazyJsch.get(), uploadUrl);
JSchSftpChannel ftpChan = session.openSftpChannel()) {
ByteArrayOutputStream sigOut = new ByteArrayOutputStream();
String rydeFilename = name + ".ryde";
BlobId rydeGcsFilename = BlobId.of(bucket, rydeFilename);
String rydeFilename = nameWithoutPrefix + ".ryde";
BlobId rydeGcsFilename = BlobId.of(bucket, name + ".ryde");
try (OutputStream ftpOutput =
ftpChan.get().put(rydeFilename, sftpProgressMonitor, OVERWRITE);
OutputStream gcsOutput = gcsUtils.openOutputStream(rydeGcsFilename);
@@ -228,14 +239,15 @@ public final class RdeUploadAction implements Runnable, EscrowTask {
new RydeEncoder.Builder()
.setRydeOutput(teeOutput, receiverKey)
.setSignatureOutput(sigOut, signingKey)
.setFileMetadata(name, xmlLength, watermark)
.setFileMetadata(nameWithoutPrefix, xmlLength, watermark)
.build()) {
long bytesCopied = ByteStreams.copy(ghostrydeDecoder, rydeEncoder);
logger.atInfo().log("Uploaded %,d bytes to path '%s'.", bytesCopied, rydeFilename);
}
String sigFilename = name + ".sig";
String sigFilename = nameWithoutPrefix + ".sig";
BlobId sigGcsFilename = BlobId.of(bucket, name + ".sig");
byte[] signature = sigOut.toByteArray();
gcsUtils.createFromBytes(BlobId.of(bucket, sigFilename), signature);
gcsUtils.createFromBytes(sigGcsFilename, signature);
ftpChan.get().put(new ByteArrayInputStream(signature), sigFilename);
logger.atInfo().log("Uploaded %,d bytes to path '%s'.", signature.length, sigFilename);
}
core/src/main/java/google/registry/reporting/billing/GenerateInvoicesAction.java
+4 -1
View File
@@ -30,6 +30,7 @@ import com.google.common.collect.ImmutableMap;
import com.google.common.flogger.FluentLogger;
import com.google.common.net.MediaType;
import google.registry.config.RegistryConfig.Config;
import google.registry.config.RegistryEnvironment;
import google.registry.model.common.DatabaseMigrationStateSchedule.PrimaryDatabase;
import google.registry.reporting.ReportingModule;
import google.registry.request.Action;
@@ -127,7 +128,9 @@ public class GenerateInvoicesAction implements Runnable {
"database",
database.name(),
"billingBucketUrl",
billingBucketUrl));
billingBucketUrl,
"registryEnvironment",
RegistryEnvironment.get().name()));
LaunchFlexTemplateResponse launchResponse =
dataflow
.projects()
core/src/main/resources/google/registry/beam/invoicing_pipeline_metadata.json
+2 -2
View File
@@ -6,9 +6,9 @@
"name": "registryEnvironment",
"label": "The Registry environment.",
"helpText": "The Registry environment, required if environment-specific initialization (such as JPA) is needed on worker VMs.",
"is_optional": true,
"is_optional": false,
"regexes": [
"^[0-9A-Z_]+$"
"^PRODUCTION|SANDBOX|CRASH|QA|ALPHA$"
]
},
{
core/src/main/resources/google/registry/beam/rde_pipeline_metadata.json
+14 -5
View File
@@ -2,12 +2,21 @@
"name": "RDE/BRDA Deposit Generation",
"description": "An Apache Beam pipeline generates RDE or BRDA deposits and deposits them to GCS with GhostRyde encryption.",
"parameters": [
{
"name": "registryEnvironment",
"label": "The Registry environment.",
"helpText": "The Registry environment, required if environment-specific initialization (such as JPA) is needed on worker VMs.",
"is_optional": false,
"regexes": [
"^PRODUCTION|SANDBOX|CRASH|QA|ALPHA$"
]
},
{
"name": "pendings",
"label": "The pendings deposits to generate.",
"helpText": "A TLD to PendingDeposit map that is serialized and Base64 URL-safe encoded.",
"regexes": [
"A-Za-z0-9\\-_"
"[A-Za-z0-9\\-_]+"
]
},
{
@@ -19,12 +28,12 @@
]
},
{
"name": "gcsBucket",
"label": "The GCS bucket that where the resulting files will be stored.",
"name": "rdeStagingBucket",
"label": "The GCS bucket that where the resulting files will be stored.",
"helpText": "Only the bucket name itself, without the leading \"gs://\".",
"is_optional": false,
"regexes": [
"^[a-zA-Z0-9_\\-]+$"
"[a-zA-Z0-9_\\-]+$"
]
},
{
@@ -32,7 +41,7 @@
"label": "The PGP public key used to encrypt the RDE/BRDA deposit files.",
"helpText": "The key is Base64 URL-safe encoded.",
"regexes": [
"A-Za-z0-9\\-_"
"[A-Za-z0-9\\-_]+"
]
}
]
core/src/main/resources/google/registry/beam/spec11_pipeline_metadata.json
+2 -2
View File
@@ -6,9 +6,9 @@
"name": "registryEnvironment",
"label": "The Registry environment.",
"helpText": "The Registry environment, required if environment-specific initialization (such as JPA) is needed on worker VMs.",
"is_optional": true,
"is_optional": false,
"regexes": [
"^[0-9A-Z_]+$"
"^PRODUCTION|SANDBOX|CRASH|QA|ALPHA$"
]
},
{