More cleanup

Author: kfaraz

server/src/main/java/org/apache/druid/metadata/SqlSegmentsMetadataQuery.java

@@ -524,7 +524,7 @@ private CloseableIterator<DataSegmentPlus> retrieveSegmentBatchById(
     } else {
       final Query<Map<String, Object>> query = handle.createQuery(
           StringUtils.format(
-              "SELECT payload, used, upgraded_from_segment_id, used_status_last_updated"
+              "SELECT payload, used, upgraded_from_segment_id, used_status_last_updated, created_date"
               + " FROM %s WHERE dataSource = :dataSource %s",
               dbTables.getSegmentsTable(), getParameterizedInConditionForColumn("id", segmentIds)
           )
@@ -538,7 +538,7 @@ private CloseableIterator<DataSegmentPlus> retrieveSegmentBatchById(
           .map(
               (index, r, ctx) -> new DataSegmentPlus(
                   JacksonUtils.readValue(jsonMapper, r.getBytes(1), DataSegment.class),
-                  null,
+                  DateTimes.of(r.getString(5)),
                   DateTimes.of(r.getString(4)),
                   r.getBoolean(2),
                   null,

server/src/main/java/org/apache/druid/metadata/segment/cache/HeapMemoryDatasourceSegmentCache.java

@@ -96,24 +96,9 @@ boolean isEmpty()
    */
   boolean shouldRefreshUsedSegment(SegmentId segmentId, @Nullable DateTime persistedUpdateTime)
   {
-    return withReadLock(() -> {
-      final DataSegmentPlus cachedState = readSegmentsFor(segmentId.getInterval())
-          .idToUsedSegment.get(segmentId);
-      return cachedState == null
-             || shouldUpdateCache(cachedState.getUsedStatusLastUpdatedDate(), persistedUpdateTime);
-    });
-  }
-
-  /**
-   * Checks if a pending segment needs to be refreshed in the cache.
-   */
-  boolean shouldRefreshPendingSegment(PendingSegmentRecord record)
-  {
-    final SegmentIdWithShardSpec segmentId = record.getId();
     return withReadLock(
-        () -> !readSegmentsFor(segmentId.getInterval())
-            .idToPendingSegment
-            .containsKey(segmentId.toString())
+        () -> readSegmentsFor(segmentId.getInterval())
+            .shouldRefreshUsedSegment(segmentId, persistedUpdateTime)
     );
   }
 
@@ -139,60 +124,6 @@ private static boolean shouldUpdateCache(
     }
   }
 
-  /**
-   * Adds or updates the given segment in the cache.
-   *
-   * @return true if the segment was updated in the cache, false if the segment
-   * was left unchanged in the cache.
-   */
-  boolean addSegment(DataSegmentPlus segmentPlus)
-  {
-    if (Boolean.TRUE.equals(segmentPlus.getUsed())) {
-      return addUsedSegment(segmentPlus);
-    } else {
-      return addUnusedSegmentId(
-          segmentPlus.getDataSegment().getId(),
-          segmentPlus.getUsedStatusLastUpdatedDate()
-      );
-    }
-  }
-
-  /**
-   * Adds or updates a used segment in the cache.
-   */
-  private boolean addUsedSegment(DataSegmentPlus segmentPlus)
-  {
-    final DataSegment segment = segmentPlus.getDataSegment();
-    final SegmentId segmentId = segment.getId();
-
-    return withWriteLock(() -> {
-      if (!shouldRefreshUsedSegment(segmentId, segmentPlus.getUsedStatusLastUpdatedDate())) {
-        return false;
-      }
-
-      final SegmentsInInterval segments = writeSegmentsFor(segmentId.getInterval());
-      segments.idToUsedSegment.put(segmentId, segmentPlus);
-      segments.unusedSegmentIdToUpdatedTime.remove(segment.getId());
-      return true;
-    });
-  }
-
-  /**
-   * Adds or updates an unused segment in the cache.
-   *
-   * @param updatedTime Last updated time of this segment as persisted in the
-   *                    metadata store. This value can be null for segments
-   *                    persisted to the metadata store before the column
-   *                    used_status_last_updated was added to the segments table.
-   */
-  boolean addUnusedSegmentId(SegmentId segmentId, @Nullable DateTime updatedTime)
-  {
-    return withWriteLock(
-        () -> writeSegmentsFor(segmentId.getInterval())
-            .addUnusedSegmentId(segmentId, updatedTime)
-    );
-  }
-
   /**
    * Atomically updates segment IDs in the cache based on the segments
    * currently present in the metadata store.
@@ -219,9 +150,7 @@ SegmentSyncResult syncSegmentIds(List<SegmentRecord> persistedSegments, DateTime
 
         if (record.isUsed()) {
           // Refresh this used segment if it has been updated in the metadata store
-          final DataSegmentPlus cachedState = intervalSegments.idToUsedSegment.get(segmentId);
-          if (cachedState == null
-              || shouldUpdateCache(cachedState.getUsedStatusLastUpdatedDate(), record.getLastUpdatedTime())) {
+          if (intervalSegments.shouldRefreshUsedSegment(segmentId, record.getLastUpdatedTime())) {
             usedSegmentIdsToRefresh.add(segmentId.toString());
           }
         } else {
@@ -241,14 +170,23 @@ SegmentSyncResult syncSegmentIds(List<SegmentRecord> persistedSegments, DateTime
     });
   }
 
-  SegmentSyncResult syncPendingSegments(List<PendingSegmentRecord> persistedPendingSegments, DateTime syncStartTime)
+  /**
+   * Atomically updates pending segments in the cache based on the segments
+   * currently present in the metadata store.
+   *
+   * @param persistedPendingSegments All pending segments present in the metadata store.
+   * @param syncStartTime            Start time of the current sync
+   * @return Summary of updates made to the cache.
+   */
+  SegmentSyncResult syncPendingSegments(
+      List<PendingSegmentRecord> persistedPendingSegments,
+      DateTime syncStartTime
+  )
   {
     return withWriteLock(() -> {
       int numSegmentsUpdated = 0;
       for (PendingSegmentRecord record : persistedPendingSegments) {
-        final boolean updated = shouldRefreshPendingSegment(record)
-                                && insertPendingSegment(record, false);
-        if (updated) {
+        if (insertPendingSegment(record, false)) {
           ++numSegmentsUpdated;
         }
       }
@@ -264,7 +202,7 @@ SegmentSyncResult syncPendingSegments(List<PendingSegmentRecord> persistedPendin
    * Removes all pending segments which are present in the cache but not present
    * in the metadata store.
    */
-  int removeUnpersistedPendingSegments(Set<String> persistedPendingSegmentIds, DateTime pollStartTime)
+  private int removeUnpersistedPendingSegments(Set<String> persistedPendingSegmentIds, DateTime pollStartTime)
   {
     return withWriteLock(() -> {
       final Set<String> unpersistedSegmentIds =
@@ -284,7 +222,7 @@ && shouldUpdateCache(record.getCreatedDate(), pollStartTime)
    * @param syncStartTime       Start time of the current sync
    * @return Number of unpersisted segments removed from cache.
    */
-  int removeUnpersistedSegments(Set<SegmentId> persistedSegmentIds, DateTime syncStartTime)
+  private int removeUnpersistedSegments(Set<SegmentId> persistedSegmentIds, DateTime syncStartTime)
   {
     return withWriteLock(() -> {
       final Set<SegmentId> unpersistedSegmentIds = new HashSet<>();
@@ -350,11 +288,17 @@ void markCacheSynced()
     });
   }
 
+  /**
+   * Must be accessed within a {@link #withReadLock} method.
+   */
   private SegmentsInInterval readSegmentsFor(Interval interval)
   {
     return intervalToSegments.getOrDefault(interval, SegmentsInInterval.EMPTY);
   }
 
+  /**
+   * Must be accessed within a {@link #withWriteLock} method.
+   */
   private SegmentsInInterval writeSegmentsFor(Interval interval)
   {
     return intervalToSegments.computeIfAbsent(interval, i -> new SegmentsInInterval());
@@ -535,7 +479,8 @@ public int insertSegments(Set<DataSegmentPlus> segments)
     return withWriteLock(() -> {
       int numInsertedSegments = 0;
       for (DataSegmentPlus segmentPlus : segments) {
-        if (addSegment(segmentPlus)) {
+        final Interval interval = segmentPlus.getDataSegment().getInterval();
+        if (writeSegmentsFor(interval).addSegment(segmentPlus)) {
           ++numInsertedSegments;
         }
       }
@@ -553,19 +498,22 @@ public int insertSegmentsWithMetadata(Set<DataSegmentPlus> segments)
   @Override
   public boolean markSegmentAsUnused(SegmentId segmentId, DateTime updateTime)
   {
-    return addUnusedSegmentId(segmentId, updateTime);
+    return writeSegmentsFor(segmentId.getInterval()).addUnusedSegmentId(segmentId, updateTime);
   }
 
   @Override
   public int markSegmentsAsUnused(Set<SegmentId> segmentIds, DateTime updateTime)
   {
-    int updatedCount = 0;
-    for (SegmentId segmentId : segmentIds) {
-      if (addUnusedSegmentId(segmentId, updateTime)) {
-        ++updatedCount;
+    return withWriteLock(() -> {
+      int updatedCount = 0;
+      for (SegmentId segmentId : segmentIds) {
+        final Interval interval = segmentId.getInterval();
+        if (writeSegmentsFor(interval).addUnusedSegmentId(segmentId, updateTime)) {
+          ++updatedCount;
+        }
       }
-    }
-    return updatedCount;
+      return updatedCount;
+    });
   }
 
   @Override
@@ -577,29 +525,38 @@ public int markSegmentsWithinIntervalAsUnused(
   {
     final Set<String> eligibleVersions = versions == null ? null : Set.copyOf(versions);
 
-    int updatedCount = 0;
-    for (DataSegmentPlus segment : findUsedSegmentsPlusOverlappingAnyOf(List.of(interval))) {
-      // Update segments with eligible versions or all versions (if eligibleVersions is null)
-      if ((eligibleVersions == null || eligibleVersions.contains(segment.getDataSegment().getVersion()))
-          && addUnusedSegmentId(segment.getDataSegment().getId(), updateTime)) {
-        ++updatedCount;
+    return withWriteLock(() -> {
+      int updatedCount = 0;
+      for (DataSegmentPlus segmentPlus : findUsedSegmentsPlusOverlappingAnyOf(List.of(interval))) {
+        // Update segments with eligible versions or all versions (if eligibleVersions is null)
+        final DataSegment segment = segmentPlus.getDataSegment();
+        final boolean isEligibleVersion = eligibleVersions == null
+                                          || eligibleVersions.contains(segment.getVersion());
+        if (isEligibleVersion
+            && writeSegmentsFor(segment.getInterval()).addUnusedSegmentId(segment.getId(), updateTime)) {
+          ++updatedCount;
+        }
       }
-    }
 
-    return updatedCount;
+      return updatedCount;
+    });
   }
 
   @Override
   public int markAllSegmentsAsUnused(DateTime updateTime)
   {
-    int updatedCount = 0;
-    for (DataSegmentPlus segment : findUsedSegmentsPlusOverlappingAnyOf(List.of())) {
-      if (addUnusedSegmentId(segment.getDataSegment().getId(), updateTime)) {
-        ++updatedCount;
+    return withWriteLock(() -> {
+      int updatedCount = 0;
+      for (DataSegmentPlus segmentPlus : findUsedSegmentsPlusOverlappingAnyOf(List.of())) {
+        final DataSegment segment = segmentPlus.getDataSegment();
+        if (writeSegmentsFor(segment.getInterval())
+            .addUnusedSegmentId(segment.getId(), updateTime)) {
+          ++updatedCount;
+        }
       }
-    }
 
-    return updatedCount;
+      return updatedCount;
+    });
   }
 
   @Override
@@ -786,6 +743,49 @@ private void updateMaxUnusedId(SegmentId segmentId)
           .merge(segmentId.getVersion(), segmentId.getPartitionNum(), Math::max);
     }
 
+    /**
+     * Adds or updates the given segment in the cache.
+     *
+     * @return true if the segment was updated in the cache, false if the segment
+     * was left unchanged in the cache.
+     */
+    boolean addSegment(DataSegmentPlus segmentPlus)
+    {
+      if (Boolean.TRUE.equals(segmentPlus.getUsed())) {
+        return addUsedSegment(segmentPlus);
+      } else {
+        return addUnusedSegmentId(
+            segmentPlus.getDataSegment().getId(),
+            segmentPlus.getUsedStatusLastUpdatedDate()
+        );
+      }
+    }
+
+    /**
+     * Adds or updates a used segment in the cache.
+     */
+    private boolean addUsedSegment(DataSegmentPlus segmentPlus)
+    {
+      final DataSegment segment = segmentPlus.getDataSegment();
+      final SegmentId segmentId = segment.getId();
+
+      if (!shouldRefreshUsedSegment(segmentId, segmentPlus.getUsedStatusLastUpdatedDate())) {
+        return false;
+      }
+
+      idToUsedSegment.put(segmentId, segmentPlus);
+      unusedSegmentIdToUpdatedTime.remove(segment.getId());
+      return true;
+    }
+
+    /**
+     * Adds or updates an unused segment in the cache.
+     *
+     * @param updatedTime Last updated time of this segment as persisted in the
+     *                    metadata store. This value can be null for segments
+     *                    persisted to the metadata store before the column
+     *                    used_status_last_updated was added to the segments table.
+     */
     private boolean addUnusedSegmentId(SegmentId segmentId, @Nullable DateTime updatedTime)
     {
       idToUsedSegment.remove(segmentId);
@@ -799,5 +799,24 @@ private boolean addUnusedSegmentId(SegmentId segmentId, @Nullable DateTime updat
         return false;
       }
     }
+
+    private boolean shouldRefreshUnusedSegment(SegmentId segmentId, DateTime newUpdateTime)
+    {
+      return !unusedSegmentIdToUpdatedTime.containsKey(segmentId)
+             || shouldUpdateCache(unusedSegmentIdToUpdatedTime.get(segmentId), newUpdateTime);
+    }
+
+    private boolean shouldRefreshUsedSegment(SegmentId segmentId, DateTime newUpdateTime)
+    {
+      final DataSegmentPlus usedSegment = idToUsedSegment.get(segmentId);
+
+      if (usedSegment == null) {
+        // Do not refresh the segment if it has recently been marked as unused in the cache
+        return shouldRefreshUnusedSegment(segmentId, newUpdateTime);
+      } else {
+        // Refresh the used segment if the entry in the cache is stale
+        return shouldUpdateCache(usedSegment.getUsedStatusLastUpdatedDate(), newUpdateTime);
+      }
+    }
   }
 }

server/src/main/java/org/apache/druid/metadata/segment/cache/HeapMemorySegmentMetadataCache.java

@@ -67,6 +67,9 @@
 /**
  * In-memory implementation of {@link SegmentMetadataCache}.
  * <p>
+ * Non-leader Overlords also keep polling the metadata store to keep the cache
+ * up-to-date in case leadership changes.
+ * <p>
  * The map {@link #datasourceToSegmentCache} contains the cache for each datasource.
  * Items are only added to this map and never removed. This is to avoid handling
  * race conditions where a thread has invoked {@link #getDatasource} but hasn't
@@ -606,9 +609,8 @@ private void updatePendingSegmentsInCache(
   }
 
   /**
-   * Retrieves all pending segments from metadata store and updates the cache if
-   * {@link HeapMemoryDatasourceSegmentCache#shouldRefreshPendingSegment} is
-   * true for it.
+   * Retrieves all pending segments from metadata store and populates them in
+   * the respective {@link DatasourceSegmentSummary}.
    */
   private void retrieveAllPendingSegments(
       Map<String, DatasourceSegmentSummary> datasourceToSummary

server/src/main/java/org/apache/druid/metadata/segment/cache/Metric.java

@@ -50,7 +50,7 @@ private Metric()
   /**
    * Number of transactions performed on the cache for a datasource.
    */
-  public static final String TRANSACTION_COUNT = "transactions";
+  public static final String TRANSACTION_COUNT = METRIC_NAME_PREFIX + "transactions";
 
 
   // CACHE SYNC TIME METRICS

server/src/main/java/org/apache/druid/metadata/segment/cache/SegmentMetadataCache.java

@@ -20,7 +20,12 @@
 package org.apache.druid.metadata.segment.cache;
 
 /**
- * Cache for metadata of pending segments and committed segments.
+ * Cache for metadata of pending segments and committed segments maintained by
+ * the Overlord to improve performance of segment allocation and other task actions.
+ * <p>
+ * Not to be confused with {@link org.apache.druid.segment.metadata.AbstractSegmentMetadataCache}
+ * which is used by Brokers to cache row signature, number of rows, etc. to aid
+ * with Druid query performance.
  */
 public interface SegmentMetadataCache
 {