Refactor SQLMetadataSegmentManager; Change contract of REST methods in DataSourcesResource

benchmarks/src/main/java/org/apache/druid/server/coordinator/NewestSegmentFirstPolicyBenchmark.java

@@ -20,6 +20,8 @@
 package org.apache.druid.server.coordinator;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import org.apache.druid.client.DataSourcesSnapshot;
 import org.apache.druid.java.util.common.DateTimes;
 import org.apache.druid.server.coordinator.helper.CompactionSegmentIterator;
 import org.apache.druid.server.coordinator.helper.CompactionSegmentSearchPolicy;
@@ -42,6 +44,7 @@
 import org.openjdk.jmh.annotations.State;
 import org.openjdk.jmh.infra.Blackhole;
 
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -102,14 +105,10 @@ public void setup()
       );
     }
 
-    dataSources = new HashMap<>();
+    List<DataSegment> segments = new ArrayList<>();
     for (int i = 0; i < numDataSources; i++) {
       final String dataSource = DATA_SOURCE_PREFIX + i;
 
-      VersionedIntervalTimeline<String, DataSegment> timeline = new VersionedIntervalTimeline<>(
-          String.CASE_INSENSITIVE_ORDER
-      );
-
       final int startYear = ThreadLocalRandom.current().nextInt(2000, 2040);
       DateTime date = DateTimes.of(startYear, 1, 1, 0, 0);
 
@@ -127,12 +126,11 @@ public void setup()
               0,
               segmentSizeBytes
           );
-          timeline.add(segment.getInterval(), segment.getVersion(), shardSpec.createChunk(segment));
+          segments.add(segment);
         }
       }
-
-      dataSources.put(dataSource, timeline);
     }
+    dataSources = DataSourcesSnapshot.fromUsedSegments(segments, ImmutableMap.of()).getUsedSegmentsTimelinesPerDataSource();
   }
 
   @Benchmark

core/src/main/java/org/apache/druid/java/util/common/logger/Logger.java

@@ -123,6 +123,11 @@ public void error(Throwable t, String message, Object... formatArgs)
     log.error(StringUtils.nonStrictFormat(message, formatArgs), t);
   }
 
+  public void assertionError(String message, Object... formatArgs)
+  {
+    log.error("ASSERTION_ERROR: " + message, formatArgs);
+  }
+
   public void wtf(String message, Object... formatArgs)
   {
     log.error(StringUtils.nonStrictFormat("WTF?!: " + message, formatArgs), new Exception());

core/src/main/java/org/apache/druid/java/util/emitter/EmittingLogger.java

@@ -28,6 +28,7 @@
 import org.apache.druid.java.util.emitter.service.AlertBuilder;
 import org.apache.druid.java.util.emitter.service.ServiceEmitter;
 
+import javax.annotation.Nullable;
 import java.io.PrintWriter;
 import java.io.StringWriter;
 
@@ -61,15 +62,22 @@ public AlertBuilder makeAlert(String message, Object... objects)
     return makeAlert(null, message, objects);
   }
 
-  public AlertBuilder makeAlert(Throwable t, String message, Object... objects)
+  public AlertBuilder makeAlert(@Nullable Throwable t, String message, Object... objects)
   {
     if (emitter == null) {
       final String errorMessage = StringUtils.format(
-          "Emitter not initialized!  Cannot alert.  Please make sure to call %s.registerEmitter()", this.getClass()
+          "Emitter not initialized!  Cannot alert.  Please make sure to call %s.registerEmitter()\n"
+          + "Message: %s",
+          this.getClass(),
+          StringUtils.nonStrictFormat(message, objects)
       );
 
       error(errorMessage);
-      throw new ISE(errorMessage);
+      ISE e = new ISE(errorMessage);
+      if (t != null) {
+        e.addSuppressed(t);
+      }
+      throw e;
     }
 
     final AlertBuilder retVal = new EmittingAlertBuilder(t, StringUtils.format(message, objects), emitter)

core/src/main/java/org/apache/druid/timeline/TimelineLookup.java

@@ -22,6 +22,7 @@
 import org.apache.druid.timeline.partition.PartitionHolder;
 import org.joda.time.Interval;
 
+import javax.annotation.Nullable;
 import java.util.List;
 
 
@@ -50,5 +51,5 @@ public interface TimelineLookup<VersionType, ObjectType>
    */
   List<TimelineObjectHolder<VersionType, ObjectType>> lookupWithIncompletePartitions(Interval interval);
 
-  PartitionHolder<ObjectType> findEntry(Interval interval, VersionType version);
+  @Nullable PartitionHolder<ObjectType> findEntry(Interval interval, VersionType version);
 }

core/src/main/java/org/apache/druid/timeline/VersionedIntervalTimeline.java

@@ -23,16 +23,18 @@
 import com.google.common.base.Function;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.Iterators;
-import com.google.common.collect.Ordering;
 import org.apache.druid.java.util.common.DateTimes;
 import org.apache.druid.java.util.common.UOE;
 import org.apache.druid.java.util.common.guava.Comparators;
 import org.apache.druid.timeline.partition.ImmutablePartitionHolder;
 import org.apache.druid.timeline.partition.PartitionChunk;
 import org.apache.druid.timeline.partition.PartitionHolder;
+import org.apache.druid.utils.CollectionUtils;
 import org.joda.time.Interval;
 
+import javax.annotation.Nullable;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Comparator;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -44,7 +46,9 @@
 import java.util.Objects;
 import java.util.Set;
 import java.util.TreeMap;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.StreamSupport;
 
 /**
  * VersionedIntervalTimeline is a data structure that manages objects on a specific timeline.
@@ -75,12 +79,11 @@ public class VersionedIntervalTimeline<VersionType, ObjectType> implements Timel
       Comparators.intervalsByStartThenEnd()
   );
   private final Map<Interval, TreeMap<VersionType, TimelineEntry>> allTimelineEntries = new HashMap<>();
+  private final AtomicInteger numObjects = new AtomicInteger();
 
   private final Comparator<? super VersionType> versionComparator;
 
-  public VersionedIntervalTimeline(
-      Comparator<? super VersionType> versionComparator
-  )
+  public VersionedIntervalTimeline(Comparator<? super VersionType> versionComparator)
   {
     this.versionComparator = versionComparator;
   }
@@ -92,7 +95,8 @@ public static VersionedIntervalTimeline<String, DataSegment> forSegments(Iterabl
 
   public static VersionedIntervalTimeline<String, DataSegment> forSegments(Iterator<DataSegment> segments)
   {
-    final VersionedIntervalTimeline<String, DataSegment> timeline = new VersionedIntervalTimeline<>(Ordering.natural());
+    final VersionedIntervalTimeline<String, DataSegment> timeline =
+        new VersionedIntervalTimeline<>(Comparator.naturalOrder());
     addSegments(timeline, segments);
     return timeline;
   }
@@ -115,6 +119,28 @@ public Map<Interval, TreeMap<VersionType, TimelineEntry>> getAllTimelineEntries(
     return allTimelineEntries;
   }
 
+  /**
+   * Returns a lazy collection with all objects (including overshadowed, see {@link #findOvershadowed}) in this
+   * VersionedIntervalTimeline to be used for iteration or {@link Collection#stream()} transformation. The order of
+   * objects in this collection is unspecified.
+   *
+   * Note: iteration over the returned collection may not be as trivially cheap as, for example, iteration over an
+   * ArrayList. Try (to some reasonable extent) to organize the code so that it iterates the returned collection only
+   * once rather than several times.
+   */
+  public Collection<ObjectType> iterateAllObjects()
+  {
+    return CollectionUtils.createLazyCollectionFromStream(
+        () -> allTimelineEntries
+            .values()
+            .stream()
+            .flatMap((TreeMap<VersionType, TimelineEntry> entryMap) -> entryMap.values().stream())
+            .flatMap((TimelineEntry entry) -> StreamSupport.stream(entry.getPartitionHolder().spliterator(), false))
+            .map(PartitionChunk::getObject),
+        numObjects.get()
+    );
+  }
+
   public void add(final Interval interval, VersionType version, PartitionChunk<ObjectType> object)
   {
     addAll(Iterators.singletonIterator(object), o -> interval, o -> version);
@@ -143,15 +169,19 @@ private void addAll(
           TreeMap<VersionType, TimelineEntry> versionEntry = new TreeMap<>(versionComparator);
           versionEntry.put(version, entry);
           allTimelineEntries.put(interval, versionEntry);
+          numObjects.incrementAndGet();
         } else {
           entry = exists.get(version);
 
           if (entry == null) {
             entry = new TimelineEntry(interval, version, new PartitionHolder<>(object));
             exists.put(version, entry);
+            numObjects.incrementAndGet();
           } else {
             PartitionHolder<ObjectType> partitionHolder = entry.getPartitionHolder();
-            partitionHolder.add(object);
+            if (partitionHolder.add(object)) {
+              numObjects.incrementAndGet();
+            }
           }
         }
 
@@ -174,6 +204,7 @@ private void addAll(
     }
   }
 
+  @Nullable
   public PartitionChunk<ObjectType> remove(Interval interval, VersionType version, PartitionChunk<ObjectType> chunk)
   {
     try {
@@ -189,7 +220,11 @@ public PartitionChunk<ObjectType> remove(Interval interval, VersionType version,
         return null;
       }
 
-      PartitionChunk<ObjectType> retVal = entry.getPartitionHolder().remove(chunk);
+      PartitionChunk<ObjectType> removedChunk = entry.getPartitionHolder().remove(chunk);
+      if (removedChunk == null) {
+        return null;
+      }
+      numObjects.decrementAndGet();
       if (entry.getPartitionHolder().isEmpty()) {
         versionEntries.remove(version);
         if (versionEntries.isEmpty()) {
@@ -201,25 +236,23 @@ public PartitionChunk<ObjectType> remove(Interval interval, VersionType version,
 
       remove(completePartitionsTimeline, interval, entry, false);
 
-      return retVal;
+      return removedChunk;
     }
     finally {
       lock.writeLock().unlock();
     }
   }
 
   @Override
-  public PartitionHolder<ObjectType> findEntry(Interval interval, VersionType version)
+  public @Nullable PartitionHolder<ObjectType> findEntry(Interval interval, VersionType version)
   {
     try {
       lock.readLock().lock();
       for (Map.Entry<Interval, TreeMap<VersionType, TimelineEntry>> entry : allTimelineEntries.entrySet()) {
         if (entry.getKey().equals(interval) || entry.getKey().contains(interval)) {
           TimelineEntry foundEntry = entry.getValue().get(version);
           if (foundEntry != null) {
-            return new ImmutablePartitionHolder<ObjectType>(
-                foundEntry.getPartitionHolder()
-            );
+            return new ImmutablePartitionHolder<>(foundEntry.getPartitionHolder());
           }
         }
       }
@@ -307,6 +340,10 @@ private TimelineObjectHolder<VersionType, ObjectType> timelineEntryToObjectHolde
     );
   }
 
+  /**
+   * This method should be deduplicated with DataSourcesSnapshot.determineOvershadowedSegments(): see
+   * https://github.com/apache/incubator-druid/issues/8070.
+   */
   public Set<TimelineObjectHolder<VersionType, ObjectType>> findOvershadowed()
   {
     try {
@@ -315,8 +352,8 @@ public Set<TimelineObjectHolder<VersionType, ObjectType>> findOvershadowed()
 
       Map<Interval, Map<VersionType, TimelineEntry>> overShadowed = new HashMap<>();
       for (Map.Entry<Interval, TreeMap<VersionType, TimelineEntry>> versionEntry : allTimelineEntries.entrySet()) {
-        Map<VersionType, TimelineEntry> versionCopy = new HashMap<>();
-        versionCopy.putAll(versionEntry.getValue());
+        @SuppressWarnings("unchecked")
+        Map<VersionType, TimelineEntry> versionCopy = (TreeMap) versionEntry.getValue().clone();
         overShadowed.put(versionEntry.getKey(), versionCopy);
       }
 

core/src/main/java/org/apache/druid/timeline/partition/ImmutablePartitionHolder.java

@@ -23,7 +23,7 @@
  */
 public class ImmutablePartitionHolder<T> extends PartitionHolder<T>
 {
-  public ImmutablePartitionHolder(PartitionHolder partitionHolder)
+  public ImmutablePartitionHolder(PartitionHolder<T> partitionHolder)
   {
     super(partitionHolder);
   }
@@ -35,7 +35,7 @@ public PartitionChunk<T> remove(PartitionChunk<T> tPartitionChunk)
   }
 
   @Override
-  public void add(PartitionChunk<T> tPartitionChunk)
+  public boolean add(PartitionChunk<T> tPartitionChunk)
   {
     throw new UnsupportedOperationException();
   }