feat: add support for push/pull/kill segments in google cloud deep storage without zip compression

docs/development/extensions-core/google.md

@@ -53,3 +53,4 @@ To configure connectivity to google cloud, run druid processes with `GOOGLE_APPL
 |`druid.google.bucket`||Google Storage bucket name.|Must be set.|
 |`druid.google.prefix`|A prefix string that will be prepended to the blob names for the segments published to Google deep storage| |""|
 |`druid.google.maxListingLength`|maximum number of input files matching a given prefix to retrieve at a time| |1024|
+|`druid.storage.zip`|Whether segments are written as directories of objects (`false`) or as a single `index.zip` object (`true`). Writing segments unzipped requires permission to delete objects under the segment path, so that a push replaces the segment already there.|`true`, `false`|`true`|

extensions-core/google-extensions/src/main/java/org/apache/druid/storage/google/GoogleDataSegmentKiller.java

@@ -61,19 +61,48 @@ public void kill(DataSegment segment) throws SegmentLoadingException
     Map<String, Object> loadSpec = segment.getLoadSpec();
     final String bucket = MapUtils.getString(loadSpec, "bucket");
     final String indexPath = MapUtils.getString(loadSpec, "path");
-    final String descriptorPath = DataSegmentKiller.descriptorPath(indexPath);
 
     try {
-      deleteIfPresent(bucket, indexPath);
-      // descriptor.json is a file to store segment metadata in deep storage. This file is deprecated and not stored
-      // anymore, but we still delete them if exists.
-      deleteIfPresent(bucket, descriptorPath);
+      if (indexPath.endsWith("/")) {
+        // segment was pushed unzipped, so the path names a directory of objects; delete every one of them
+        deleteObjectsInPath(bucket, indexPath);
+      } else {
+        deleteIfPresent(bucket, indexPath);
+        // descriptor.json is a file to store segment metadata in deep storage. This file is deprecated and not stored
+        // anymore, but we still delete them if exists.
+        deleteIfPresent(bucket, DataSegmentKiller.descriptorPath(indexPath));
+      }
     }
     catch (StorageException e) {
       throw new SegmentLoadingException(e, "Couldn't kill segment[%s]: [%s]", segment.getId(), e.getMessage());
     }
   }
 
+  private void deleteObjectsInPath(String bucket, String pathPrefix) throws SegmentLoadingException
+  {
+    try {
+      GoogleUtils.deleteObjectsInPath(
+          storage,
+          inputDataConfig,
+          bucket,
+          pathPrefix,
+          Predicates.alwaysTrue()
+      );
+    }
+    catch (StorageException e) {
+      throw e;
+    }
+    catch (Exception e) {
+      throw new SegmentLoadingException(
+          e,
+          "Couldn't delete objects under [gs://%s/%s]: [%s]",
+          bucket,
+          pathPrefix,
+          e.getMessage()
+      );
+    }
+  }
+
   private void deleteIfPresent(String bucket, String path)
   {
     try {

extensions-core/google-extensions/src/main/java/org/apache/druid/storage/google/GoogleDataSegmentPuller.java

@@ -20,8 +20,11 @@
 package org.apache.druid.storage.google;
 
 import com.google.common.base.Predicate;
+import com.google.common.collect.ImmutableList;
 import com.google.inject.Inject;
+import org.apache.druid.data.input.impl.CloudObjectLocation;
 import org.apache.druid.java.util.common.FileUtils;
+import org.apache.druid.java.util.common.RetryUtils;
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.java.util.common.logger.Logger;
 import org.apache.druid.segment.loading.SegmentLoadingException;
@@ -32,17 +35,21 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.net.URI;
+import java.nio.file.Paths;
+import java.util.Iterator;
 
 public class GoogleDataSegmentPuller implements URIDataPuller
 {
   private static final Logger LOG = new Logger(GoogleDataSegmentPuller.class);
 
   protected final GoogleStorage storage;
+  private final GoogleInputDataConfig inputDataConfig;
 
   @Inject
-  public GoogleDataSegmentPuller(final GoogleStorage storage)
+  public GoogleDataSegmentPuller(final GoogleStorage storage, final GoogleInputDataConfig inputDataConfig)
   {
     this.storage = storage;
+    this.inputDataConfig = inputDataConfig;
   }
 
   FileUtils.FileCopyResult getSegmentFiles(final String bucket, final String path, File outDir)
@@ -53,13 +60,12 @@ FileUtils.FileCopyResult getSegmentFiles(final String bucket, final String path,
     try {
       FileUtils.mkdirp(outDir);
 
-      final GoogleByteSource byteSource = new GoogleByteSource(storage, bucket, path);
-      final FileUtils.FileCopyResult result = CompressionUtils.unzip(
-          byteSource,
-          outDir,
-          GoogleUtils::isRetryable,
-          false
-      );
+      // A trailing slash means the segment was pushed unzipped (druid.storage.zip=false), so the path names a
+      // directory of objects to pull individually rather than a single zip to unpack.
+      final FileUtils.FileCopyResult result = path.endsWith("/")
+                                              ? getSegmentFilesFromDirectory(bucket, path, outDir)
+                                              : unzipSegmentFiles(bucket, path, outDir);
+
       LOG.info("Loaded %d bytes from [%s] to [%s]", result.size(), path, outDir.getAbsolutePath());
       return result;
     }
@@ -79,6 +85,44 @@ FileUtils.FileCopyResult getSegmentFiles(final String bucket, final String path,
     }
   }
 
+  private FileUtils.FileCopyResult unzipSegmentFiles(final String bucket, final String path, final File outDir)
+      throws IOException
+  {
+    final GoogleByteSource byteSource = new GoogleByteSource(storage, bucket, path);
+    return CompressionUtils.unzip(
+        byteSource,
+        outDir,
+        GoogleUtils::isRetryable,
+        false
+    );
+  }
+
+  private FileUtils.FileCopyResult getSegmentFilesFromDirectory(
+      final String bucket,
+      final String pathPrefix,
+      final File outDir
+  )
+  {
+    final Iterator<GoogleStorageObjectMetadata> objects = GoogleUtils.lazyFetchingStorageObjectsIterator(
+        storage,
+        ImmutableList.of(new CloudObjectLocation(bucket, pathPrefix).toUri(GoogleStorageDruidModule.SCHEME_GS))
+                     .iterator(),
+        inputDataConfig.getMaxListingLength()
+    );
+
+    final FileUtils.FileCopyResult copyResult = new FileUtils.FileCopyResult();
+    while (objects.hasNext()) {
+      final GoogleStorageObjectMetadata object = objects.next();
+      final GoogleByteSource byteSource = new GoogleByteSource(storage, bucket, object.getName());
+      final File outFile = new File(outDir, Paths.get(object.getName()).getFileName().toString());
+      copyResult.addFiles(
+          FileUtils.retryCopy(byteSource, outFile, GoogleUtils.GOOGLE_RETRY, RetryUtils.DEFAULT_MAX_TRIES).getFiles()
+      );
+    }
+
+    return copyResult;
+  }
+
   @Override
   public InputStream getInputStream(URI uri) throws IOException
   {

extensions-core/google-extensions/src/main/java/org/apache/druid/storage/google/GoogleDataSegmentPusher.java

@@ -24,35 +24,53 @@
 import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableMap;
 import com.google.inject.Inject;
+import org.apache.druid.java.util.common.IOE;
 import org.apache.druid.java.util.common.RE;
 import org.apache.druid.java.util.common.RetryUtils;
 import org.apache.druid.java.util.common.logger.Logger;
 import org.apache.druid.segment.SegmentUtils;
 import org.apache.druid.segment.loading.DataSegmentPusher;
+import org.apache.druid.segment.loading.DeepStorageSegmentConfig;
 import org.apache.druid.timeline.DataSegment;
 import org.apache.druid.utils.CompressionUtils;
 
 import java.io.File;
 import java.io.IOException;
 import java.net.URI;
 import java.nio.file.Files;
+import java.util.HashSet;
 import java.util.Map;
+import java.util.Set;
 
 public class GoogleDataSegmentPusher implements DataSegmentPusher
 {
   private static final Logger log = new Logger(GoogleDataSegmentPusher.class);
 
+  static final String INDEX_ZIP_FILE_NAME = "index.zip";
+
+  /**
+   * Originally Google deep storage always wrote segments zipped, so that is what {@code druid.storage.zip} falls back
+   * to.
+   */
+  private static final boolean DEFAULT_ZIP = true;
+
   private final GoogleStorage storage;
   private final GoogleAccountConfig config;
+  private final GoogleInputDataConfig inputDataConfig;
+  private final boolean zip;
 
   @Inject
   public GoogleDataSegmentPusher(
       final GoogleStorage storage,
-      final GoogleAccountConfig config
+      final GoogleAccountConfig config,
+      final GoogleInputDataConfig inputDataConfig,
+      final DeepStorageSegmentConfig deepStorageConfig
   )
   {
     this.storage = storage;
     this.config = config;
+    this.inputDataConfig = inputDataConfig;
+    this.zip = deepStorageConfig.isZip(DEFAULT_ZIP);
   }
 
   public void insert(final File file, final String contentType, final String path)
@@ -91,12 +109,29 @@ public DataSegment push(final File indexFilesDir, final DataSegment segment, fin
   public DataSegment pushToPath(File indexFilesDir, DataSegment segment, String storageDirSuffix) throws IOException
   {
     final int version = SegmentUtils.getVersionFromDir(indexFilesDir);
+    final String basePath = buildPath(storageDirSuffix);
+
+    try {
+      if (zip) {
+        return pushZip(indexFilesDir, segment, version, basePath);
+      } else {
+        return pushNoZip(indexFilesDir, segment, version, basePath);
+      }
+    }
+    catch (Exception e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  private DataSegment pushZip(File indexFilesDir, DataSegment segment, int version, String basePath)
+      throws IOException
+  {
     File indexFile = null;
 
     try {
       indexFile = Files.createTempFile("index", ".zip").toFile();
       final long indexSize = CompressionUtils.zip(indexFilesDir, indexFile);
-      final String indexPath = buildPath(storageDirSuffix + "/" + "index.zip");
+      final String indexPath = basePath + "/" + INDEX_ZIP_FILE_NAME;
 
       final DataSegment outSegment = segment
           .withSize(indexSize)
@@ -107,9 +142,6 @@ public DataSegment pushToPath(File indexFilesDir, DataSegment segment, String st
 
       return outSegment;
     }
-    catch (Exception e) {
-      throw new RuntimeException(e);
-    }
     finally {
       if (indexFile != null) {
         log.debug("Deleting file [%s]", indexFile);
@@ -118,6 +150,76 @@ public DataSegment pushToPath(File indexFilesDir, DataSegment segment, String st
     }
   }
 
+  /**
+   * Uploads the segment files as they are, one object per file, under {@code basePath}. The resulting loadSpec path is
+   * the directory itself, with a trailing slash to tell {@link GoogleDataSegmentPuller} and
+   * {@link GoogleDataSegmentKiller} that it names a directory of files rather than a single object.
+   */
+  private DataSegment pushNoZip(File indexFilesDir, DataSegment segment, int version, String basePath)
+      throws IOException
+  {
+    final File[] files = indexFilesDir.listFiles();
+    if (files == null) {
+      throw new IOE("Cannot list directory [%s]", indexFilesDir);
+    }
+
+    final String dirPath = basePath + "/";
+    final Set<String> pushedPaths = new HashSet<>();
+
+    long size = 0;
+    for (final File file : files) {
+      if (file.isFile()) {
+        size += file.length();
+        final String path = dirPath + file.getName();
+        insert(file, "application/octet-stream", path);
+        pushedPaths.add(path);
+      } else {
+        // Segment directories are expected to be flat.
+        throw new IOE("Unexpected subdirectory [%s]", file.getName());
+      }
+    }
+
+    deleteStaleObjects(dirPath, pushedPaths);
+
+    return segment.withSize(size)
+                  .withLoadSpec(makeLoadSpec(config.getBucket(), dirPath))
+                  .withBinaryVersion(version);
+  }
+
+  /**
+   * Removes everything under {@code dirPath} that this push did not write.
+   * <p>
+   * A zipped push replaces the previous segment outright, because one {@code index.zip} object overwrites another, and
+   * {@link #push} with {@code useUniquePath = false} is expected to replace a previous push the same way. Uploading
+   * file by file only overwrites the names the new segment happens to share, so without this a re-push could leave
+   * behind objects of whatever was there before: a stale {@code index.zip} from a zipped push, or smoosh chunks from a
+   * larger prior v9 segment.
+   * <p>
+   * Note that this makes an unzipped push require permission to delete objects under the segment path, which a zipped
+   * push does not.
+   */
+  private void deleteStaleObjects(final String dirPath, final Set<String> pushedPaths) throws IOException
+  {
+    try {
+      GoogleUtils.deleteObjectsInPath(
+          storage,
+          inputDataConfig,
+          config.getBucket(),
+          dirPath,
+          object -> !pushedPaths.contains(object.getName())
+      );
+    }
+    catch (Exception e) {
+      throw new IOE(
+          e,
+          "Could not remove objects left under [gs://%s/%s] by a previous push, which would be loaded as part of this"
+          + " segment",
+          config.getBucket(),
+          dirPath
+      );
+    }
+  }
+
   @VisibleForTesting
   String buildPath(final String path)
   {

extensions-core/google-extensions/src/main/java/org/apache/druid/storage/google/GoogleTimestampVersionedDataFinder.java

@@ -35,9 +35,9 @@ public class GoogleTimestampVersionedDataFinder extends GoogleDataSegmentPuller
   private static final long MAX_LISTING_KEYS = 1000;
 
   @Inject
-  public GoogleTimestampVersionedDataFinder(final GoogleStorage storage)
+  public GoogleTimestampVersionedDataFinder(final GoogleStorage storage, final GoogleInputDataConfig inputDataConfig)
   {
-    super(storage);
+    super(storage, inputDataConfig);
   }
 
   @Override

extensions-core/google-extensions/src/test/java/org/apache/druid/storage/google/GoogleDataSegmentKillerTest.java

@@ -95,6 +95,38 @@ public void killTest() throws SegmentLoadingException
     verifyAll();
   }
 
+  @Test
+  public void test_kill_unzippedSegment_deletesEveryObjectInTheDirectory() throws SegmentLoadingException, IOException
+  {
+    // pushed with druid.storage.zip=false, so the path is the directory holding the segment files
+    final String unzippedPath = "test/2015-04-12T00:00:00.000Z_2015-04-13T00:00:00.000Z/1/0/";
+    final DataSegment segment = DataSegment.builder(DATA_SEGMENT)
+                                           .loadSpec(ImmutableMap.of("bucket", BUCKET, "path", unzippedPath))
+                                           .build();
+
+    EasyMock.expect(inputDataConfig.getMaxListingLength()).andReturn(MAX_KEYS).anyTimes();
+    EasyMock.expect(storage.list(EasyMock.eq(BUCKET), EasyMock.eq(unzippedPath), EasyMock.anyObject(), EasyMock.anyObject()))
+            .andReturn(new GoogleStorageObjectPage(
+                ImmutableList.of(
+                    new GoogleStorageObjectMetadata(BUCKET, unzippedPath + "version.bin", 4L, TIME_0),
+                    new GoogleStorageObjectMetadata(BUCKET, unzippedPath + "meta.smoosh", 8L, TIME_0)
+                ),
+                null
+            ));
+    storage.delete(BUCKET, unzippedPath + "version.bin");
+    EasyMock.expectLastCall();
+    storage.delete(BUCKET, unzippedPath + "meta.smoosh");
+    EasyMock.expectLastCall();
+
+    replayAll();
+
+    GoogleDataSegmentKiller killer = new GoogleDataSegmentKiller(storage, accountConfig, inputDataConfig);
+
+    killer.kill(segment);
+
+    verifyAll();
+  }
+
   @Test
   public void killWithErrorTest()
   {