perf: speed-up coordinator segment metadata cache delta syncs

extensions-core/mysql-metadata-storage/src/main/java/org/apache/druid/metadata/storage/mysql/MySQLConnector.java

@@ -191,6 +191,13 @@ public String limitClause(int limit)
     return String.format(Locale.ENGLISH, "LIMIT %d", limit);
   }
 
+  @Override
+  protected String getDropIndexStatement(String indexName, String tableName)
+  {
+    // MySQL requires the target table in a DROP INDEX statement.
+    return StringUtils.format("DROP INDEX %s ON %s", indexName, tableName);
+  }
+
   @Override
   public boolean tableExists(Handle handle, String tableName)
   {

server/src/main/java/org/apache/druid/client/DataSourcesSnapshot.java

@@ -35,6 +35,7 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -91,6 +92,67 @@ public static DataSourcesSnapshot fromUsedSegments(Iterable<DataSegment> segment
     );
   }
 
+  /**
+   * Builds a new snapshot from this one by recomputing only the datasources that
+   * changed, reusing this snapshot's {@link ImmutableDruidDataSource}, timeline,
+   * and overshadowed-segment computation for every unchanged datasource.
+   * <p>
+   * Note that a changed datasource has its entire timeline rebuilt even if only a
+   * single one of its segments changed; the saving is that untouched datasources
+   * are reused as is. Overshadowing is computed per-datasource, so reusing the
+   * prior overshadowed segments for unchanged datasources is correct. The result
+   * is identical to a full {@link #fromUsedSegments} rebuild of the same final
+   * state, at a cost proportional to the changed datasources rather than all of them.
+   *
+   * @param datasourcesToRefresh Datasource → its complete current set of used
+   *                             segments, for each datasource that changed. Sets
+   *                             must be non-empty; removals are conveyed via
+   *                             {@code removedDataSources}.
+   * @param removedDataSources   Datasources whose caches are now empty/gone.
+   * @param snapshotTime         Time of this snapshot (poll start time).
+   */
+  public DataSourcesSnapshot updateSnapshotForDataSources(
+      Map<String, Set<DataSegment>> datasourcesToRefresh,
+      Set<String> removedDataSources,
+      DateTime snapshotTime
+  )
+  {
+    final Map<String, String> properties = Map.of("created", snapshotTime.toString());
+    final Map<String, ImmutableDruidDataSource> dataSources = new HashMap<>(dataSourcesWithAllUsedSegments);
+    final Map<String, SegmentTimeline> timelines = new HashMap<>(usedSegmentsTimelinesPerDataSource);
+
+    final Set<String> changedDataSources = new HashSet<>(removedDataSources);
+    changedDataSources.addAll(datasourcesToRefresh.keySet());
+
+    removedDataSources.forEach(ds -> {
+      dataSources.remove(ds);
+      timelines.remove(ds);
+    });
+    datasourcesToRefresh.forEach((ds, segments) -> {
+      dataSources.put(ds, new ImmutableDruidDataSource(ds, properties, segments));
+      timelines.put(ds, SegmentTimeline.forSegments(segments));
+    });
+
+    // Reuse prior overshadowed segments for unchanged datasources; recompute only the changed ones.
+    final List<DataSegment> overshadowed = new ArrayList<>();
+    for (DataSegment segment : overshadowedSegments) {
+      if (!changedDataSources.contains(segment.getDataSource())) {
+        overshadowed.add(segment);
+      }
+    }
+    for (String ds : datasourcesToRefresh.keySet()) {
+      final ImmutableDruidDataSource dataSource = dataSources.get(ds);
+      final SegmentTimeline timeline = timelines.get(ds);
+      for (DataSegment segment : dataSource.getSegments()) {
+        if (timeline.isOvershadowed(segment)) {
+          overshadowed.add(segment);
+        }
+      }
+    }
+
+    return new DataSourcesSnapshot(snapshotTime, dataSources, timelines, overshadowed);
+  }
+
   private final DateTime snapshotTime;
   private final Map<String, ImmutableDruidDataSource> dataSourcesWithAllUsedSegments;
   private final Map<String, SegmentTimeline> usedSegmentsTimelinesPerDataSource;
@@ -110,6 +172,23 @@ private DataSourcesSnapshot(
     this.overshadowedSegments = ImmutableSet.copyOf(determineOvershadowedSegments());
   }
 
+  /**
+   * Constructor used by {@link #updateSnapshotForDataSources} where timelines and
+   * overshadowed segments have already been computed incrementally.
+   */
+  private DataSourcesSnapshot(
+      DateTime snapshotTime,
+      Map<String, ImmutableDruidDataSource> dataSourcesWithAllUsedSegments,
+      Map<String, SegmentTimeline> usedSegmentsTimelinesPerDataSource,
+      Collection<DataSegment> overshadowedSegments
+  )
+  {
+    this.snapshotTime = snapshotTime;
+    this.dataSourcesWithAllUsedSegments = dataSourcesWithAllUsedSegments;
+    this.usedSegmentsTimelinesPerDataSource = usedSegmentsTimelinesPerDataSource;
+    this.overshadowedSegments = ImmutableSet.copyOf(overshadowedSegments);
+  }
+
   /**
    * Time when this snapshot was taken. Since polling segments from the database
    * may be a slow operation, this represents the poll start time.

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

@@ -389,11 +389,6 @@ tableName, getPayloadType(), getQuoteString(), getCollation()
         )
     );
 
-    createIndex(
-        tableName,
-        "IDX_%S_USED",
-        List.of("used")
-    );
     createIndex(
         tableName,
         "IDX_%S_DATASOURCE_USED_END_START",
@@ -404,6 +399,17 @@ tableName, getPayloadType(), getQuoteString(), getCollation()
             "start"
         )
     );
+    // Covering index for the used-segment ID scan performed on every metadata
+    // cache sync (SELECT id, dataSource, used_status_last_updated WHERE used=true).
+    // Includes id explicitly so the scan is index-only on all backends (not only
+    // engines like InnoDB that append the primary key to secondary indexes).
+    // Its leading 'used' column also serves plain 'WHERE used = ?' lookups, so a
+    // separate IDX_%S_USED index is not created.
+    createIndex(
+        tableName,
+        "IDX_%S_USED_USLU_DATASOURCE",
+        List.of("used", "used_status_last_updated", "dataSource", "id")
+    );
   }
 
   private void createUpgradeSegmentsTable(final String tableName)
@@ -663,6 +669,16 @@ protected void alterSegmentTable()
         "IDX_%S_DATASOURCE_UPGRADED_FROM_SEGMENT_ID",
         List.of("dataSource", "upgraded_from_segment_id")
     );
+    // Migration for existing tables: covering index backing the used-segment ID
+    // scan on every cache sync (see createSegmentTable).
+    createIndex(
+        tableName,
+        "IDX_%S_USED_USLU_DATASOURCE",
+        List.of("used", "used_status_last_updated", "dataSource", "id")
+    );
+    // The covering index above leads with 'used', so it supersedes the single-column
+    // IDX_%S_USED. Drop the now-redundant index on existing tables.
+    dropIndex(tableName, "IDX_%S_USED", List.of("used"));
   }
 
   @Override
@@ -1256,6 +1272,59 @@ public void createIndex(
     }
   }
 
+  /**
+   * Drops an index on {@code tableName} if it exists, under either the full or
+   * short naming convention (see {@link #createIndex}). No-op if absent, so it is
+   * safe to call on every startup.
+   *
+   * @param tableName           Name of the table the index is on
+   * @param fullIndexNameFormat Same format string that was passed to {@link #createIndex}
+   * @param indexCols           Same columns that were passed to {@link #createIndex}
+   */
+  public void dropIndex(
+      final String tableName,
+      final String fullIndexNameFormat,
+      final List<String> indexCols
+  )
+  {
+    final Set<String> createdIndexSet = getIndexOnTable(tableName);
+    final String shortIndexName = generateShortIndexName(tableName, indexCols);
+    final String fullIndexName = StringUtils.toUpperCase(StringUtils.format(fullIndexNameFormat, tableName));
+
+    final String indexName;
+    if (createdIndexSet.contains(fullIndexName)) {
+      indexName = fullIndexName;
+    } else if (createdIndexSet.contains(shortIndexName)) {
+      indexName = shortIndexName;
+    } else {
+      log.info("Index[%s] on table[%s] does not exist, skipping drop.", fullIndexName, tableName);
+      return;
+    }
+
+    try {
+      retryWithHandle(
+          (HandleCallback<Void>) handle -> {
+            final String dropSQL = getDropIndexStatement(indexName, tableName);
+            log.info("Dropping index[%s] on table[%s] using SQL[%s].", indexName, tableName, dropSQL);
+            handle.execute(dropSQL);
+            return null;
+          }
+      );
+    }
+    catch (Exception e) {
+      log.warn(e, "Could not drop index[%s] on table[%s]", indexName, tableName);
+    }
+  }
+
+  /**
+   * SQL to drop an index. Standard SQL / Derby / PostgreSQL use {@code DROP INDEX <name>};
+   * MySQL requires the {@code ON <table>} clause (see {@code MySQLConnector}).
+   */
+  protected String getDropIndexStatement(String indexName, String tableName)
+  {
+    return StringUtils.format("DROP INDEX %s", indexName);
+  }
+
   /**
    * Checks table metadata to determine if the given column exists in the table.
    *

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

@@ -38,6 +38,7 @@
 import java.util.Objects;
 import java.util.Set;
 import java.util.TreeMap;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
@@ -68,6 +69,15 @@ class HeapMemoryDatasourceSegmentCache extends ReadWriteCache implements AutoClo
    */
   private final AtomicInteger references = new AtomicInteger(0);
 
+  /**
+   * Set when a write-through transaction mutates this cache directly (via
+   * {@link HeapMemorySegmentMetadataCache#writeCacheForDataSource}), so that the
+   * incremental snapshot rebuild in the next sync refreshes this datasource even
+   * though the metadata store and cache already agree (the DB diff would report
+   * no change). Cleared while the snapshot is being rebuilt.
+   */
+  private final AtomicBoolean hasNewWrites = new AtomicBoolean(false);
+
   HeapMemoryDatasourceSegmentCache(String dataSource)
   {
     super(true);
@@ -116,6 +126,46 @@ boolean isBeingUsedByTransaction()
     return references.get() > 0;
   }
 
+  /**
+   * Marks that a write-through transaction has mutated this cache directly, so
+   * the next snapshot rebuild must refresh this datasource. Called under the
+   * write lock held by {@link HeapMemorySegmentMetadataCache#writeCacheForDataSource}.
+   */
+  void markHasNewWrites()
+  {
+    hasNewWrites.set(true);
+  }
+
+  /**
+   * If a write-through mutation occurred since the last snapshot rebuild, returns
+   * all used segments and clears the flag (atomically under the write lock);
+   * otherwise returns null. Used to catch datasources whose cache diverged from
+   * the last published snapshot without a corresponding metadata-store diff.
+   */
+  @Nullable
+  Set<DataSegment> getUsedSegmentsIfHasNewWrites()
+  {
+    return withWriteLock(() -> {
+      if (!hasNewWrites.getAndSet(false)) {
+        return null;
+      }
+      return findUsedSegmentsOverlappingAnyOf(List.of());
+    });
+  }
+
+  /**
+   * Returns all used segments and clears the new-writes flag, atomically under
+   * the write lock. Used for datasources already being refreshed from the
+   * metadata-store diff, so a concurrent write-through is not lost.
+   */
+  Set<DataSegment> getUsedSegmentsAndClearNewWrites()
+  {
+    return withWriteLock(() -> {
+      hasNewWrites.set(false);
+      return findUsedSegmentsOverlappingAnyOf(List.of());
+    });
+  }
+
   /**
    * Checks if a record in the cache needs to be updated.
    *