Break out EppResourceInput subclasses

This frankenclass was starting to get a little cluttered, especially once I add the ChildEntityInput/Reader.
-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=117266882

java/com/google/domain/registry/export/BUILD

@@ -26,6 +26,7 @@ java_library(
         "//java/com/google/domain/registry/gcs",
         "//java/com/google/domain/registry/groups",
         "//java/com/google/domain/registry/mapreduce",
+        "//java/com/google/domain/registry/mapreduce/inputs",
         "//java/com/google/domain/registry/model",
         "//java/com/google/domain/registry/request",
         "//java/com/google/domain/registry/storage/drive",

java/com/google/domain/registry/export/ExportDomainListsAction.java

@@ -15,7 +15,7 @@
 package com.google.domain.registry.export;
 
 import static com.google.appengine.tools.cloudstorage.GcsServiceFactory.createGcsService;
-import static com.google.domain.registry.mapreduce.EppResourceInputs.createEntityInput;
+import static com.google.domain.registry.mapreduce.inputs.EppResourceInputs.createEntityInput;
 import static com.google.domain.registry.model.EppResourceUtils.isActive;
 import static com.google.domain.registry.model.registry.Registries.getTldsOfType;
 import static com.google.domain.registry.util.PipelineUtils.createJobPath;

java/com/google/domain/registry/flows/BUILD

@@ -26,6 +26,7 @@ java_library(
         "//java/com/google/domain/registry/config",
         "//java/com/google/domain/registry/dns",
         "//java/com/google/domain/registry/mapreduce",
+        "//java/com/google/domain/registry/mapreduce/inputs",
         "//java/com/google/domain/registry/model",
         "//java/com/google/domain/registry/monitoring/whitebox",
         "//java/com/google/domain/registry/request",

java/com/google/domain/registry/flows/async/DeleteEppResourceAction.java

@@ -29,10 +29,10 @@ import com.google.appengine.tools.mapreduce.Reducer;
 import com.google.appengine.tools.mapreduce.ReducerInput;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterators;
-import com.google.domain.registry.mapreduce.EppResourceInputs;
 import com.google.domain.registry.mapreduce.MapreduceAction;
 import com.google.domain.registry.mapreduce.MapreduceRunner;
-import com.google.domain.registry.mapreduce.NullInput;
+import com.google.domain.registry.mapreduce.inputs.EppResourceInputs;
+import com.google.domain.registry.mapreduce.inputs.NullInput;
 import com.google.domain.registry.model.EppResource;
 import com.google.domain.registry.model.annotations.ExternalMessagingName;
 import com.google.domain.registry.model.domain.DomainBase;

java/com/google/domain/registry/flows/async/DnsRefreshForHostRenameAction.java

@@ -22,9 +22,9 @@ import static com.google.domain.registry.util.PreconditionsUtils.checkArgumentNo
 import com.google.appengine.tools.mapreduce.Mapper;
 import com.google.common.collect.ImmutableList;
 import com.google.domain.registry.dns.DnsQueue;
-import com.google.domain.registry.mapreduce.EppResourceInputs;
 import com.google.domain.registry.mapreduce.MapreduceAction;
 import com.google.domain.registry.mapreduce.MapreduceRunner;
+import com.google.domain.registry.mapreduce.inputs.EppResourceInputs;
 import com.google.domain.registry.model.domain.DomainResource;
 import com.google.domain.registry.model.domain.ReferenceUnion;
 import com.google.domain.registry.model.host.HostResource;

java/com/google/domain/registry/mapreduce/BUILD

@@ -11,7 +11,7 @@ java_library(
         "//java/com/google/common/annotations",
         "//java/com/google/common/base",
         "//java/com/google/common/collect",
-        "//java/com/google/domain/registry/config",
+        "//java/com/google/domain/registry/mapreduce/inputs",
         "//java/com/google/domain/registry/model",
         "//java/com/google/domain/registry/request",
         "//java/com/google/domain/registry/util",

java/com/google/domain/registry/mapreduce/MapreduceRunner.java

@@ -34,6 +34,7 @@ import com.google.appengine.tools.pipeline.Job0;
 import com.google.appengine.tools.pipeline.JobSetting;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Optional;
+import com.google.domain.registry.mapreduce.inputs.ConcatenatingInput;
 import com.google.domain.registry.request.Parameter;
 import com.google.domain.registry.util.FormattingLogger;
 import com.google.domain.registry.util.PipelineUtils;

java/com/google/domain/registry/mapreduce/inputs/BUILD

@@ -0,0 +1,25 @@
+package(
+    default_visibility = ["//java/com/google/domain/registry:registry_project"],
+)
+
+
+java_library(
+    name = "inputs",
+    srcs = glob(["*.java"]),
+    visibility = ["//visibility:public"],
+    deps = [
+        "//java/com/google/common/annotations",
+        "//java/com/google/common/base",
+        "//java/com/google/common/collect",
+        "//java/com/google/domain/registry/model",
+        "//java/com/google/domain/registry/util",
+        "//third_party/java/appengine:appengine-api",
+        "//third_party/java/appengine_mapreduce2:appengine_mapreduce",
+        "//third_party/java/appengine_pipeline",
+        "//third_party/java/dagger",
+        "//third_party/java/joda_time",
+        "//third_party/java/jsr330_inject",
+        "//third_party/java/objectify:objectify-v4_1",
+        "//third_party/java/servlet/servlet_api",
+    ],
+)

java/com/google/domain/registry/mapreduce/inputs/ChunkingKeyInput.java

@@ -0,0 +1,113 @@
+// Copyright 2016 Google Inc. All Rights Reserved.
+//
+// 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.domain.registry.mapreduce.inputs;
+
+import com.google.appengine.api.datastore.Key;
+import com.google.appengine.tools.mapreduce.Input;
+import com.google.appengine.tools.mapreduce.InputReader;
+import com.google.common.collect.ImmutableList;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.NoSuchElementException;
+
+/** A MapReduce {@link Input} adapter that chunks an input of keys into sublists of keys. */
+public class ChunkingKeyInput extends Input<List<Key>> {
+
+  private static final long serialVersionUID = 1670202385246824694L;
+
+  private final Input<Key> input;
+  private final int chunkSize;
+
+  public ChunkingKeyInput(Input<Key> input, int chunkSize) {
+    this.input = input;
+    this.chunkSize = chunkSize;
+  }
+
+  /**
+   * An input reader that wraps around another input reader and returns its contents in chunks of
+   * a given size.
+   */
+  private static class ChunkingKeyInputReader extends InputReader<List<Key>> {
+
+    private static final long serialVersionUID = 53502324675703263L;
+
+    private final InputReader<Key> reader;
+    private final int chunkSize;
+
+    ChunkingKeyInputReader(InputReader<Key> reader, int chunkSize) {
+      this.reader = reader;
+      this.chunkSize = chunkSize;
+    }
+
+    @Override
+    public List<Key> next() throws IOException {
+      ImmutableList.Builder<Key> chunk = new ImmutableList.Builder<>();
+      try {
+        for (int i = 0; i < chunkSize; i++) {
+          chunk.add(reader.next());
+        }
+      } catch (NoSuchElementException e) {
+        // Amazingly this is the recommended (and only) way to test for hasNext().
+      }
+      ImmutableList<Key> builtChunk = chunk.build();
+      if (builtChunk.isEmpty()) {
+        throw new NoSuchElementException();  // Maintain the contract.
+      }
+      return builtChunk;
+    }
+
+    @Override
+    public Double getProgress() {
+      return reader.getProgress();
+    }
+
+    @Override
+    public void beginShard() throws IOException {
+      reader.beginShard();
+    }
+
+    @Override
+    public void beginSlice() throws IOException {
+      reader.beginSlice();
+    }
+
+    @Override
+    public void endSlice() throws IOException {
+      reader.endSlice();
+    }
+
+    @Override
+    public void endShard() throws IOException {
+      reader.endShard();
+    }
+
+    @Override
+    public long estimateMemoryRequirement() {
+      // The reader's memory requirement plus the memory for this chunk's worth of buffered keys.
+      // 256 comes from DatastoreKeyInputReader.AVERAGE_KEY_SIZE.
+      return reader.estimateMemoryRequirement() + chunkSize * 256;
+    }
+  }
+
+  @Override
+  public List<InputReader<List<Key>>> createReaders() throws IOException {
+    ImmutableList.Builder<InputReader<List<Key>>> readers = new ImmutableList.Builder<>();
+    for (InputReader<Key> reader : input.createReaders()) {
+      readers.add(new ChunkingKeyInputReader(reader, chunkSize));
+    }
+    return readers.build();
+  }
+}

java/com/google/domain/registry/mapreduce/inputs/ConcatenatingInput.java

@@ -0,0 +1,66 @@
+// Copyright 2016 Google Inc. All Rights Reserved.
+//
+// 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.domain.registry.mapreduce.inputs;
+
+import com.google.appengine.tools.mapreduce.Input;
+import com.google.appengine.tools.mapreduce.InputReader;
+import com.google.appengine.tools.mapreduce.inputs.ConcatenatingInputReader;
+import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.ListMultimap;
+
+import java.io.IOException;
+import java.util.Collection;
+import java.util.List;
+import java.util.Set;
+
+/**
+ * A MapReduce {@link Input} adapter that joins multiple inputs.
+ *
+ * @param <T> input type
+ */
+public class ConcatenatingInput<T> extends Input<T> {
+
+  private static final long serialVersionUID = 1225981408139437077L;
+
+  private final Set<? extends Input<? extends T>> inputs;
+  private final int numShards;
+
+  public ConcatenatingInput(Iterable<? extends Input<? extends T>> inputs, int numShards) {
+    this.inputs = ImmutableSet.copyOf(inputs);
+    this.numShards = numShards;
+  }
+
+  @Override
+  public List<InputReader<T>> createReaders() throws IOException {
+    ListMultimap<Integer, InputReader<T>> shards = ArrayListMultimap.create();
+    int i = 0;
+    for (Input<? extends T> input : inputs) {
+      for (InputReader<? extends T> reader : input.createReaders()) {
+        // Covariant cast is safe because an InputReader<I> only outputs I and never consumes it.
+        @SuppressWarnings("unchecked")
+        InputReader<T> typedReader = (InputReader<T>) reader;
+        shards.put(i % numShards, typedReader);
+        i++;
+      }
+    }
+    ImmutableList.Builder<InputReader<T>> concatenatingReaders = new ImmutableList.Builder<>();
+    for (Collection<InputReader<T>> shard : shards.asMap().values()) {
+      concatenatingReaders.add(new ConcatenatingInputReader<>(ImmutableList.copyOf(shard)));
+    }
+    return concatenatingReaders.build();
+  }
+}

java/com/google/domain/registry/mapreduce/inputs/EppResourceBaseInput.java

@@ -0,0 +1,62 @@
+// Copyright 2016 Google Inc. All Rights Reserved.
+//
+// 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.domain.registry.mapreduce.inputs;
+
+import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.domain.registry.util.CollectionUtils.difference;
+
+import com.google.appengine.tools.mapreduce.Input;
+import com.google.appengine.tools.mapreduce.InputReader;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
+import com.google.domain.registry.model.EppResource;
+import com.google.domain.registry.model.index.EppResourceIndex;
+import com.google.domain.registry.model.index.EppResourceIndexBucket;
+
+import com.googlecode.objectify.Key;
+
+import java.util.List;
+
+/** Base class for {@link Input} classes that map over {@link EppResourceIndex}. */
+abstract class EppResourceBaseInput<I> extends Input<I> {
+
+  private static final long serialVersionUID = -6681886718929462122L;
+
+  @Override
+  public List<InputReader<I>> createReaders() {
+    ImmutableList.Builder<InputReader<I>> readers = new ImmutableList.Builder<>();
+    for (Key<EppResourceIndexBucket> bucketKey : EppResourceIndexBucket.getAllBuckets()) {
+      readers.add(bucketToReader(bucketKey));
+    }
+    return readers.build();
+  }
+
+  /** Creates a reader that returns the resources under a bucket. */
+  protected abstract InputReader<I> bucketToReader(Key<EppResourceIndexBucket> bucketKey);
+
+  static <R extends EppResource> void checkResourceClassesForInheritance(
+      ImmutableSet<Class<? extends R>> resourceClasses) {
+    for (Class<? extends R> resourceClass : resourceClasses) {
+      for (Class<? extends R> potentialSuperclass : difference(resourceClasses, resourceClass)) {
+        checkArgument(
+            !potentialSuperclass.isAssignableFrom(resourceClass),
+            "Cannot specify resource classes with inheritance relationship: %s extends %s",
+            resourceClass,
+            potentialSuperclass);
+      }
+    }
+  }
+}
+

java/com/google/domain/registry/mapreduce/inputs/EppResourceBaseReader.java

@@ -0,0 +1,146 @@
+// Copyright 2016 Google Inc. All Rights Reserved.
+//
+// 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.domain.registry.mapreduce.inputs;
+
+import static com.google.domain.registry.model.EntityClasses.CLASS_TO_KIND_FUNCTION;
+import static com.google.domain.registry.model.ofy.ObjectifyService.ofy;
+
+import com.google.appengine.api.datastore.Cursor;
+import com.google.appengine.api.datastore.QueryResultIterator;
+import com.google.appengine.tools.mapreduce.InputReader;
+import com.google.common.collect.FluentIterable;
+import com.google.common.collect.ImmutableSet;
+import com.google.domain.registry.model.EppResource;
+import com.google.domain.registry.model.index.EppResourceIndex;
+import com.google.domain.registry.model.index.EppResourceIndexBucket;
+import com.google.domain.registry.util.FormattingLogger;
+
+import com.googlecode.objectify.Key;
+import com.googlecode.objectify.cmd.Query;
+
+import java.util.NoSuchElementException;
+
+/** Base class for {@link InputReader} classes that map over {@link EppResourceIndex}. */
+abstract class EppResourceBaseReader<T> extends InputReader<T> {
+
+  static final FormattingLogger logger = FormattingLogger.getLoggerForCallerClass();
+
+  /** Number of bytes in 1MB of memory, used for memory estimates. */
+  static final long ONE_MB = 1024 * 1024;
+
+  private static final long serialVersionUID = -2970253037856017147L;
+
+  /**
+   * The resource kinds to filter for.
+   *
+   * <p>This can be empty, or any of {"ContactResource", "HostResource", "DomainBase"}. It will
+   * never contain "EppResource", "DomainResource" or "DomainApplication" since these aren't
+   * actual kinds in Datastore.
+   */
+  private final ImmutableSet<String> filterKinds;
+
+  private final Key<EppResourceIndexBucket> bucketKey;
+  private final long memoryEstimate;
+
+  private Cursor cursor;
+  private int total;
+  private int loaded;
+
+  private transient QueryResultIterator<EppResourceIndex> queryIterator;
+
+  EppResourceBaseReader(
+      Key<EppResourceIndexBucket>
+      bucketKey,
+      long memoryEstimate,
+      ImmutableSet<String> filterKinds) {
+    this.bucketKey = bucketKey;
+    this.memoryEstimate = memoryEstimate;
+    this.filterKinds = filterKinds;
+  }
+
+  /** Called once at start. Cache the expected size. */
+  @Override
+  public void beginShard() {
+    total = query().count();
+  }
+
+  /** Called every time we are deserialized. Create a new query or resume an existing one. */
+  @Override
+  public void beginSlice() {
+    Query<EppResourceIndex> query = query();
+    if (cursor != null) {
+      // The underlying query is strongly consistent, and according to the documentation at
+      // https://cloud.google.com/appengine/docs/java/datastore/queries#Java_Data_consistency
+      // "strongly consistent queries are always transactionally consistent". However, each time
+      // we restart the query at a cursor we have a new effective query, and "if the results for a
+      // query change between uses of a cursor, the query notices only changes that occur in
+      // results after the cursor. If a new result appears before the cursor's position for the
+      // query, it will not be returned when the results after the cursor are fetched."
+      // What this means in practice is that entities that are created after the initial query
+      // begins may or may not be seen by this reader, depending on whether the query was
+      // paused and restarted with a cursor before it would have reached the new entity.
+      query = query.startAt(cursor);
+    }
+    queryIterator = query.iterator();
+  }
+
+  /** Called occasionally alongside {@link #next}. */
+  @Override
+  public Double getProgress() {
+    // Cap progress at 1.0, since the query's count() can increase during the run of the mapreduce
+    // if more entities are written, but we've cached the value once in "total".
+    return Math.min(1.0, ((double) loaded) / total);
+  }
+
+  /** Called before we are serialized. Save a serializable cursor for this query. */
+  @Override
+  public void endSlice() {
+    cursor = queryIterator.getCursor();
+  }
+
+  /** Query for children of this bucket. */
+  Query<EppResourceIndex> query() {
+    Query<EppResourceIndex> query = ofy().load().type(EppResourceIndex.class).ancestor(bucketKey);
+    return filterKinds.isEmpty() ? query : query.filter("kind in", filterKinds);
+  }
+
+  /** Returns the estimated memory that will be used by this reader in bytes. */
+  @Override
+  public long estimateMemoryRequirement() {
+    return memoryEstimate;
+  }
+
+  /**
+   * Get the next {@link EppResourceIndex} from the query.
+   *
+   * @throws NoSuchElementException if there are no more elements.
+   */
+  EppResourceIndex nextEri() {
+    loaded++;
+    try {
+      return queryIterator.next();
+    } finally {
+      ofy().clearSessionCache();  // Try not to leak memory.
+    }
+  }
+
+  static <R extends EppResource> ImmutableSet<String> varargsToKinds(
+      ImmutableSet<Class<? extends R>> resourceClasses) {
+    // Ignore EppResource when finding kinds, since it doesn't have one and doesn't imply filtering.
+    return resourceClasses.contains(EppResource.class)
+          ? ImmutableSet.<String>of()
+          : FluentIterable.from(resourceClasses).transform(CLASS_TO_KIND_FUNCTION).toSet();
+  }
+}

java/com/google/domain/registry/mapreduce/inputs/EppResourceEntityInput.java

@@ -0,0 +1,43 @@
+// Copyright 2016 Google Inc. All Rights Reserved.
+//
+// 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.domain.registry.mapreduce.inputs;
+
+import com.google.appengine.tools.mapreduce.Input;
+import com.google.appengine.tools.mapreduce.InputReader;
+import com.google.common.collect.ImmutableSet;
+import com.google.domain.registry.model.EppResource;
+import com.google.domain.registry.model.index.EppResourceIndexBucket;
+
+import com.googlecode.objectify.Key;
+
+/** A MapReduce {@link Input} that loads all {@link EppResource} objects of a given type. */
+class EppResourceEntityInput<R extends EppResource> extends EppResourceBaseInput<R> {
+
+  private static final long serialVersionUID = 8162607479124406226L;
+
+  private final ImmutableSet<Class<? extends R>> resourceClasses;
+
+  public EppResourceEntityInput(ImmutableSet<Class<? extends R>> resourceClasses) {
+    this.resourceClasses = resourceClasses;
+    checkResourceClassesForInheritance(resourceClasses);
+  }
+
+  @Override
+  protected InputReader<R> bucketToReader(Key<EppResourceIndexBucket> bucketKey) {
+    return new EppResourceEntityReader<R>(bucketKey, resourceClasses);
+  }
+}
+
+

java/com/google/domain/registry/mapreduce/inputs/EppResourceEntityReader.java

@@ -0,0 +1,78 @@
+// Copyright 2016 Google Inc. All Rights Reserved.
+//
+// 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.domain.registry.mapreduce.inputs;
+
+import com.google.appengine.tools.mapreduce.InputReader;
+import com.google.common.collect.ImmutableSet;
+import com.google.domain.registry.model.EppResource;
+import com.google.domain.registry.model.index.EppResourceIndex;
+import com.google.domain.registry.model.index.EppResourceIndexBucket;
+
+import com.googlecode.objectify.Key;
+import com.googlecode.objectify.Ref;
+
+import java.util.NoSuchElementException;
+
+/** Reader that maps over {@link EppResourceIndex} and returns resources. */
+class EppResourceEntityReader<R extends EppResource> extends EppResourceBaseReader<R> {
+
+  private static final long serialVersionUID = -8042933349899971801L;
+
+  /**
+   * The resource classes to postfilter for.
+   *
+   * <p>This can be {@link EppResource} or any descendant classes, regardless of whether those
+   * classes map directly to a kind in datastore, with the restriction that none of the classes
+   * is a supertype of any of the others.
+   */
+  private final ImmutableSet<Class<? extends R>> resourceClasses;
+
+  public EppResourceEntityReader(
+      Key<EppResourceIndexBucket> bucketKey,
+      ImmutableSet<Class<? extends R>> resourceClasses) {
+    super(
+        bucketKey,
+        ONE_MB * 2,  // Estimate 2MB of memory for this reader, since it loads a (max 1MB) entity.
+        varargsToKinds(resourceClasses));
+    this.resourceClasses = resourceClasses;
+  }
+
+  /**
+   * Called for each map invocation.
+   *
+   * @throws NoSuchElementException if there are no more elements, as specified in the
+   *         {@link InputReader#next} Javadoc.
+   */
+  @Override
+  public R next() throws NoSuchElementException {
+    // Loop until we find a value, or nextRef() throws a NoSuchElementException.
+    while (true) {
+      Ref<? extends EppResource> reference = nextEri().getReference();
+      EppResource resource = reference.get();
+      if (resource == null) {
+        logger.severefmt("Broken ERI reference: %s", reference.getKey());
+        continue;
+      }
+      // Postfilter to distinguish polymorphic types (e.g. DomainBase and DomainResource).
+      for (Class<? extends R> resourceClass : resourceClasses) {
+        if (resourceClass.isAssignableFrom(resource.getClass())) {
+          @SuppressWarnings("unchecked")
+          R r = (R) resource;
+          return r;
+        }
+      }
+    }
+  }
+}

java/com/google/domain/registry/mapreduce/inputs/EppResourceIndexInput.java

@@ -0,0 +1,35 @@
+// Copyright 2016 Google Inc. All Rights Reserved.
+//
+// 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.domain.registry.mapreduce.inputs;
+
+import com.google.appengine.tools.mapreduce.Input;
+import com.google.appengine.tools.mapreduce.InputReader;
+import com.google.domain.registry.model.index.EppResourceIndex;
+import com.google.domain.registry.model.index.EppResourceIndexBucket;
+
+import com.googlecode.objectify.Key;
+
+/**
+ * A MapReduce {@link Input} that loads all {@link EppResourceIndex} entities.
+ */
+class EppResourceIndexInput extends EppResourceBaseInput<EppResourceIndex> {
+
+  private static final long serialVersionUID = -1231269296567279059L;
+
+  @Override
+  protected InputReader<EppResourceIndex> bucketToReader(Key<EppResourceIndexBucket> bucketKey) {
+    return new EppResourceIndexReader(bucketKey);
+  }
+}

java/com/google/domain/registry/mapreduce/inputs/EppResourceIndexReader.java

@@ -0,0 +1,47 @@
+// Copyright 2016 Google Inc. All Rights Reserved.
+//
+// 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.domain.registry.mapreduce.inputs;
+
+import com.google.appengine.tools.mapreduce.InputReader;
+import com.google.common.collect.ImmutableSet;
+import com.google.domain.registry.model.index.EppResourceIndex;
+import com.google.domain.registry.model.index.EppResourceIndexBucket;
+
+import com.googlecode.objectify.Key;
+
+import java.util.NoSuchElementException;
+
+/** Reader that maps over {@link EppResourceIndex} and returns the index objects themselves. */
+class EppResourceIndexReader extends EppResourceBaseReader<EppResourceIndex> {
+
+  private static final long serialVersionUID = -4816383426796766911L;
+
+  public EppResourceIndexReader(Key<EppResourceIndexBucket> bucketKey) {
+    // Estimate 1MB of memory for this reader, which is massive overkill.
+    // Use an empty set for the filter kinds, which disables filtering.
+    super(bucketKey, ONE_MB, ImmutableSet.<String>of());
+  }
+
+  /**
+   * Called for each map invocation.
+   *
+   * @throws NoSuchElementException if there are no more elements, as specified in the
+   *         {@link InputReader#next} Javadoc.
+   */
+  @Override
+  public EppResourceIndex next() throws NoSuchElementException {
+    return nextEri();
+  }
+}

java/com/google/domain/registry/mapreduce/inputs/EppResourceInputs.java

@@ -0,0 +1,81 @@
+// Copyright 2016 Google Inc. All Rights Reserved.
+//
+// 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.domain.registry.mapreduce.inputs;
+
+import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Predicates.not;
+import static com.google.common.collect.Iterables.all;
+import static com.google.common.collect.Lists.asList;
+import static com.google.domain.registry.util.TypeUtils.hasAnnotation;
+
+import com.google.appengine.tools.mapreduce.Input;
+import com.google.common.collect.ImmutableSet;
+import com.google.domain.registry.model.EppResource;
+import com.google.domain.registry.model.index.EppResourceIndex;
+
+import com.googlecode.objectify.Key;
+import com.googlecode.objectify.annotation.EntitySubclass;
+
+/**
+ * Mapreduce helpers for {@link EppResource} keys and objects.
+ *
+ * <p>The inputs provided by this class are not deletion-aware and do not project the resources
+ * forward in time. That is the responsibility of mappers that use these inputs.
+ */
+public final class EppResourceInputs {
+
+  private EppResourceInputs() {}
+
+  /** Returns a MapReduce {@link Input} that loads all {@link EppResourceIndex} objects. */
+  public static <R extends EppResource> Input<EppResourceIndex> createIndexInput() {
+    return new EppResourceIndexInput();
+  }
+
+  /**
+   * Returns a MapReduce {@link Input} that loads all {@link EppResource} objects of a given type,
+   * including deleted resources.
+   *
+   * <p>Note: Do not concatenate multiple EntityInputs together (this is inefficient as it iterates
+   * through all buckets multiple times). Specify the types in a single input, or load all types by
+   * specifying {@link EppResource} as the class.
+   */
+  @SafeVarargs
+  public static <R extends EppResource> Input<R> createEntityInput(
+      Class<? extends R> resourceClass,
+      Class<? extends R>... moreResourceClasses) {
+    return new EppResourceEntityInput<R>(
+        ImmutableSet.copyOf(asList(resourceClass, moreResourceClasses)));
+  }
+
+  /**
+   * Returns a MapReduce {@link Input} that loads keys to all {@link EppResource} objects of a given
+   * type, including deleted resources.
+   *
+   * <p>Note: Do not concatenate multiple KeyInputs together (this is inefficient as it iterates
+   * through all buckets multiple times). Specify the types in a single input, or load all types by
+   * specifying {@link EppResource} as the class.
+   */
+  @SafeVarargs
+  public static <R extends EppResource> Input<Key<R>> createKeyInput(
+      Class<? extends R> resourceClass,
+      Class<? extends R>... moreResourceClasses) {
+    ImmutableSet<Class<? extends R>> resourceClasses =
+        ImmutableSet.copyOf(asList(resourceClass, moreResourceClasses));
+    checkArgument(
+        all(resourceClasses, not(hasAnnotation(EntitySubclass.class))),
+        "Mapping over keys requires a non-polymorphic Entity");
+    return new EppResourceKeyInput<>(resourceClasses);
+  }
+}

java/com/google/domain/registry/mapreduce/inputs/EppResourceKeyInput.java

@@ -0,0 +1,45 @@
+// Copyright 2016 Google Inc. All Rights Reserved.
+//
+// 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.domain.registry.mapreduce.inputs;
+
+import com.google.appengine.tools.mapreduce.Input;
+import com.google.appengine.tools.mapreduce.InputReader;
+import com.google.common.collect.ImmutableSet;
+import com.google.domain.registry.model.EppResource;
+import com.google.domain.registry.model.index.EppResourceIndexBucket;
+
+import com.googlecode.objectify.Key;
+
+/**
+ * A MapReduce {@link Input} that loads keys to all {@link EppResource} objects of a given type.
+ *
+ * <p>When mapping over keys we can't distinguish between Objectify polymorphic types.
+ */
+class EppResourceKeyInput<R extends EppResource> extends EppResourceBaseInput<Key<R>> {
+
+  private static final long serialVersionUID = -5426821384707653743L;
+
+  private final ImmutableSet<Class<? extends R>> resourceClasses;
+
+  public EppResourceKeyInput(ImmutableSet<Class<? extends R>> resourceClasses) {
+    this.resourceClasses = resourceClasses;
+    checkResourceClassesForInheritance(resourceClasses);
+  }
+
+  @Override
+  protected InputReader<Key<R>> bucketToReader(Key<EppResourceIndexBucket> bucketKey) {
+    return new EppResourceKeyReader<>(bucketKey, resourceClasses);
+  }
+}

java/com/google/domain/registry/mapreduce/inputs/EppResourceKeyReader.java

@@ -0,0 +1,56 @@
+// Copyright 2016 Google Inc. All Rights Reserved.
+//
+// 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.domain.registry.mapreduce.inputs;
+
+import com.google.appengine.tools.mapreduce.InputReader;
+import com.google.common.collect.ImmutableSet;
+import com.google.domain.registry.model.EppResource;
+import com.google.domain.registry.model.index.EppResourceIndex;
+import com.google.domain.registry.model.index.EppResourceIndexBucket;
+
+import com.googlecode.objectify.Key;
+
+import java.util.NoSuchElementException;
+
+/**
+ * Reader that maps over {@link EppResourceIndex} and returns resource keys.
+ *
+ * <p>When mapping over keys we can't distinguish between Objectify polymorphic types.
+ */
+class EppResourceKeyReader<R extends EppResource> extends EppResourceBaseReader<Key<R>> {
+
+  private static final long serialVersionUID = -428232054739189774L;
+
+  public EppResourceKeyReader(
+      Key<EppResourceIndexBucket> bucketKey, ImmutableSet<Class<? extends R>> resourceClasses) {
+    super(
+        bucketKey,
+        ONE_MB,  // Estimate 1MB of memory for this reader, which is massive overkill.
+        varargsToKinds(resourceClasses));
+  }
+
+  /**
+   * Called for each map invocation.
+   *
+   * @throws NoSuchElementException if there are no more elements, as specified in the
+   *         {@link InputReader#next} Javadoc.
+   */
+  @Override
+  @SuppressWarnings("unchecked")
+  public Key<R> next() throws NoSuchElementException {
+    // This is a safe cast because we filtered on kind inside the query.
+    return (Key<R>) nextEri().getReference().getKey();
+  }
+}

java/com/google/domain/registry/mapreduce/inputs/NullInput.java

@@ -0,0 +1,54 @@
+// Copyright 2016 Google Inc. All Rights Reserved.
+//
+// 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.domain.registry.mapreduce.inputs;
+
+import com.google.appengine.tools.mapreduce.Input;
+import com.google.appengine.tools.mapreduce.InputReader;
+import com.google.common.collect.ImmutableList;
+
+import java.util.List;
+import java.util.NoSuchElementException;
+
+/** An input that returns a single {@code null} value. */
+public class NullInput<T> extends Input<T> {
+
+  private static final long serialVersionUID = 1816836937031979851L;
+
+  private static final class NullReader<T> extends InputReader<T> {
+
+    private static final long serialVersionUID = -8176201363578913125L;
+
+    boolean read = false;
+
+    @Override
+    public T next() throws NoSuchElementException {
+      if (read) {
+        throw new NoSuchElementException();
+      }
+      read = true;
+      return null;
+    }
+
+    @Override
+    public Double getProgress() {
+      return read ? 1.0 : 0.0;
+    }
+  }
+
+  @Override
+  public List<? extends InputReader<T>> createReaders() {
+    return ImmutableList.of(new NullReader<T>());
+  }
+}

java/com/google/domain/registry/rde/BUILD

@@ -17,6 +17,7 @@ java_library(
         "//java/com/google/domain/registry/gcs",
         "//java/com/google/domain/registry/keyring/api",
         "//java/com/google/domain/registry/mapreduce",
+        "//java/com/google/domain/registry/mapreduce/inputs",
         "//java/com/google/domain/registry/model",
         "//java/com/google/domain/registry/request",
         "//java/com/google/domain/registry/tldconfig/idn",

java/com/google/domain/registry/rde/RdeStagingAction.java

@@ -22,9 +22,9 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSetMultimap;
 import com.google.common.collect.Multimaps;
 import com.google.domain.registry.config.ConfigModule.Config;
-import com.google.domain.registry.mapreduce.EppResourceInputs;
 import com.google.domain.registry.mapreduce.MapreduceRunner;
-import com.google.domain.registry.mapreduce.NullInput;
+import com.google.domain.registry.mapreduce.inputs.EppResourceInputs;
+import com.google.domain.registry.mapreduce.inputs.NullInput;
 import com.google.domain.registry.model.EppResource;
 import com.google.domain.registry.model.contact.ContactResource;
 import com.google.domain.registry.model.host.HostResource;

java/com/google/domain/registry/tools/mapreduce/BUILD

@@ -13,8 +13,8 @@ java_library(
         "//java/com/google/common/net",
         "//java/com/google/common/primitives",
         "//java/com/google/common/util/concurrent",
-        "//java/com/google/domain/registry/flows",
         "//java/com/google/domain/registry/mapreduce",
+        "//java/com/google/domain/registry/mapreduce/inputs",
         "//java/com/google/domain/registry/model",
         "//java/com/google/domain/registry/request",
         "//java/com/google/domain/registry/util",

java/com/google/domain/registry/tools/mapreduce/DeleteProberDataAction.java

@@ -27,8 +27,8 @@ import com.google.common.collect.FluentIterable;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
-import com.google.domain.registry.mapreduce.EppResourceInputs;
 import com.google.domain.registry.mapreduce.MapreduceRunner;
+import com.google.domain.registry.mapreduce.inputs.EppResourceInputs;
 import com.google.domain.registry.model.domain.DomainApplication;
 import com.google.domain.registry.model.domain.DomainBase;
 import com.google.domain.registry.model.index.EppResourceIndex;

java/com/google/domain/registry/tools/mapreduce/ResaveAllEppResourcesAction.java

@@ -14,7 +14,7 @@
 
 package com.google.domain.registry.tools.mapreduce;
 
-import static com.google.domain.registry.mapreduce.EppResourceInputs.createEntityInput;
+import static com.google.domain.registry.mapreduce.inputs.EppResourceInputs.createEntityInput;
 import static com.google.domain.registry.model.ofy.ObjectifyService.ofy;
 import static com.google.domain.registry.util.PipelineUtils.createJobPath;
 

java/com/google/domain/registry/tools/server/BUILD

@@ -16,6 +16,7 @@ java_library(
         "//java/com/google/domain/registry/gcs",
         "//java/com/google/domain/registry/groups",
         "//java/com/google/domain/registry/mapreduce",
+        "//java/com/google/domain/registry/mapreduce/inputs",
         "//java/com/google/domain/registry/model",
         "//java/com/google/domain/registry/request",
         "//java/com/google/domain/registry/util",

java/com/google/domain/registry/tools/server/GenerateZoneFilesAction.java

@@ -19,7 +19,7 @@ import static com.google.common.base.Predicates.notNull;
 import static com.google.common.collect.Iterables.transform;
 import static com.google.common.collect.Iterators.filter;
 import static com.google.common.io.BaseEncoding.base16;
-import static com.google.domain.registry.mapreduce.EppResourceInputs.createEntityInput;
+import static com.google.domain.registry.mapreduce.inputs.EppResourceInputs.createEntityInput;
 import static com.google.domain.registry.model.EppResourceUtils.loadAtPointInTime;
 import static com.google.domain.registry.model.ofy.ObjectifyService.ofy;
 import static com.google.domain.registry.request.Action.Method.POST;
@@ -40,7 +40,7 @@ import com.google.common.collect.ImmutableSet;
 import com.google.domain.registry.config.ConfigModule.Config;
 import com.google.domain.registry.gcs.GcsUtils;
 import com.google.domain.registry.mapreduce.MapreduceRunner;
-import com.google.domain.registry.mapreduce.NullInput;
+import com.google.domain.registry.mapreduce.inputs.NullInput;
 import com.google.domain.registry.model.EppResource;
 import com.google.domain.registry.model.domain.DomainResource;
 import com.google.domain.registry.model.domain.ReferenceUnion;

java/com/google/domain/registry/tools/server/KillAllEppResourcesAction.java

@@ -23,9 +23,9 @@ import static com.google.domain.registry.util.PipelineUtils.createJobPath;
 import com.google.appengine.tools.mapreduce.Mapper;
 import com.google.common.collect.ImmutableList;
 import com.google.domain.registry.config.RegistryEnvironment;
-import com.google.domain.registry.mapreduce.EppResourceInputs;
 import com.google.domain.registry.mapreduce.MapreduceAction;
 import com.google.domain.registry.mapreduce.MapreduceRunner;
+import com.google.domain.registry.mapreduce.inputs.EppResourceInputs;
 import com.google.domain.registry.model.EppResource;
 import com.google.domain.registry.model.domain.DomainApplication;
 import com.google.domain.registry.model.index.DomainApplicationIndex;

java/com/google/domain/registry/tools/server/javascrap/AnnihilateNonDefaultNamespacesAction.java

@@ -30,9 +30,9 @@ import com.google.appengine.tools.mapreduce.Mapper;
 import com.google.appengine.tools.mapreduce.inputs.DatastoreKeyInput;
 import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableSet;
-import com.google.domain.registry.mapreduce.ChunkingKeyInput;
 import com.google.domain.registry.mapreduce.MapreduceAction;
 import com.google.domain.registry.mapreduce.MapreduceRunner;
+import com.google.domain.registry.mapreduce.inputs.ChunkingKeyInput;
 import com.google.domain.registry.request.Action;
 import com.google.domain.registry.request.Parameter;
 import com.google.domain.registry.request.Response;

javatests/com/google/domain/registry/mapreduce/inputs/BUILD

@@ -6,12 +6,12 @@ load("//java/com/google/testing/builddefs:GenTestRules.bzl", "GenTestRules")
 
 
 java_library(
-    name = "mapreduce",
+    name = "inputs",
     srcs = glob(["*.java"]),
     deps = [
         "//java/com/google/common/base",
         "//java/com/google/domain/registry/config",
-        "//java/com/google/domain/registry/mapreduce",
+        "//java/com/google/domain/registry/mapreduce/inputs",
         "//java/com/google/domain/registry/model",
         "//javatests/com/google/domain/registry/testing",
         "//third_party/java/appengine:appengine-api-testonly",
@@ -28,5 +28,5 @@ GenTestRules(
     default_test_size = "medium",
     jvm_flags = ["-XX:MaxPermSize=256m"],
     test_files = glob(["*Test.java"]),
-    deps = [":mapreduce"],
+    deps = [":inputs"],
 )

javatests/com/google/domain/registry/mapreduce/inputs/EppResourceInputsTest.java

@@ -12,12 +12,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package com.google.domain.registry.mapreduce;
+package com.google.domain.registry.mapreduce.inputs;
 
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assert_;
-import static com.google.domain.registry.mapreduce.EppResourceInputs.createEntityInput;
+import static com.google.domain.registry.mapreduce.inputs.EppResourceInputs.createEntityInput;
-import static com.google.domain.registry.mapreduce.EppResourceInputs.createKeyInput;
+import static com.google.domain.registry.mapreduce.inputs.EppResourceInputs.createKeyInput;
 import static com.google.domain.registry.model.index.EppResourceIndexBucket.getBucketKey;
 import static com.google.domain.registry.testing.DatastoreHelper.createTld;
 import static com.google.domain.registry.testing.DatastoreHelper.newDomainApplication;