Cache segment metadata on the Overlord to speed up segment allocation and other task actions

docs/configuration/index.md

@@ -1168,6 +1168,15 @@ If autoscaling is enabled, you can set these additional configs:
 
 The `druid.supervisor.idleConfig.*` specification in the Overlord runtime properties defines the default behavior for the entire cluster. See [Idle Configuration in Kafka Supervisor IOConfig](../ingestion/kinesis-ingestion.md#io-configuration) to override it for an individual supervisor.
 
+##### Segment metadata cache (EXPERIMENTAL)
+
+The following properties pertain to segment metadata caching on the Overlord that may be used to speed up segment allocation and other metadata operations.
+
+|Property|Description|Default|
+|--------|-----------|-------|
+|`druid.manager.segments.useCache`|If `true`, segment metadata is cached on the Overlord to speed up metadata operations.|false|
+|`druid.manager.segments.pollDuration`|Duration (in ISO 8601 format) between successive syncs of the cache with the metadata store. This property is used only when `druid.manager.segments.useCache` is `true`.|`PT1M` (1 minute)|
+
 #### Overlord dynamic configuration
 
 The Overlord has dynamic configurations to tune how Druid assigns tasks to workers.

docs/operations/metrics.md

@@ -313,6 +313,25 @@ If the JVM does not support CPU time measurement for the current thread, `ingest
 |`worker/task/success/count`|Number of tasks that succeeded on an indexer during the emission period. This metric is only available if the `WorkerTaskCountStatsMonitor` module is included.|`dataSource`|Varies|
 |`worker/task/running/count`|Number of tasks running on an indexer per emission period. This metric is only available if the `WorkerTaskCountStatsMonitor` module is included.|`dataSource`|Varies|
 
+### Segment metadata cache
+
+The following metrics are emitted only when [segment metadata caching](../configuration/index.md#segment-metadata-cache-experimental) is enabled on the Overlord.
+
+|Metric|Description|Dimensions|
+|------|-----------|----------|
+|`segment/used/count`|Number of used segments currently present in the metadata store.|`dataSource`|
+|`segment/unused/count`|Number of unused segments currently present in the metadata store.|`dataSource`|
+|`segment/pending/count`|Number of pending segments currently present in the metadata store.|`dataSource`|
+|`segment/metadataCache/sync/time`|Number of milliseconds taken for the cache to sync with the metadata store.||
+|`segment/metadataCache/deleted`|Total number of segments deleted from the cache during the latest sync.||
+|`segment/metadataCache/skipped`|Total number of unparseable segment records that were skipped in the latest sync.||
+|`segment/metadataCache/used/stale`|Number of used segments in the cache which are out-of-date and need to be refreshed.|`dataSource`|
+|`segment/metadataCache/used/updated`|Number of used segments updated in the cache during the latest sync.|`dataSource`|
+|`segment/metadataCache/unused/updated`|Number of unused segments updated in the cache during the latest sync.|`dataSource`|
+|`segment/metadataCache/pending/deleted`|Number of pending segments deleted from the cache during the latest sync.|`dataSource`|
+|`segment/metadataCache/pending/updated`|Number of pending segments updated in the cache during the latest sync.|`dataSource`|
+|`segment/metadataCache/pending/skipped`|Number of unparseable pending segment records that were skipped in the latest sync.|`dataSource`|
+
 ## Shuffle metrics (Native parallel task)
 
 The shuffle metrics can be enabled by adding `org.apache.druid.indexing.worker.shuffle.ShuffleMonitor` in `druid.monitoring.monitors`.

extensions-contrib/materialized-view-maintenance/src/test/java/org/apache/druid/indexing/materializedview/MaterializedViewSupervisorTest.java

@@ -46,13 +46,16 @@
 import org.apache.druid.metadata.MetadataSupervisorManager;
 import org.apache.druid.metadata.SqlSegmentsMetadataManager;
 import org.apache.druid.metadata.TestDerbyConnector;
+import org.apache.druid.metadata.segment.SqlSegmentMetadataTransactionFactory;
+import org.apache.druid.metadata.segment.cache.NoopSegmentMetadataCache;
 import org.apache.druid.query.aggregation.AggregatorFactory;
 import org.apache.druid.query.aggregation.LongSumAggregatorFactory;
 import org.apache.druid.segment.TestHelper;
 import org.apache.druid.segment.indexing.DataSchema;
 import org.apache.druid.segment.metadata.CentralizedDatasourceSchemaConfig;
 import org.apache.druid.segment.metadata.SegmentSchemaManager;
 import org.apache.druid.segment.realtime.ChatHandlerProvider;
+import org.apache.druid.server.coordinator.simulate.TestDruidLeaderSelector;
 import org.apache.druid.server.security.AuthorizerMapper;
 import org.apache.druid.timeline.DataSegment;
 import org.apache.druid.timeline.partition.HashBasedNumberedShardSpec;
@@ -105,6 +108,13 @@ public void setUp()
         derbyConnector
     );
     indexerMetadataStorageCoordinator = new IndexerSQLMetadataStorageCoordinator(
+        new SqlSegmentMetadataTransactionFactory(
+            objectMapper,
+            derbyConnectorRule.metadataTablesConfigSupplier().get(),
+            derbyConnector,
+            new TestDruidLeaderSelector(),
+            new NoopSegmentMetadataCache()
+        ),
         objectMapper,
         derbyConnectorRule.metadataTablesConfigSupplier().get(),
         derbyConnector,

extensions-contrib/materialized-view-selection/src/test/java/org/apache/druid/query/materializedview/DatasourceOptimizerTest.java

@@ -42,6 +42,8 @@
 import org.apache.druid.java.util.http.client.HttpClient;
 import org.apache.druid.metadata.IndexerSQLMetadataStorageCoordinator;
 import org.apache.druid.metadata.TestDerbyConnector;
+import org.apache.druid.metadata.segment.SqlSegmentMetadataTransactionFactory;
+import org.apache.druid.metadata.segment.cache.NoopSegmentMetadataCache;
 import org.apache.druid.query.Query;
 import org.apache.druid.query.QueryRunnerFactoryConglomerate;
 import org.apache.druid.query.QueryRunnerTestHelper;
@@ -56,6 +58,7 @@
 import org.apache.druid.segment.realtime.appenderator.SegmentSchemas;
 import org.apache.druid.server.coordination.DruidServerMetadata;
 import org.apache.druid.server.coordination.ServerType;
+import org.apache.druid.server.coordinator.simulate.TestDruidLeaderSelector;
 import org.apache.druid.server.initialization.ZkPathsConfig;
 import org.apache.druid.server.metrics.NoopServiceEmitter;
 import org.apache.druid.timeline.DataSegment;
@@ -112,6 +115,13 @@ public void setUp() throws Exception
     );
 
     metadataStorageCoordinator = new IndexerSQLMetadataStorageCoordinator(
+        new SqlSegmentMetadataTransactionFactory(
+            jsonMapper,
+            derbyConnectorRule.metadataTablesConfigSupplier().get(),
+            derbyConnector,
+            new TestDruidLeaderSelector(),
+            new NoopSegmentMetadataCache()
+        ),
         jsonMapper,
         derbyConnectorRule.metadataTablesConfigSupplier().get(),
         derbyConnector,

extensions-core/mysql-metadata-storage/src/test/java/org/apache/druid/metadata/storage/mysql/MySQLMetadataStorageModuleTest.java

@@ -26,6 +26,8 @@
 import com.google.inject.Key;
 import com.google.inject.Module;
 import com.google.inject.Provides;
+import org.apache.druid.client.indexing.IndexingService;
+import org.apache.druid.discovery.DruidLeaderSelector;
 import org.apache.druid.guice.GuiceInjectors;
 import org.apache.druid.guice.JsonConfigProvider;
 import org.apache.druid.guice.JsonConfigurator;
@@ -129,6 +131,14 @@ public ServiceEmitter getEmitter()
               {
                 return new ServiceEmitter("test", "localhost", new NoopEmitter());
               }
+
+              @Provides
+              @IndexingService
+              public DruidLeaderSelector getLeaderSelector()
+              {
+                // A provider for DruidLeaderSelector is needed by SqlSegmentMetadataTransactionFactory
+                return null;
+              }
             }
         )
     );

indexing-service/src/main/java/org/apache/druid/indexing/overlord/DruidOverlord.java

@@ -41,6 +41,7 @@
 import org.apache.druid.java.util.common.lifecycle.LifecycleStop;
 import org.apache.druid.java.util.emitter.EmittingLogger;
 import org.apache.druid.java.util.emitter.service.ServiceEmitter;
+import org.apache.druid.metadata.segment.cache.SegmentMetadataCache;
 import org.apache.druid.server.DruidNode;
 import org.apache.druid.server.coordinator.CoordinatorOverlordServiceConfig;
 
@@ -89,6 +90,7 @@ public DruidOverlord(
       final OverlordDutyExecutor overlordDutyExecutor,
       @IndexingService final DruidLeaderSelector overlordLeaderSelector,
       final SegmentAllocationQueue segmentAllocationQueue,
+      final SegmentMetadataCache segmentMetadataCache,
       final CompactionScheduler compactionScheduler,
       final ScheduledBatchTaskManager scheduledBatchTaskManager,
       final ObjectMapper mapper,
@@ -141,6 +143,7 @@ public void becomeLeader()
                 @Override
                 public void start()
                 {
+                  segmentMetadataCache.becomeLeader();
                   segmentAllocationQueue.becomeLeader();
                   taskMaster.becomeHalfLeader(taskRunner, taskQueue);
                 }
@@ -150,6 +153,7 @@ public void stop()
                 {
                   taskMaster.stopBeingLeader();
                   segmentAllocationQueue.stopBeingLeader();
+                  segmentMetadataCache.stopBeingLeader();
                 }
               }
           );

indexing-service/src/test/java/org/apache/druid/indexing/common/actions/RetrieveSegmentsActionsTest.java

@@ -29,8 +29,8 @@
 import org.apache.druid.timeline.partition.NoneShardSpec;
 import org.joda.time.Interval;
 import org.junit.Assert;
-import org.junit.BeforeClass;
-import org.junit.ClassRule;
+import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 
 import java.util.HashSet;
@@ -42,15 +42,15 @@ public class RetrieveSegmentsActionsTest
   private static final String UNUSED_V0 = "v0";
   private static final String UNUSED_V1 = "v1";
 
-  @ClassRule
-  public static TaskActionTestKit actionTestKit = new TaskActionTestKit();
+  @Rule
+  public TaskActionTestKit actionTestKit = new TaskActionTestKit();
 
   private static Task task;
   private static Set<DataSegment> expectedUnusedSegments;
   private static Set<DataSegment> expectedUsedSegments;
 
-  @BeforeClass
-  public static void setup()
+  @Before
+  public void setup()
   {
     task = NoopTask.create();
 
@@ -99,6 +99,7 @@ private static DataSegment createSegment(Interval interval, String version)
   @Test
   public void testRetrieveUsedSegmentsAction()
   {
+    actionTestKit.syncSegmentMetadataCache();
     final RetrieveUsedSegmentsAction action =
         new RetrieveUsedSegmentsAction(task.getDataSource(), ImmutableList.of(INTERVAL));
     final Set<DataSegment> observedUsedSegments = new HashSet<>(action.perform(task, actionTestKit.getTaskActionToolbox()));

indexing-service/src/test/java/org/apache/druid/indexing/common/actions/SegmentAllocateActionTest.java

@@ -56,6 +56,7 @@
 import org.joda.time.Period;
 import org.junit.After;
 import org.junit.Assert;
+import org.junit.Assume;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
@@ -93,28 +94,32 @@ public class SegmentAllocateActionTest
 
   private SegmentAllocationQueue allocationQueue;
 
-  @Parameterized.Parameters(name = "granularity = {0}, useBatch = {1}, skipSegmentPayloadFetchForAllocation = {2}")
+  @Parameterized.Parameters(name = "lock={0}, useBatch={1}, useSegmentCache={2}, reduceMetadataIO={3}")
   public static Iterable<Object[]> constructorFeeder()
   {
+    // reduceMetadataIO is applicable only with batch allocation
     return ImmutableList.of(
-        new Object[]{LockGranularity.SEGMENT, true, true},
-        new Object[]{LockGranularity.SEGMENT, true, false},
-        new Object[]{LockGranularity.SEGMENT, false, false},
-        new Object[]{LockGranularity.TIME_CHUNK, true, true},
-        new Object[]{LockGranularity.TIME_CHUNK, true, false},
-        new Object[]{LockGranularity.TIME_CHUNK, false, false}
+        new Object[]{LockGranularity.SEGMENT, true, true, true},
+        new Object[]{LockGranularity.SEGMENT, true, false, false},
+        new Object[]{LockGranularity.SEGMENT, false, false, false},
+        new Object[]{LockGranularity.TIME_CHUNK, true, true, true},
+        new Object[]{LockGranularity.TIME_CHUNK, true, false, false},
+        new Object[]{LockGranularity.TIME_CHUNK, false, false, false},
+        new Object[]{LockGranularity.TIME_CHUNK, false, true, false}
     );
   }
 
   public SegmentAllocateActionTest(
       LockGranularity lockGranularity,
       boolean useBatch,
+      boolean useSegmentMetadataCache,
       boolean skipSegmentPayloadFetchForAllocation
   )
   {
     this.lockGranularity = lockGranularity;
     this.useBatch = useBatch;
     this.taskActionTestKit.setSkipSegmentPayloadFetchForAllocation(skipSegmentPayloadFetchForAllocation);
+    this.taskActionTestKit.setUseSegmentMetadataCache(useSegmentMetadataCache);
   }
 
   @Before
@@ -141,9 +146,7 @@ public void tearDown()
   @Test
   public void testManySegmentsSameInterval_noLineageCheck() throws Exception
   {
-    if (lockGranularity == LockGranularity.SEGMENT) {
-      return;
-    }
+    Assume.assumeTrue(lockGranularity == LockGranularity.TIME_CHUNK);
 
     final Task task = NoopTask.create();
     final int numTasks = 2;
@@ -1170,7 +1173,7 @@ public void testSegmentIdMustNotBeReused()
     // Allocate another id and ensure that it doesn't exist in the druid_segments table
     final SegmentIdWithShardSpec theId =
         allocate(task1, DateTimes.nowUtc(), Granularities.NONE, Granularities.ALL, "seq", "3");
-    Assert.assertNull(coordinator.retrieveSegmentForId(theId.asSegmentId().toString(), true));
+    Assert.assertNull(coordinator.retrieveSegmentForId(theId.getDataSource(), theId.asSegmentId().toString()));
 
     lockbox.unlock(task1, Intervals.ETERNITY);
   }

indexing-service/src/test/java/org/apache/druid/indexing/common/actions/SegmentAllocationQueueTest.java

@@ -59,15 +59,22 @@ public class SegmentAllocationQueueTest
 
   private final boolean reduceMetadataIO;
 
-  @Parameterized.Parameters(name = "reduceMetadataIO = {0}")
+  @Parameterized.Parameters(name = "reduceMetadataIO = {0}, useSegmentCache = {1}")
   public static Object[][] getTestParameters()
   {
-    return new Object[][]{{true}, {false}};
+    return new Object[][]{
+        {true, true},
+        {true, false},
+        {false, true},
+        {false, false}
+    };
   }
 
-  public SegmentAllocationQueueTest(boolean reduceMetadataIO)
+  public SegmentAllocationQueueTest(boolean reduceMetadataIO, boolean useSegmentMetadataCache)
   {
     this.reduceMetadataIO = reduceMetadataIO;
+
+    taskActionTestKit.setUseSegmentMetadataCache(useSegmentMetadataCache);
   }
 
   @Before