diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskLockbox.java b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskLockbox.java index 5d71940d4704..1a1369ba1e52 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskLockbox.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskLockbox.java @@ -1242,7 +1242,10 @@ public void remove(final Task task) idsInSameGroup.remove(task.getId()); if (idsInSameGroup.isEmpty()) { final int pendingSegmentsDeleted - = metadataStorageCoordinator.deletePendingSegmentsForTaskAllocatorId(taskAllocatorId); + = metadataStorageCoordinator.deletePendingSegmentsForTaskAllocatorId( + task.getDataSource(), + taskAllocatorId + ); log.info( "Deleted [%d] entries from pendingSegments table for pending segments group [%s] with APPEND locks.", pendingSegmentsDeleted, taskAllocatorId diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskLockboxTest.java b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskLockboxTest.java index 3af74235e823..3c95709a9fca 100644 --- a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskLockboxTest.java +++ b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskLockboxTest.java @@ -1948,8 +1948,10 @@ public void testCleanupOnUnlock() // Only the replaceTask should attempt a delete on the upgradeSegments table EasyMock.expect(coordinator.deleteUpgradeSegmentsForTask(replaceTask.getId())).andReturn(0).once(); // Any task may attempt pending segment clean up - EasyMock.expect(coordinator.deletePendingSegmentsForTaskAllocatorId(replaceTask.getId())).andReturn(0).once(); - EasyMock.expect(coordinator.deletePendingSegmentsForTaskAllocatorId(appendTask.getId())).andReturn(0).once(); + EasyMock.expect(coordinator.deletePendingSegmentsForTaskAllocatorId(replaceTask.getDataSource(), replaceTask.getId())) + .andReturn(0).once(); + EasyMock.expect(coordinator.deletePendingSegmentsForTaskAllocatorId(appendTask.getDataSource(), appendTask.getId())) + .andReturn(0).once(); EasyMock.replay(coordinator); final TaskLockbox taskLockbox = new TaskLockbox(taskStorage, coordinator); diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/test/TestIndexerMetadataStorageCoordinator.java b/indexing-service/src/test/java/org/apache/druid/indexing/test/TestIndexerMetadataStorageCoordinator.java index 6c4f556133ec..3aceae494c68 100644 --- a/indexing-service/src/test/java/org/apache/druid/indexing/test/TestIndexerMetadataStorageCoordinator.java +++ b/indexing-service/src/test/java/org/apache/druid/indexing/test/TestIndexerMetadataStorageCoordinator.java @@ -297,7 +297,7 @@ public int deleteUpgradeSegmentsForTask(final String taskId) } @Override - public int deletePendingSegmentsForTaskAllocatorId(final String taskGroup) + public int deletePendingSegmentsForTaskAllocatorId(final String datasource, final String taskAllocatorId) { throw new UnsupportedOperationException(); } diff --git a/server/src/main/java/org/apache/druid/indexing/overlord/IndexerMetadataStorageCoordinator.java b/server/src/main/java/org/apache/druid/indexing/overlord/IndexerMetadataStorageCoordinator.java index aea2674f6b85..da6dd9ffd951 100644 --- a/server/src/main/java/org/apache/druid/indexing/overlord/IndexerMetadataStorageCoordinator.java +++ b/server/src/main/java/org/apache/druid/indexing/overlord/IndexerMetadataStorageCoordinator.java @@ -492,10 +492,11 @@ SegmentPublishResult commitMetadataOnly( /** * Delete pending segment for a give task group after all the tasks belonging to it have completed. + * @param datasource datasource of the task * @param taskAllocatorId task id / task group / replica group for an appending task * @return number of pending segments deleted from the metadata store */ - int deletePendingSegmentsForTaskAllocatorId(String taskAllocatorId); + int deletePendingSegmentsForTaskAllocatorId(String datasource, String taskAllocatorId); /** * Fetches all the pending segments of the datasource that overlap with a given interval. diff --git a/server/src/main/java/org/apache/druid/metadata/IndexerSQLMetadataStorageCoordinator.java b/server/src/main/java/org/apache/druid/metadata/IndexerSQLMetadataStorageCoordinator.java index c5a36656c9a2..618173a5db77 100644 --- a/server/src/main/java/org/apache/druid/metadata/IndexerSQLMetadataStorageCoordinator.java +++ b/server/src/main/java/org/apache/druid/metadata/IndexerSQLMetadataStorageCoordinator.java @@ -2813,17 +2813,18 @@ public DataSegment retrieveSegmentForId(final String id, boolean includeUnused) } @Override - public int deletePendingSegmentsForTaskAllocatorId(final String pendingSegmentsGroup) + public int deletePendingSegmentsForTaskAllocatorId(final String datasource, final String taskAllocatorId) { return connector.getDBI().inTransaction( (handle, status) -> handle .createStatement( StringUtils.format( - "DELETE FROM %s WHERE task_allocator_id = :task_allocator_id", + "DELETE FROM %s WHERE dataSource = :dataSource AND task_allocator_id = :task_allocator_id", dbTables.getPendingSegmentsTable() ) ) - .bind("task_allocator_id", pendingSegmentsGroup) + .bind("dataSource", datasource) + .bind("task_allocator_id", taskAllocatorId) .execute() ); } diff --git a/server/src/main/java/org/apache/druid/metadata/SQLMetadataConnector.java b/server/src/main/java/org/apache/druid/metadata/SQLMetadataConnector.java index 94452de00ee1..cb85548d5ba9 100644 --- a/server/src/main/java/org/apache/druid/metadata/SQLMetadataConnector.java +++ b/server/src/main/java/org/apache/druid/metadata/SQLMetadataConnector.java @@ -294,7 +294,7 @@ tableName, getPayloadType(), getQuoteString(), getCollation() ) ) ); - alterPendingSegmentsTableAddParentIdAndTaskGroup(tableName); + alterPendingSegmentsTable(tableName); } public void createDataSourceTable(final String tableName) @@ -481,7 +481,16 @@ private void alterEntryTableAddTypeAndGroupId(final String tableName) } } - private void alterPendingSegmentsTableAddParentIdAndTaskGroup(final String tableName) + /** + * Adds the following columns to the pending segments table to clean up unused records, + * and to faciliatate concurrent append and replace. + * 1) task_allocator_id -> The task id / task group id / task replica group id of the task that allocated it. + * 2) upgraded_from_segment_id -> The id of the segment from which the entry was upgraded upon concurrent replace. + * + * Also, adds an index on (dataSource, task_allocator_id) + * @param tableName name of the pending segments table + */ + private void alterPendingSegmentsTable(final String tableName) { List statements = new ArrayList<>(); if (tableHasColumn(tableName, "upgraded_from_segment_id")) { @@ -499,6 +508,14 @@ private void alterPendingSegmentsTableAddParentIdAndTaskGroup(final String table if (!statements.isEmpty()) { alterTable(tableName, statements); } + + final Set createdIndexSet = getIndexOnTable(tableName); + createIndex( + tableName, + StringUtils.format("idx_%1$s_datasource_task_allocator_id", tableName), + ImmutableList.of("dataSource", "task_allocator_id"), + createdIndexSet + ); } public void createLogTable(final String tableName, final String entryTypeName)