[Test] [Do Not Merge] Run all tests with batch segment allocation enabled

core/src/main/java/org/apache/druid/java/util/common/Intervals.java

@@ -20,10 +20,13 @@
 package org.apache.druid.java.util.common;
 
 import com.google.common.collect.ImmutableList;
+import org.apache.druid.java.util.common.guava.Comparators;
 import org.joda.time.DateTime;
 import org.joda.time.Interval;
 import org.joda.time.chrono.ISOChronology;
 
+import javax.annotation.Nullable;
+
 public final class Intervals
 {
   public static final Interval ETERNITY = utc(JodaUtils.MIN_INSTANT, JodaUtils.MAX_INSTANT);
@@ -68,6 +71,32 @@ public static boolean isEternity(final Interval interval)
     return ETERNITY.equals(interval);
   }
 
+  /**
+   * Finds an interval from the given set of sortedIntervals which overlaps with
+   * the searchInterval. If multiple candidate intervals overlap with the
+   * searchInterval, the "smallest" interval based on the
+   * {@link Comparators#intervalsByStartThenEnd()} is returned.
+   *
+   * @param searchInterval  Interval which should overlap with the result
+   * @param sortedIntervals Candidate overlapping intervals, sorted in ascending
+   *                        order, using {@link Comparators#intervalsByStartThenEnd()}.
+   * @return The first overlapping interval, if one exists, otherwise null.
+   */
+  @Nullable
+  public static Interval findOverlappingInterval(Interval searchInterval, Interval[] sortedIntervals)
+  {
+    for (Interval interval : sortedIntervals) {
+      if (interval.overlaps(searchInterval)) {
+        return interval;
+      } else if (interval.getStart().isAfter(searchInterval.getEnd())) {
+        // Intervals after this cannot have an overlap
+        return null;
+      }
+    }
+
+    return null;
+  }
+
   private Intervals()
   {
   }

core/src/test/java/org/apache/druid/java/util/common/IntervalsTest.java

@@ -0,0 +1,81 @@
+/*
+ * 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.java.util.common;
+
+import org.apache.druid.java.util.common.guava.Comparators;
+import org.joda.time.Interval;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.util.Arrays;
+
+public class IntervalsTest
+{
+
+  @Test
+  public void testFindOverlappingInterval()
+  {
+    final Interval[] sortedIntervals = new Interval[]{
+        Intervals.of("2019/2020"),
+        Intervals.of("2021/2022"),
+        Intervals.of("2021-04-01/2021-05-01"),
+        Intervals.of("2022/2023")
+    };
+    Arrays.sort(sortedIntervals, Comparators.intervalsByStartThenEnd());
+
+    // Search interval outside the bounds of the sorted intervals
+    Assert.assertNull(
+        Intervals.findOverlappingInterval(Intervals.of("2018/2019"), sortedIntervals)
+    );
+    Assert.assertNull(
+        Intervals.findOverlappingInterval(Intervals.of("2023/2024"), sortedIntervals)
+    );
+
+    // Search interval within bounds, overlap exists
+    // Fully overlapping interval
+    Assert.assertEquals(
+        Intervals.of("2021/2022"),
+        Intervals.findOverlappingInterval(Intervals.of("2021/2022"), sortedIntervals)
+    );
+
+    // Partially overlapping interval
+    Assert.assertEquals(
+        Intervals.of("2022/2023"),
+        Intervals.findOverlappingInterval(Intervals.of("2022-01-01/2022-01-02"), sortedIntervals)
+    );
+
+    Assert.assertEquals(
+        Intervals.of("2021/2022"),
+        Intervals.findOverlappingInterval(Intervals.of("2021-06-01/2021-07-01"), sortedIntervals)
+    );
+
+    // Overlap with multiple intervals, "smallest" one is returned
+    Assert.assertEquals(
+        Intervals.of("2021/2022"),
+        Intervals.findOverlappingInterval(Intervals.of("2021-03-01/2021-04-01"), sortedIntervals)
+    );
+
+    // Search interval within bounds, no overlap
+    Assert.assertNull(
+        Intervals.findOverlappingInterval(Intervals.of("2020-01-02/2020-03-03"), sortedIntervals)
+    );
+  }
+
+}

indexing-service/src/main/java/org/apache/druid/indexing/common/TaskToolbox.java

@@ -20,7 +20,6 @@
 package org.apache.druid.indexing.common;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
-import com.google.common.base.Function;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Multimap;
@@ -329,14 +328,7 @@ public void publishSegments(Iterable<DataSegment> segments) throws IOException
     // Request segment pushes for each set
     final Multimap<Interval, DataSegment> segmentMultimap = Multimaps.index(
         segments,
-        new Function<DataSegment, Interval>()
-        {
-          @Override
-          public Interval apply(DataSegment segment)
-          {
-            return segment.getInterval();
-          }
-        }
+        DataSegment::getInterval
     );
     for (final Collection<DataSegment> segmentCollection : segmentMultimap.asMap().values()) {
       getTaskActionClient().submit(new SegmentInsertAction(ImmutableSet.copyOf(segmentCollection)));

indexing-service/src/main/java/org/apache/druid/indexing/common/actions/LocalTaskActionClient.java

@@ -30,6 +30,7 @@
 
 import javax.annotation.Nullable;
 import java.util.Map;
+import java.util.concurrent.TimeUnit;
 
 public class LocalTaskActionClient implements TaskActionClient
 {
@@ -76,11 +77,28 @@ public <RetType> RetType submit(TaskAction<RetType> taskAction)
     }
 
     final long performStartTime = System.currentTimeMillis();
-    final RetType result = taskAction.perform(task, toolbox);
+    final RetType result = performAction(taskAction);
     emitTimerMetric("task/action/run/time", taskAction, System.currentTimeMillis() - performStartTime);
     return result;
   }
 
+  private <R> R performAction(TaskAction<R> taskAction)
+  {
+    try {
+      final R result;
+      if (taskAction.canPerformAsync(task, toolbox)) {
+        result = taskAction.performAsync(task, toolbox).get(5, TimeUnit.MINUTES);
+      } else {
+        result = taskAction.perform(task, toolbox);
+      }
+
+      return result;
+    }
+    catch (Throwable t) {
+      throw new RuntimeException(t);
+    }
+  }
+
   private void emitTimerMetric(final String metric, final TaskAction<?> action, final long time)
   {
     final ServiceMetricEvent.Builder metricBuilder = ServiceMetricEvent.builder();

indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentAllocateAction.java

@@ -47,6 +47,7 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
+import java.util.concurrent.Future;
 import java.util.concurrent.ThreadLocalRandom;
 import java.util.stream.Collectors;
 
@@ -180,6 +181,23 @@ public TypeReference<SegmentIdWithShardSpec> getReturnTypeReference()
     };
   }
 
+  @Override
+  public boolean canPerformAsync(Task task, TaskActionToolbox toolbox)
+  {
+    return toolbox.canBatchSegmentAllocation();
+  }
+
+  @Override
+  public Future<SegmentIdWithShardSpec> performAsync(Task task, TaskActionToolbox toolbox)
+  {
+    if (!toolbox.canBatchSegmentAllocation()) {
+      throw new ISE("Batched segment allocation is disabled");
+    }
+    return toolbox.getSegmentAllocationQueue().add(
+        new SegmentAllocateRequest(task, this, MAX_ATTEMPTS)
+    );
+  }
+
   @Override
   public SegmentIdWithShardSpec perform(
       final Task task,

indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentAllocateRequest.java

@@ -0,0 +1,77 @@
+/*
+ * 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.indexing.common.actions;
+
+import org.apache.druid.indexing.common.task.Task;
+import org.joda.time.Interval;
+import org.joda.time.chrono.ISOChronology;
+
+/**
+ * Request received by the overlord for segment allocation.
+ */
+public class SegmentAllocateRequest
+{
+  private final Task task;
+  private final SegmentAllocateAction action;
+  private final int maxAttempts;
+  private final Interval rowInterval;
+
+  private int attempts;
+
+  public SegmentAllocateRequest(Task task, SegmentAllocateAction action, int maxAttempts)
+  {
+    this.task = task;
+    this.action = action;
+    this.maxAttempts = maxAttempts;
+    this.rowInterval = action.getQueryGranularity()
+                             .bucket(action.getTimestamp())
+                             .withChronology(ISOChronology.getInstanceUTC());
+  }
+
+  public Task getTask()
+  {
+    return task;
+  }
+
+  public SegmentAllocateAction getAction()
+  {
+    return action;
+  }
+
+  public void incrementAttempts()
+  {
+    ++attempts;
+  }
+
+  public boolean canRetry()
+  {
+    return attempts < maxAttempts;
+  }
+
+  public int getAttempts()
+  {
+    return attempts;
+  }
+
+  public Interval getRowInterval()
+  {
+    return rowInterval;
+  }
+}

indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentAllocateResult.java

@@ -0,0 +1,49 @@
+/*
+ * 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.indexing.common.actions;
+
+import org.apache.druid.segment.realtime.appenderator.SegmentIdWithShardSpec;
+
+public class SegmentAllocateResult
+{
+  private final SegmentIdWithShardSpec segmentId;
+  private final String errorMessage;
+
+  public SegmentAllocateResult(SegmentIdWithShardSpec segmentId, String errorMessage)
+  {
+    this.segmentId = segmentId;
+    this.errorMessage = errorMessage;
+  }
+
+  public SegmentIdWithShardSpec getSegmentId()
+  {
+    return segmentId;
+  }
+
+  public String getErrorMessage()
+  {
+    return errorMessage;
+  }
+
+  public boolean isSuccess()
+  {
+    return segmentId != null;
+  }
+}