Ability to Delete task logs and segments from S3

core/src/main/java/org/apache/druid/common/utils/CurrentTimeMillisSupplier.java

@@ -0,0 +1,32 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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 org.apache.druid.common.utils;
+
+import java.util.function.LongSupplier;
+
+public class CurrentTimeMillisSupplier implements LongSupplier
+{
+  @Override
+  public long getAsLong()
+  {
+    return System.currentTimeMillis();
+  }
+}

extensions-core/s3-extensions/src/main/java/org/apache/druid/storage/s3/S3DataSegmentKiller.java

@@ -20,27 +20,38 @@
 package org.apache.druid.storage.s3;
 
 import com.amazonaws.AmazonServiceException;
+import com.google.common.base.Predicates;
 import com.google.inject.Inject;
 import org.apache.druid.java.util.common.MapUtils;
 import org.apache.druid.java.util.common.logger.Logger;
 import org.apache.druid.segment.loading.DataSegmentKiller;
 import org.apache.druid.segment.loading.SegmentLoadingException;
 import org.apache.druid.timeline.DataSegment;
 
+import java.io.IOException;
 import java.util.Map;
 
 /**
+ *
  */
 public class S3DataSegmentKiller implements DataSegmentKiller
 {
   private static final Logger log = new Logger(S3DataSegmentKiller.class);
 
   private final ServerSideEncryptingAmazonS3 s3Client;
+  private final S3DataSegmentPusherConfig segmentPusherConfig;
+  private final S3InputDataConfig inputDataConfig;
 
   @Inject
-  public S3DataSegmentKiller(ServerSideEncryptingAmazonS3 s3Client)
+  public S3DataSegmentKiller(
+      ServerSideEncryptingAmazonS3 s3Client,
+      S3DataSegmentPusherConfig segmentPusherConfig,
+      S3InputDataConfig inputDataConfig
+  )
   {
     this.s3Client = s3Client;
+    this.segmentPusherConfig = segmentPusherConfig;
+    this.inputDataConfig = inputDataConfig;
   }
 
   @Override
@@ -69,8 +80,23 @@ public void kill(DataSegment segment) throws SegmentLoadingException
   }
 
   @Override
-  public void killAll()
+  public void killAll() throws IOException
   {
-    throw new UnsupportedOperationException("not implemented");
+    log.info("Deleting all segment files from s3 location [bucket: '%s' prefix: '%s']",
+             segmentPusherConfig.getBucket(), segmentPusherConfig.getBaseKey()
+    );
+    try {
+      S3Utils.deleteObjectsInPath(
+          s3Client,
+          inputDataConfig,
+          segmentPusherConfig.getBucket(),
+          segmentPusherConfig.getBaseKey(),
+          Predicates.alwaysTrue()
+      );
+    }
+    catch (Exception e) {
+      log.error("Error occurred while deleting segment files from s3. Error: %s", e.getMessage());
+      throw new IOException(e);
+    }
   }
 }

extensions-core/s3-extensions/src/main/java/org/apache/druid/storage/s3/S3StorageDruidModule.java

@@ -155,6 +155,7 @@ public void configure(Binder binder)
            .to(S3DataSegmentArchiver.class)
            .in(LazySingleton.class);
     Binders.dataSegmentPusherBinder(binder).addBinding(SCHEME).to(S3DataSegmentPusher.class).in(LazySingleton.class);
+    JsonConfigProvider.bind(binder, "druid.storage", S3InputDataConfig.class);
     JsonConfigProvider.bind(binder, "druid.storage", S3DataSegmentPusherConfig.class);
     JsonConfigProvider.bind(binder, "druid.storage", S3DataSegmentArchiverConfig.class);
     JsonConfigProvider.bind(binder, "druid.storage", S3StorageConfig.class);

extensions-core/s3-extensions/src/main/java/org/apache/druid/storage/s3/S3TaskLogs.java

@@ -27,6 +27,7 @@
 import com.google.common.base.Throwables;
 import com.google.common.io.ByteSource;
 import com.google.inject.Inject;
+import org.apache.druid.common.utils.CurrentTimeMillisSupplier;
 import org.apache.druid.java.util.common.IOE;
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.java.util.common.logger.Logger;
@@ -35,6 +36,7 @@
 import java.io.File;
 import java.io.IOException;
 import java.io.InputStream;
+import java.util.Date;
 
 /**
  * Provides task logs archived on S3.
@@ -45,12 +47,21 @@ public class S3TaskLogs implements TaskLogs
 
   private final ServerSideEncryptingAmazonS3 service;
   private final S3TaskLogsConfig config;
+  private final S3InputDataConfig inputDataConfig;
+  private final CurrentTimeMillisSupplier timeSupplier;
 
   @Inject
-  public S3TaskLogs(ServerSideEncryptingAmazonS3 service, S3TaskLogsConfig config)
+  public S3TaskLogs(
+      ServerSideEncryptingAmazonS3 service,
+      S3TaskLogsConfig config,
+      S3InputDataConfig inputDataConfig,
+      CurrentTimeMillisSupplier timeSupplier
+  )
   {
     this.service = service;
     this.config = config;
+    this.inputDataConfig = inputDataConfig;
+    this.timeSupplier = timeSupplier;
   }
 
   @Override
@@ -152,14 +163,34 @@ String getTaskLogKey(String taskid, String filename)
   }
 
   @Override
-  public void killAll()
+  public void killAll() throws IOException
   {
-    throw new UnsupportedOperationException("not implemented");
+    log.info("Deleting all task logs from s3 location [bucket: %s    prefix: %s].",
+             config.getS3Bucket(), config.getS3Prefix()
+    );
+
+    long now = timeSupplier.getAsLong();
+    killOlderThan(now);
   }
 
   @Override
-  public void killOlderThan(long timestamp)
+  public void killOlderThan(long timestamp) throws IOException
   {
-    throw new UnsupportedOperationException("not implemented");
+    log.info("Deleting all task logs from s3 location [bucket: '%s' prefix: '%s'] older than %s.",
+             config.getS3Bucket(), config.getS3Prefix(), new Date(timestamp)
+    );
+    try {
+      S3Utils.deleteObjectsInPath(
+          service,
+          inputDataConfig,
+          config.getS3Bucket(),
+          config.getS3Prefix(),
+          (object) -> object.getLastModified().getTime() < timestamp
+      );
+    }
+    catch (Exception e) {
+      log.error("Error occurred while deleting task log files from s3. Error: %s", e.getMessage());
+      throw new IOException(e);
+    }
   }
 }

extensions-core/s3-extensions/src/main/java/org/apache/druid/storage/s3/S3TaskLogsConfig.java

@@ -61,6 +61,12 @@ public String getS3Prefix()
     return s3Prefix;
   }
 
+  @VisibleForTesting
+  void setS3Prefix(String s3Prefix)
+  {
+    this.s3Prefix = s3Prefix;
+  }
+
   public boolean getDisableAcl()
   {
     return disableAcl;

extensions-core/s3-extensions/src/main/java/org/apache/druid/storage/s3/S3Utils.java

@@ -23,6 +23,7 @@
 import com.amazonaws.services.s3.model.AccessControlList;
 import com.amazonaws.services.s3.model.AmazonS3Exception;
 import com.amazonaws.services.s3.model.CanonicalGrantee;
+import com.amazonaws.services.s3.model.DeleteObjectsRequest;
 import com.amazonaws.services.s3.model.Grant;
 import com.amazonaws.services.s3.model.ListObjectsV2Request;
 import com.amazonaws.services.s3.model.ListObjectsV2Result;
@@ -31,6 +32,7 @@
 import com.amazonaws.services.s3.model.S3ObjectSummary;
 import com.google.common.base.Joiner;
 import com.google.common.base.Predicate;
+import com.google.common.collect.ImmutableList;
 import org.apache.druid.data.input.impl.CloudObjectLocation;
 import org.apache.druid.java.util.common.ISE;
 import org.apache.druid.java.util.common.RetryUtils;
@@ -41,7 +43,9 @@
 import java.io.File;
 import java.io.IOException;
 import java.net.URI;
+import java.util.ArrayList;
 import java.util.Iterator;
+import java.util.List;
 
 /**
  *
@@ -200,6 +204,54 @@ public static S3ObjectSummary getSingleObjectSummary(ServerSideEncryptingAmazonS
     return objectSummary;
   }
 
+  public static void deleteObjectsInPath(
+      ServerSideEncryptingAmazonS3 s3Client,
+      S3InputDataConfig config,
+      String bucket,
+      String prefix,
+      Predicate<S3ObjectSummary> filter
+  )
+      throws Exception
+  {
+    final List<DeleteObjectsRequest.KeyVersion> keysToDelete = new ArrayList<>(config.getMaxListingLength());
+    final ObjectSummaryIterator iterator = new ObjectSummaryIterator(
+        s3Client,
+        ImmutableList.of(new CloudObjectLocation(bucket, prefix).toUri("s3")),
+        config.getMaxListingLength()
+    );
+
+    while (iterator.hasNext()) {
+      final S3ObjectSummary nextObject = iterator.next();
+      if (filter.apply(nextObject)) {
+        keysToDelete.add(new DeleteObjectsRequest.KeyVersion(nextObject.getKey()));
+        if (keysToDelete.size() == config.getMaxListingLength()) {
+          deleteBucketKeys(s3Client, bucket, keysToDelete);
+          log.info("Deleted %d files", keysToDelete.size());
+          keysToDelete.clear();
+        }
+      }
+    }
+
+    if (keysToDelete.size() > 0) {
+      deleteBucketKeys(s3Client, bucket, keysToDelete);
+      log.info("Deleted %d files", keysToDelete.size());
+    }
+  }
+
+  public static void deleteBucketKeys(
+      ServerSideEncryptingAmazonS3 s3Client,
+      String bucket,
+      List<DeleteObjectsRequest.KeyVersion> keysToDelete
+  )
+      throws Exception
+  {
+    DeleteObjectsRequest deleteRequest = new DeleteObjectsRequest(bucket).withKeys(keysToDelete);
+    S3Utils.retryS3Operation(() -> {
+      s3Client.deleteObjects(deleteRequest);
+      return null;
+    });
+  }
+
   /**
    * Uploads a file to S3 if possible. First trying to set ACL to give the bucket owner full control of the file before uploading.
    *

extensions-core/s3-extensions/src/main/java/org/apache/druid/storage/s3/ServerSideEncryptingAmazonS3.java

@@ -25,6 +25,7 @@
 import com.amazonaws.services.s3.model.AccessControlList;
 import com.amazonaws.services.s3.model.CopyObjectRequest;
 import com.amazonaws.services.s3.model.CopyObjectResult;
+import com.amazonaws.services.s3.model.DeleteObjectsRequest;
 import com.amazonaws.services.s3.model.GetObjectMetadataRequest;
 import com.amazonaws.services.s3.model.GetObjectRequest;
 import com.amazonaws.services.s3.model.ListObjectsV2Request;
@@ -128,6 +129,11 @@ public void deleteObject(String bucket, String key)
     amazonS3.deleteObject(bucket, key);
   }
 
+  public void deleteObjects(DeleteObjectsRequest request)
+  {
+    amazonS3.deleteObjects(request);
+  }
+
   public static class Builder
   {
     private AmazonS3ClientBuilder amazonS3ClientBuilder = AmazonS3Client.builder();