Fix hanging threads in GcsDiffFileLister (#1243)

* Fix hanging threads in GcsDiffFileLister

Basically, whenever we request threads using the request thread factory,
we must be on the request thread itself. Dagger doesn't guarantee this
for us if we provide the ExecutorService directly in the action (or in
the GcsDiffFileLister), but we can gurantee that we're on the request
thread itself by simply injecting a Lazy, so that the executor is
instantiated inside the request itself.

In addition, add a timeout on the futures just in case.

core/src/main/java/google/registry/backup/BackupModule.java

@@ -35,6 +35,8 @@ import google.registry.request.HttpException.BadRequestException;
 import google.registry.request.Parameter;
 import java.lang.annotation.Documented;
 import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
 import javax.inject.Qualifier;
 import javax.servlet.http.HttpServletRequest;
 import org.joda.time.DateTime;
@@ -101,4 +103,9 @@ public final class BackupModule {
   static ListeningExecutorService provideListeningExecutorService() {
     return listeningDecorator(newFixedThreadPool(NUM_THREADS, currentRequestThreadFactory()));
   }
+
+  @Provides
+  static ScheduledExecutorService provideScheduledExecutorService() {
+    return Executors.newSingleThreadScheduledExecutor();
+  }
 }

core/src/main/java/google/registry/backup/GcsDiffFileLister.java

@@ -28,11 +28,15 @@ import com.google.common.flogger.FluentLogger;
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.ListeningExecutorService;
+import com.google.common.util.concurrent.UncheckedExecutionException;
+import dagger.Lazy;
 import google.registry.backup.BackupModule.Backups;
 import google.registry.gcs.GcsUtils;
 import java.io.IOException;
+import java.time.Duration;
 import java.util.Map;
 import java.util.TreeMap;
+import java.util.concurrent.ScheduledExecutorService;
 import javax.annotation.Nullable;
 import javax.inject.Inject;
 import org.joda.time.DateTime;
@@ -42,15 +46,25 @@ class GcsDiffFileLister {
 
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
 
+  /** Timeout for retrieving per-file information from GCS. */
+  private static final Duration FILE_INFO_TIMEOUT_DURATION = Duration.ofMinutes(1);
+
   @Inject GcsUtils gcsUtils;
 
-  @Inject @Backups ListeningExecutorService executor;
+  @Inject @Backups Lazy<ListeningExecutorService> lazyExecutor;
-  @Inject GcsDiffFileLister() {}
+  @Inject ScheduledExecutorService scheduledExecutorService;
+
+  @Inject
+  GcsDiffFileLister() {}
 
   /**
    * Traverses the sequence of diff files backwards from checkpointTime and inserts the file
    * metadata into "sequence". Returns true if a complete sequence was discovered, false if one or
    * more files are missing.
+   *
+   * @throws UncheckedExecutionException wrapping a {@link java.util.concurrent.TimeoutException} if
+   *     the GCS call fails to finish within one minute, or wrapping any other exception if
+   *     something else goes wrong.
    */
   private boolean constructDiffSequence(
       String gcsBucket,
@@ -62,7 +76,12 @@ class GcsDiffFileLister {
     while (isBeforeOrAt(fromTime, checkpointTime)) {
       BlobInfo blobInfo;
       if (upperBoundTimesToBlobInfo.containsKey(checkpointTime)) {
-        blobInfo = Futures.getUnchecked(upperBoundTimesToBlobInfo.get(checkpointTime));
+        blobInfo =
+            Futures.getUnchecked(
+                Futures.withTimeout(
+                    upperBoundTimesToBlobInfo.get(checkpointTime),
+                    FILE_INFO_TIMEOUT_DURATION,
+                    scheduledExecutorService));
       } else {
         String filename = DIFF_FILE_PREFIX + checkpointTime;
         logger.atInfo().log("Patching GCS list; discovered file: %s", filename);
@@ -102,12 +121,14 @@ class GcsDiffFileLister {
     }
     DateTime lastUpperBoundTime = START_OF_TIME;
 
+    TreeMap<DateTime, BlobInfo> sequence = new TreeMap<>();
+    try {
       for (String strippedFilename : strippedFilenames) {
         final String filename = DIFF_FILE_PREFIX + strippedFilename;
         DateTime upperBoundTime = DateTime.parse(strippedFilename);
         if (isInRange(upperBoundTime, fromTime, toTime)) {
           upperBoundTimesToBlobInfo.put(
-            upperBoundTime, executor.submit(() -> getBlobInfo(gcsBucket, filename)));
+              upperBoundTime, lazyExecutor.get().submit(() -> getBlobInfo(gcsBucket, filename)));
           lastUpperBoundTime = latestOf(upperBoundTime, lastUpperBoundTime);
         }
       }
@@ -116,15 +137,14 @@ class GcsDiffFileLister {
         return ImmutableList.of();
       }
 
-    // Reconstruct the sequence of files by traversing backwards from "lastUpperBoundTime" (i.e. the
+      // Reconstruct the sequence of files by traversing backwards from "lastUpperBoundTime" (i.e.
-    // last file that we found) and finding its previous file until we either run out of files or
+      // the last file that we found) and finding its previous file until we either run out of files
-    // get to one that precedes "fromTime".
+      // or get to one that precedes "fromTime".
       //
       // GCS file listing is eventually consistent, so it's possible that we are missing a file. The
       // metadata of a file is sufficient to identify the preceding file, so if we start from the
       // last file and work backwards we can verify that we have no holes in our chain (although we
       // may be missing files at the end).
-    TreeMap<DateTime, BlobInfo> sequence = new TreeMap<>();
       logger.atInfo().log("Restoring until: %s", lastUpperBoundTime);
       boolean inconsistentFileSet =
           !constructDiffSequence(
@@ -152,6 +172,9 @@ class GcsDiffFileLister {
           !inconsistentFileSet,
           "Unable to compute commit diff history, there are either gaps or forks in the history "
               + "file set.  Check log for details.");
+    } finally {
+      lazyExecutor.get().shutdown();
+    }
 
     logger.atInfo().log(
         "Actual restore from time: %s", getLowerBoundTime(sequence.firstEntry().getValue()));

core/src/test/java/google/registry/backup/GcsDiffFileListerTest.java

@@ -16,7 +16,6 @@ package google.registry.backup;
 
 import static com.google.common.collect.Iterables.transform;
 import static com.google.common.truth.Truth.assertThat;
-import static com.google.common.util.concurrent.MoreExecutors.newDirectExecutorService;
 import static google.registry.backup.BackupUtils.GcsMetadataKeys.LOWER_BOUND_CHECKPOINT;
 import static google.registry.backup.ExportCommitLogDiffAction.DIFF_FILE_PREFIX;
 import static org.joda.time.DateTimeZone.UTC;
@@ -33,10 +32,12 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.flogger.LoggerConfig;
 import com.google.common.testing.TestLogHandler;
+import com.google.common.util.concurrent.MoreExecutors;
 import google.registry.gcs.GcsUtils;
 import google.registry.gcs.backport.LocalStorageHelper;
 import google.registry.testing.AppEngineExtension;
 import java.io.IOException;
+import java.util.concurrent.Executors;
 import java.util.logging.LogRecord;
 import org.joda.time.DateTime;
 import org.junit.jupiter.api.BeforeEach;
@@ -60,7 +61,8 @@ public class GcsDiffFileListerTest {
   @BeforeEach
   void beforeEach() throws Exception {
     diffLister.gcsUtils = gcsUtils;
-    diffLister.executor = newDirectExecutorService();
+    diffLister.lazyExecutor = MoreExecutors::newDirectExecutorService;
+    diffLister.scheduledExecutorService = Executors.newSingleThreadScheduledExecutor();
     for (int i = 0; i < 5; i++) {
       addGcsFile(i, i + 1);
     }

core/src/test/java/google/registry/backup/ReplayCommitLogsToSqlActionTest.java

@@ -16,7 +16,6 @@ package google.registry.backup;
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.truth.Truth.assertThat;
-import static com.google.common.util.concurrent.MoreExecutors.newDirectExecutorService;
 import static google.registry.backup.RestoreCommitLogsActionTest.createCheckpoint;
 import static google.registry.backup.RestoreCommitLogsActionTest.saveDiffFile;
 import static google.registry.backup.RestoreCommitLogsActionTest.saveDiffFileNotToRestore;
@@ -42,6 +41,7 @@ import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.truth.Truth8;
+import com.google.common.util.concurrent.MoreExecutors;
 import com.googlecode.objectify.Key;
 import google.registry.gcs.GcsUtils;
 import google.registry.gcs.backport.LocalStorageHelper;
@@ -72,6 +72,7 @@ import google.registry.testing.FakeResponse;
 import google.registry.testing.TestObject;
 import google.registry.util.RequestStatusChecker;
 import java.io.IOException;
+import java.util.concurrent.Executors;
 import org.joda.time.DateTime;
 import org.joda.time.Duration;
 import org.junit.jupiter.api.AfterEach;
@@ -131,7 +132,8 @@ public class ReplayCommitLogsToSqlActionTest {
     action.gcsBucket = "gcs bucket";
     action.diffLister = new GcsDiffFileLister();
     action.diffLister.gcsUtils = gcsUtils;
-    action.diffLister.executor = newDirectExecutorService();
+    action.diffLister.lazyExecutor = MoreExecutors::newDirectExecutorService;
+    action.diffLister.scheduledExecutorService = Executors.newSingleThreadScheduledExecutor();
     ofyTm()
         .transact(
             () ->

core/src/test/java/google/registry/backup/RestoreCommitLogsActionTest.java

@@ -17,7 +17,6 @@ package google.registry.backup;
 import static com.google.common.collect.Iterables.transform;
 import static com.google.common.collect.Maps.toMap;
 import static com.google.common.truth.Truth.assertThat;
-import static com.google.common.util.concurrent.MoreExecutors.newDirectExecutorService;
 import static google.registry.backup.BackupUtils.GcsMetadataKeys.LOWER_BOUND_CHECKPOINT;
 import static google.registry.backup.BackupUtils.serializeEntity;
 import static google.registry.backup.ExportCommitLogDiffAction.DIFF_FILE_PREFIX;
@@ -34,6 +33,7 @@ import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Lists;
 import com.google.common.io.Resources;
 import com.google.common.primitives.Longs;
+import com.google.common.util.concurrent.MoreExecutors;
 import com.googlecode.objectify.Key;
 import google.registry.gcs.GcsUtils;
 import google.registry.gcs.backport.LocalStorageHelper;
@@ -55,6 +55,7 @@ import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Optional;
+import java.util.concurrent.Executors;
 import org.joda.time.DateTime;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -87,7 +88,8 @@ public class RestoreCommitLogsActionTest {
     action.gcsBucketOverride = Optional.empty();
     action.diffLister = new GcsDiffFileLister();
     action.diffLister.gcsUtils = gcsUtils;
-    action.diffLister.executor = newDirectExecutorService();
+    action.diffLister.lazyExecutor = MoreExecutors::newDirectExecutorService;
+    action.diffLister.scheduledExecutorService = Executors.newSingleThreadScheduledExecutor();
   }
 
   @Test