FrankChen021 commented on code in PR #19311: URL: https://github.com/apache/druid/pull/19311#discussion_r4036975698
########## docs/ingestion/kafka-share-group-ingestion.md: ########## @@ -0,0 +1,349 @@ +--- +id: kafka-share-group-ingestion +title: "Kafka share group ingestion" +sidebar_label: "Kafka share group ingestion" +description: "Queue-semantics ingestion from Apache Kafka using share groups (KIP-932). Scale consumers beyond partition count with at-least-once delivery." +--- + +<!-- + ~ Licensed to the Apache Software Foundation (ASF) under one + ~ or more contributor license agreements. See the NOTICE file + ~ distributed with this work for additional information + ~ regarding copyright ownership. The ASF licenses this file + ~ to you under the Apache License, Version 2.0 (the + ~ "License"); you may not use this file except in compliance + ~ with the License. You may obtain a copy of the License at + ~ + ~ http://www.apache.org/licenses/LICENSE-2.0 + ~ + ~ Unless required by applicable law or agreed to in writing, + ~ software distributed under the License is distributed on an + ~ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + ~ KIND, either express or implied. See the License for the + ~ specific language governing permissions and limitations + ~ under the License. + --> + +:::info +Requires Apache Kafka 4.0 or higher with share groups (KIP-932) enabled on the broker. +::: + +## Overview + +Kafka share groups (KIP-932) let multiple consumers read from the same partition concurrently. The broker manages per-record acquisition locks and explicit acknowledgement, so consumer count is not capped by partition count, joining or leaving consumers does not pause the group, and a slow record does not block its partition. + +Druid's `ShareGroupIndexTask` consumes from a share group and publishes segments with at-least-once delivery: records are acknowledged only after their segments are atomically registered in the metadata store. + +## When to use share group ingestion + +| Scenario | Consumer group | Share group | +|----------|---------------|-------------| +| Workers needed exceed partition count | Idle workers | All workers active | +| Elastic scaling (auto-scale events) | Rebalancing pause (30-60s) | Zero pause | +| Per-message processing time varies | Head-of-line blocking | Independent processing | +| Ordered processing required per partition | Yes | No (delivery order not guaranteed) | + +Choose share groups when throughput and elastic scaling matter more than strict per-partition ordering. + +## Task spec + +Submit a `ShareGroupIndexTask` to the Overlord. There are no start/end offsets -- the broker tracks them. + +```json +{ + "type": "index_kafka_share_group", + "dataSchema": { + "dataSource": "my_datasource", + "timestampSpec": { + "column": "__time", + "format": "auto" + }, + "dimensionsSpec": { + "useSchemaDiscovery": true + }, + "granularitySpec": { + "segmentGranularity": "DAY", + "queryGranularity": "NONE" + } + }, + "ioConfig": { + "topic": "my_topic", + "groupId": "druid-share-group", + "consumerProperties": { + "bootstrap.servers": "kafka-broker:9092" + }, + "inputFormat": { + "type": "json" + }, + "pollTimeout": 2000 + }, + "tuningConfig": { + "type": "KafkaTuningConfig", + "maxRowsPerSegment": 5000000 + } +} +``` + +## IO configuration + +| Property | Type | Required | Default | Description | +|----------|------|----------|---------|-------------| +| `topic` | String | Yes | -- | Kafka topic to consume from. | +| `groupId` | String | Yes | -- | Share group identifier. Multiple tasks with the same `groupId` share the workload. | +| `consumerProperties` | Map | Yes | -- | Kafka consumer properties. Must include `bootstrap.servers`. See [Consumer property restrictions](#consumer-property-restrictions). | +| `inputFormat` | Object | Yes | -- | Input format for parsing records (json, csv, avro, etc.). | +| `pollTimeout` | Long | No | 2000 | Poll timeout in milliseconds. | Review Comment: [P3] Document the actual pollTimeout default **Finding:** The configuration table says an omitted pollTimeout defaults to 2000 ms, but ShareGroupIndexTaskIOConfig uses KafkaSupervisorIOConfig.DEFAULT_POLL_TIMEOUT_MILLIS, which is 100 ms. Users relying on this documentation will poll twenty times more frequently than the implementation indicates in the docs. **Suggestion:** Make the documentation match the implementation, or change the new configuration default and its test to the intended 2000 ms value. ########## extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/ShareGroupIndexTaskRunner.java: ########## @@ -0,0 +1,535 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.kafka; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Supplier; +import com.google.common.collect.ImmutableMap; +import org.apache.druid.common.config.Configs; +import org.apache.druid.data.input.Committer; +import org.apache.druid.data.input.InputFormat; +import org.apache.druid.data.input.InputRow; +import org.apache.druid.data.input.InputRowSchema; +import org.apache.druid.data.input.kafka.KafkaRecordEntity; +import org.apache.druid.data.input.kafka.KafkaTopicPartition; +import org.apache.druid.indexer.TaskStatus; +import org.apache.druid.indexer.partitions.DynamicPartitionsSpec; +import org.apache.druid.indexing.appenderator.ActionBasedPublishedSegmentRetriever; +import org.apache.druid.indexing.appenderator.ActionBasedSegmentAllocator; +import org.apache.druid.indexing.common.LockGranularity; +import org.apache.druid.indexing.common.TaskLock; +import org.apache.druid.indexing.common.TaskLockType; +import org.apache.druid.indexing.common.TaskToolbox; +import org.apache.druid.indexing.common.actions.LockReleaseAction; +import org.apache.druid.indexing.common.actions.SegmentAllocateAction; +import org.apache.druid.indexing.common.actions.SegmentLockAcquireAction; +import org.apache.druid.indexing.common.actions.SegmentTransactionalAppendAction; +import org.apache.druid.indexing.common.actions.TaskLocks; +import org.apache.druid.indexing.common.actions.TimeChunkLockAcquireAction; +import org.apache.druid.indexing.common.task.InputRowFilter; +import org.apache.druid.indexing.common.task.Tasks; +import org.apache.druid.indexing.input.InputRowSchemas; +import org.apache.druid.indexing.overlord.SegmentPublishResult; +import org.apache.druid.indexing.seekablestream.SeekableStreamAppenderatorConfig; +import org.apache.druid.indexing.seekablestream.StreamChunkReader; +import org.apache.druid.indexing.seekablestream.common.AcknowledgeType; +import org.apache.druid.indexing.seekablestream.common.AcknowledgingRecordSupplier; +import org.apache.druid.indexing.seekablestream.common.OrderedPartitionableRecord; +import org.apache.druid.java.util.common.ISE; +import org.apache.druid.java.util.common.logger.Logger; +import org.apache.druid.segment.SegmentSchemaMapping; +import org.apache.druid.segment.incremental.ParseExceptionHandler; +import org.apache.druid.segment.incremental.RowIngestionMeters; +import org.apache.druid.segment.indexing.DataSchema; +import org.apache.druid.segment.realtime.SegmentGenerationMetrics; +import org.apache.druid.segment.realtime.appenderator.Appenderator; +import org.apache.druid.segment.realtime.appenderator.AppenderatorDriverAddResult; +import org.apache.druid.segment.realtime.appenderator.SegmentIdWithShardSpec; +import org.apache.druid.segment.realtime.appenderator.SegmentsAndCommitMetadata; +import org.apache.druid.segment.realtime.appenderator.StreamAppenderatorDriver; +import org.apache.druid.segment.realtime.appenderator.TransactionalSegmentPublisher; +import org.apache.druid.timeline.DataSegment; +import org.apache.druid.timeline.partition.NumberedPartialShardSpec; +import org.apache.kafka.common.errors.WakeupException; + +import javax.annotation.Nullable; +import java.io.IOException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; + +/** + * Single-threaded ingestion loop for {@link ShareGroupIndexTask}. Records are + * {@code ACCEPT}-acknowledged only after the containing segment is published; + * any earlier failure leaves them unacknowledged so the broker redelivers + * them once the acquisition lock expires. + */ +public class ShareGroupIndexTaskRunner +{ + private static final Logger log = new Logger(ShareGroupIndexTaskRunner.class); + + static final String METRIC_COMMIT_FAILURES = "ingest/shareGroup/commitFailures"; + + private final ShareGroupIndexTask task; + private final TaskToolbox toolbox; + private final ObjectMapper configMapper; + private final String sequenceName; + private final LockGranularity lockGranularity; + private final TaskLockType lockType; + private final Function<ShareGroupIndexTaskIOConfig, + AcknowledgingRecordSupplier<KafkaTopicPartition, Long, KafkaRecordEntity>> supplierFactory; + + private final AtomicReference<AcknowledgingRecordSupplier<KafkaTopicPartition, Long, KafkaRecordEntity>> + activeSupplier = new AtomicReference<>(); + + ShareGroupIndexTaskRunner( + ShareGroupIndexTask task, + TaskToolbox toolbox, + ObjectMapper configMapper + ) + { + this(task, toolbox, configMapper, null); + } + + @VisibleForTesting + ShareGroupIndexTaskRunner( + ShareGroupIndexTask task, + TaskToolbox toolbox, + ObjectMapper configMapper, + @Nullable Function<ShareGroupIndexTaskIOConfig, + AcknowledgingRecordSupplier<KafkaTopicPartition, Long, KafkaRecordEntity>> supplierFactory + ) + { + this.task = task; + this.toolbox = toolbox; + this.configMapper = configMapper; + this.sequenceName = task.getId(); + final boolean useTimeChunkLocks = Configs.valueOrDefault( + task.getContextValue(Tasks.FORCE_TIME_CHUNK_LOCK_KEY, Tasks.DEFAULT_FORCE_TIME_CHUNK_LOCK), + Tasks.DEFAULT_FORCE_TIME_CHUNK_LOCK + ); + this.lockGranularity = useTimeChunkLocks ? LockGranularity.TIME_CHUNK : LockGranularity.SEGMENT; + this.lockType = TaskLocks.determineLockTypeForAppend(task.getContext()); + this.supplierFactory = supplierFactory != null ? supplierFactory : this::createDefaultRecordSupplier; + } + + public TaskStatus run() throws Exception + { + final DataSchema dataSchema = task.getDataSchema(); + final KafkaIndexTaskTuningConfig tuningConfig = task.getTuningConfig(); + final ShareGroupIndexTaskIOConfig ioConfig = task.getIOConfig(); + + final RowIngestionMeters rowIngestionMeters = + toolbox.getRowIngestionMetersFactory().createRowIngestionMeters(); + final ParseExceptionHandler parseExceptionHandler = new ParseExceptionHandler( + rowIngestionMeters, + tuningConfig.isLogParseExceptions(), + tuningConfig.getMaxParseExceptions(), + tuningConfig.getMaxSavedParseExceptions() + ); + final SegmentGenerationMetrics segmentGenerationMetrics = new SegmentGenerationMetrics(); + + final InputRowSchema inputRowSchema = InputRowSchemas.fromDataSchema(dataSchema); + final InputFormat inputFormat = ioConfig.getInputFormat(); + if (inputFormat == null) { + throw new ISE("inputFormat must be specified in ioConfig"); + } + + // Raw type mirrors SeekableStreamIndexTaskRunner; works around + // OrderedPartitionableRecord.getData() returning a wildcard list. + @SuppressWarnings({"rawtypes", "unchecked"}) + final StreamChunkReader chunkReader = new StreamChunkReader<KafkaRecordEntity>( + inputFormat, + inputRowSchema, + dataSchema.getTransformSpec(), + toolbox.getIndexingTmpDir(), + InputRowFilter.allowAll(), + rowIngestionMeters, + parseExceptionHandler + ); + + final Appenderator appenderator = toolbox.getAppenderatorsManager().createRealtimeAppenderatorForTask( + toolbox.getSegmentLoaderConfig(), + task.getId(), + dataSchema, + SeekableStreamAppenderatorConfig.fromTuningConfig( + tuningConfig.withBasePersistDirectory(toolbox.getPersistDir()), + toolbox.getProcessingConfig() + ), + toolbox.getConfig(), + segmentGenerationMetrics, + toolbox.getSegmentPusher(), + toolbox.getJsonMapper(), + toolbox.getIndexIO(), + toolbox.getIndexMerger(), + toolbox.getQueryRunnerFactoryConglomerate(), + toolbox.getSegmentAnnouncer(), + toolbox.getEmitter(), + toolbox.getQueryProcessingPool(), + toolbox.getJoinableFactory(), + toolbox.getCache(), + toolbox.getCacheConfig(), + toolbox.getCachePopulatorStats(), + toolbox.getPolicyEnforcer(), + rowIngestionMeters, + parseExceptionHandler, + toolbox.getCentralizedTableSchemaConfig(), + interval -> { + toolbox.getTaskActionClient().submit(new LockReleaseAction(interval)); + } + ); + + final StreamAppenderatorDriver driver = new StreamAppenderatorDriver( + appenderator, + new ActionBasedSegmentAllocator( + toolbox.getTaskActionClient(), + dataSchema, + (schema, row, sequenceName, previousSegmentId, skipSegmentLineageCheck) -> new SegmentAllocateAction( + schema.getDataSource(), + row.getTimestamp(), + schema.getGranularitySpec().getQueryGranularity(), + schema.getGranularitySpec().getSegmentGranularity(), + sequenceName, + previousSegmentId, + skipSegmentLineageCheck, + NumberedPartialShardSpec.instance(), + lockGranularity, + lockType + ) + ), + toolbox.getSegmentHandoffNotifierFactory(), + new ActionBasedPublishedSegmentRetriever(toolbox.getTaskActionClient()), + toolbox.getDataSegmentKiller(), + toolbox.getJsonMapper(), + segmentGenerationMetrics + ); + + final org.apache.druid.indexing.common.stats.TaskRealtimeMetricsMonitor metricsMonitor = + new org.apache.druid.indexing.common.stats.TaskRealtimeMetricsMonitor( + segmentGenerationMetrics, + rowIngestionMeters, + task.getMetricBuilder() + ); + toolbox.addMonitor(metricsMonitor); + + boolean runLoopCompleted = false; + try (final AcknowledgingRecordSupplier<KafkaTopicPartition, Long, KafkaRecordEntity> recordSupplier = + supplierFactory.apply(ioConfig)) { + activeSupplier.set(recordSupplier); + final TaskStatus status = runLoop(driver, appenderator, recordSupplier, chunkReader, ioConfig, tuningConfig, toolbox); + runLoopCompleted = true; + return status; + } + finally { + activeSupplier.set(null); + if (!runLoopCompleted) { + try { + appenderator.closeNow(); + } + catch (Exception e) { + log.warn(e, "Exception during emergency closeNow() of Appenderator in run(); continuing teardown."); + } + } + try { + toolbox.removeMonitor(metricsMonitor); + } + catch (Exception e) { + log.warn(e, "Exception removing TaskRealtimeMetricsMonitor; continuing teardown."); + } + try { + driver.close(); + } + catch (Exception e) { + log.warn(e, "Exception closing StreamAppenderatorDriver; continuing teardown."); + } + } + } + + @VisibleForTesting + TaskStatus runLoop( + StreamAppenderatorDriver driver, + Appenderator appenderator, + AcknowledgingRecordSupplier<KafkaTopicPartition, Long, KafkaRecordEntity> recordSupplier, + StreamChunkReader chunkReader, + ShareGroupIndexTaskIOConfig ioConfig, + KafkaIndexTaskTuningConfig tuningConfig, + TaskToolbox toolbox + ) throws Exception + { + // Share groups manage delivery state on the broker; the Committer is a + // no-op placeholder required by the appenderator driver contract. + final Supplier<Committer> committerSupplier = () -> new Committer() + { + @Override + public Object getMetadata() + { + return ImmutableMap.of("type", "share_group"); + } + + @Override + public void run() + { + } + }; + + log.info("Starting share group ingestion loop for topic[%s], group[%s].", + ioConfig.getTopic(), ioConfig.getGroupId()); + + long totalRowsIngested = 0; + boolean lockTimeoutLogged = false; + boolean appenderatorClosedNormally = false; + + try { + recordSupplier.subscribe(Collections.singleton(ioConfig.getTopic())); + driver.startJob(this::acquireLockForRestoredSegment); + while (!task.isStopRequested()) { + + final List<OrderedPartitionableRecord<KafkaTopicPartition, Long, KafkaRecordEntity>> records; + try { + records = recordSupplier.poll(ioConfig.getPollTimeout()); + } + catch (WakeupException e) { + if (task.isStopRequested()) { + log.info("Wake-up received while stopGracefully is in progress; exiting ingestion loop."); + break; + } + throw e; + } + + if (!lockTimeoutLogged) { + final Optional<Integer> effectiveLockMs = recordSupplier.acquisitionLockTimeoutMs(); + log.info( + "Effective broker acquisition lock timeout for share-group[%s]: %s ms", + ioConfig.getGroupId(), + effectiveLockMs.map(String::valueOf).orElse("<unknown>") + ); + lockTimeoutLogged = true; + } + + if (records.isEmpty()) { + continue; + } + + final Map<KafkaTopicPartition, List<Long>> batchOffsets = new HashMap<>(); + boolean midBatchPersistNeeded = false; + boolean pushThresholdLogged = false; + + for (OrderedPartitionableRecord<KafkaTopicPartition, Long, KafkaRecordEntity> record : records) { + batchOffsets.computeIfAbsent(record.getPartitionId(), k -> new ArrayList<>()) + .add(record.getSequenceNumber()); + + final List<InputRow> rows = chunkReader.parse(record.getData(), false); + + for (InputRow row : rows) { + final AppenderatorDriverAddResult addResult = driver.add( + row, + sequenceName, + committerSupplier, + true, + false + ); + + if (!addResult.isOk()) { + // Throw without acknowledging so the broker redelivers this batch. + throw new ISE( + "Could not allocate segment for row with timestamp[%s]", + row.getTimestamp() + ); + } + + totalRowsIngested++; + midBatchPersistNeeded |= addResult.isPersistRequired(); + + if (!pushThresholdLogged && addResult.isPushRequired( + tuningConfig.getPartitionsSpec().getMaxRowsPerSegment(), + tuningConfig.getPartitionsSpec() + .getMaxTotalRowsOr(DynamicPartitionsSpec.DEFAULT_MAX_TOTAL_ROWS) + )) { + pushThresholdLogged = true; + log.info( + "Segment row threshold reached during batch (current segment[%s], rows[%d]); " + + "will publish at end-of-batch boundary.", + addResult.getSegmentIdentifier(), + addResult.getNumRowsInSegment() + ); + } + } + } + + if (midBatchPersistNeeded) { + driver.persist(committerSupplier.get()); + } + + driver.persist(committerSupplier.get()); + + final TransactionalSegmentPublisher publisher = new TransactionalSegmentPublisher() + { + @Override + public SegmentPublishResult publishAnnotatedSegments( + @Nullable Set<DataSegment> mustBeNullOrEmptyOverwriteSegments, + Set<DataSegment> segmentsToPush, + @Nullable Object commitMetadata, + SegmentSchemaMapping segmentSchemaMapping + ) throws IOException + { + if (segmentsToPush.isEmpty()) { + return SegmentPublishResult.ok(segmentsToPush); + } + return toolbox.getTaskActionClient().submit( + SegmentTransactionalAppendAction.forSegments(segmentsToPush, segmentSchemaMapping) Review Comment: [P2] Select the publisher for the configured lock type **Finding:** ShareGroupIndexTask only defaults useConcurrentLocks with putIfAbsent, so a task may explicitly set it false or choose useSharedLock/taskLockType and acquire EXCLUSIVE or SHARED locks. This publisher always submits SegmentTransactionalAppendAction, whose perform method rejects every task lock that is not APPEND, so the first non-empty batch fails during publish and its records are never acknowledged for that otherwise accepted task context. **Suggestion:** Either reject or normalize non-APPEND lock contexts for this task, or select the publish action from the actual lock type as the existing batch-task path does; add coverage for useConcurrentLocks=false. ########## extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/ShareGroupIndexTaskRunner.java: ########## @@ -0,0 +1,535 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.kafka; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Supplier; +import com.google.common.collect.ImmutableMap; +import org.apache.druid.common.config.Configs; +import org.apache.druid.data.input.Committer; +import org.apache.druid.data.input.InputFormat; +import org.apache.druid.data.input.InputRow; +import org.apache.druid.data.input.InputRowSchema; +import org.apache.druid.data.input.kafka.KafkaRecordEntity; +import org.apache.druid.data.input.kafka.KafkaTopicPartition; +import org.apache.druid.indexer.TaskStatus; +import org.apache.druid.indexer.partitions.DynamicPartitionsSpec; +import org.apache.druid.indexing.appenderator.ActionBasedPublishedSegmentRetriever; +import org.apache.druid.indexing.appenderator.ActionBasedSegmentAllocator; +import org.apache.druid.indexing.common.LockGranularity; +import org.apache.druid.indexing.common.TaskLock; +import org.apache.druid.indexing.common.TaskLockType; +import org.apache.druid.indexing.common.TaskToolbox; +import org.apache.druid.indexing.common.actions.LockReleaseAction; +import org.apache.druid.indexing.common.actions.SegmentAllocateAction; +import org.apache.druid.indexing.common.actions.SegmentLockAcquireAction; +import org.apache.druid.indexing.common.actions.SegmentTransactionalAppendAction; +import org.apache.druid.indexing.common.actions.TaskLocks; +import org.apache.druid.indexing.common.actions.TimeChunkLockAcquireAction; +import org.apache.druid.indexing.common.task.InputRowFilter; +import org.apache.druid.indexing.common.task.Tasks; +import org.apache.druid.indexing.input.InputRowSchemas; +import org.apache.druid.indexing.overlord.SegmentPublishResult; +import org.apache.druid.indexing.seekablestream.SeekableStreamAppenderatorConfig; +import org.apache.druid.indexing.seekablestream.StreamChunkReader; +import org.apache.druid.indexing.seekablestream.common.AcknowledgeType; +import org.apache.druid.indexing.seekablestream.common.AcknowledgingRecordSupplier; +import org.apache.druid.indexing.seekablestream.common.OrderedPartitionableRecord; +import org.apache.druid.java.util.common.ISE; +import org.apache.druid.java.util.common.logger.Logger; +import org.apache.druid.segment.SegmentSchemaMapping; +import org.apache.druid.segment.incremental.ParseExceptionHandler; +import org.apache.druid.segment.incremental.RowIngestionMeters; +import org.apache.druid.segment.indexing.DataSchema; +import org.apache.druid.segment.realtime.SegmentGenerationMetrics; +import org.apache.druid.segment.realtime.appenderator.Appenderator; +import org.apache.druid.segment.realtime.appenderator.AppenderatorDriverAddResult; +import org.apache.druid.segment.realtime.appenderator.SegmentIdWithShardSpec; +import org.apache.druid.segment.realtime.appenderator.SegmentsAndCommitMetadata; +import org.apache.druid.segment.realtime.appenderator.StreamAppenderatorDriver; +import org.apache.druid.segment.realtime.appenderator.TransactionalSegmentPublisher; +import org.apache.druid.timeline.DataSegment; +import org.apache.druid.timeline.partition.NumberedPartialShardSpec; +import org.apache.kafka.common.errors.WakeupException; + +import javax.annotation.Nullable; +import java.io.IOException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; + +/** + * Single-threaded ingestion loop for {@link ShareGroupIndexTask}. Records are + * {@code ACCEPT}-acknowledged only after the containing segment is published; + * any earlier failure leaves them unacknowledged so the broker redelivers + * them once the acquisition lock expires. + */ +public class ShareGroupIndexTaskRunner +{ + private static final Logger log = new Logger(ShareGroupIndexTaskRunner.class); + + static final String METRIC_COMMIT_FAILURES = "ingest/shareGroup/commitFailures"; + + private final ShareGroupIndexTask task; + private final TaskToolbox toolbox; + private final ObjectMapper configMapper; + private final String sequenceName; + private final LockGranularity lockGranularity; + private final TaskLockType lockType; + private final Function<ShareGroupIndexTaskIOConfig, + AcknowledgingRecordSupplier<KafkaTopicPartition, Long, KafkaRecordEntity>> supplierFactory; + + private final AtomicReference<AcknowledgingRecordSupplier<KafkaTopicPartition, Long, KafkaRecordEntity>> + activeSupplier = new AtomicReference<>(); + + ShareGroupIndexTaskRunner( + ShareGroupIndexTask task, + TaskToolbox toolbox, + ObjectMapper configMapper + ) + { + this(task, toolbox, configMapper, null); + } + + @VisibleForTesting + ShareGroupIndexTaskRunner( + ShareGroupIndexTask task, + TaskToolbox toolbox, + ObjectMapper configMapper, + @Nullable Function<ShareGroupIndexTaskIOConfig, + AcknowledgingRecordSupplier<KafkaTopicPartition, Long, KafkaRecordEntity>> supplierFactory + ) + { + this.task = task; + this.toolbox = toolbox; + this.configMapper = configMapper; + this.sequenceName = task.getId(); + final boolean useTimeChunkLocks = Configs.valueOrDefault( + task.getContextValue(Tasks.FORCE_TIME_CHUNK_LOCK_KEY, Tasks.DEFAULT_FORCE_TIME_CHUNK_LOCK), + Tasks.DEFAULT_FORCE_TIME_CHUNK_LOCK + ); + this.lockGranularity = useTimeChunkLocks ? LockGranularity.TIME_CHUNK : LockGranularity.SEGMENT; + this.lockType = TaskLocks.determineLockTypeForAppend(task.getContext()); + this.supplierFactory = supplierFactory != null ? supplierFactory : this::createDefaultRecordSupplier; + } + + public TaskStatus run() throws Exception + { + final DataSchema dataSchema = task.getDataSchema(); + final KafkaIndexTaskTuningConfig tuningConfig = task.getTuningConfig(); + final ShareGroupIndexTaskIOConfig ioConfig = task.getIOConfig(); + + final RowIngestionMeters rowIngestionMeters = + toolbox.getRowIngestionMetersFactory().createRowIngestionMeters(); + final ParseExceptionHandler parseExceptionHandler = new ParseExceptionHandler( + rowIngestionMeters, + tuningConfig.isLogParseExceptions(), + tuningConfig.getMaxParseExceptions(), + tuningConfig.getMaxSavedParseExceptions() + ); + final SegmentGenerationMetrics segmentGenerationMetrics = new SegmentGenerationMetrics(); + + final InputRowSchema inputRowSchema = InputRowSchemas.fromDataSchema(dataSchema); + final InputFormat inputFormat = ioConfig.getInputFormat(); + if (inputFormat == null) { + throw new ISE("inputFormat must be specified in ioConfig"); + } + + // Raw type mirrors SeekableStreamIndexTaskRunner; works around + // OrderedPartitionableRecord.getData() returning a wildcard list. + @SuppressWarnings({"rawtypes", "unchecked"}) + final StreamChunkReader chunkReader = new StreamChunkReader<KafkaRecordEntity>( + inputFormat, + inputRowSchema, + dataSchema.getTransformSpec(), + toolbox.getIndexingTmpDir(), + InputRowFilter.allowAll(), + rowIngestionMeters, + parseExceptionHandler + ); + + final Appenderator appenderator = toolbox.getAppenderatorsManager().createRealtimeAppenderatorForTask( + toolbox.getSegmentLoaderConfig(), + task.getId(), + dataSchema, + SeekableStreamAppenderatorConfig.fromTuningConfig( + tuningConfig.withBasePersistDirectory(toolbox.getPersistDir()), + toolbox.getProcessingConfig() + ), + toolbox.getConfig(), + segmentGenerationMetrics, + toolbox.getSegmentPusher(), + toolbox.getJsonMapper(), + toolbox.getIndexIO(), + toolbox.getIndexMerger(), + toolbox.getQueryRunnerFactoryConglomerate(), + toolbox.getSegmentAnnouncer(), + toolbox.getEmitter(), + toolbox.getQueryProcessingPool(), + toolbox.getJoinableFactory(), + toolbox.getCache(), + toolbox.getCacheConfig(), + toolbox.getCachePopulatorStats(), + toolbox.getPolicyEnforcer(), + rowIngestionMeters, + parseExceptionHandler, + toolbox.getCentralizedTableSchemaConfig(), + interval -> { + toolbox.getTaskActionClient().submit(new LockReleaseAction(interval)); + } + ); + + final StreamAppenderatorDriver driver = new StreamAppenderatorDriver( + appenderator, + new ActionBasedSegmentAllocator( + toolbox.getTaskActionClient(), + dataSchema, + (schema, row, sequenceName, previousSegmentId, skipSegmentLineageCheck) -> new SegmentAllocateAction( + schema.getDataSource(), + row.getTimestamp(), + schema.getGranularitySpec().getQueryGranularity(), + schema.getGranularitySpec().getSegmentGranularity(), + sequenceName, + previousSegmentId, + skipSegmentLineageCheck, + NumberedPartialShardSpec.instance(), + lockGranularity, + lockType + ) + ), + toolbox.getSegmentHandoffNotifierFactory(), + new ActionBasedPublishedSegmentRetriever(toolbox.getTaskActionClient()), + toolbox.getDataSegmentKiller(), + toolbox.getJsonMapper(), + segmentGenerationMetrics + ); + + final org.apache.druid.indexing.common.stats.TaskRealtimeMetricsMonitor metricsMonitor = + new org.apache.druid.indexing.common.stats.TaskRealtimeMetricsMonitor( + segmentGenerationMetrics, + rowIngestionMeters, + task.getMetricBuilder() + ); + toolbox.addMonitor(metricsMonitor); + + boolean runLoopCompleted = false; + try (final AcknowledgingRecordSupplier<KafkaTopicPartition, Long, KafkaRecordEntity> recordSupplier = + supplierFactory.apply(ioConfig)) { + activeSupplier.set(recordSupplier); + final TaskStatus status = runLoop(driver, appenderator, recordSupplier, chunkReader, ioConfig, tuningConfig, toolbox); + runLoopCompleted = true; + return status; + } + finally { + activeSupplier.set(null); + if (!runLoopCompleted) { + try { + appenderator.closeNow(); + } + catch (Exception e) { + log.warn(e, "Exception during emergency closeNow() of Appenderator in run(); continuing teardown."); + } + } + try { + toolbox.removeMonitor(metricsMonitor); + } + catch (Exception e) { + log.warn(e, "Exception removing TaskRealtimeMetricsMonitor; continuing teardown."); + } + try { + driver.close(); + } + catch (Exception e) { + log.warn(e, "Exception closing StreamAppenderatorDriver; continuing teardown."); + } + } + } + + @VisibleForTesting + TaskStatus runLoop( + StreamAppenderatorDriver driver, + Appenderator appenderator, + AcknowledgingRecordSupplier<KafkaTopicPartition, Long, KafkaRecordEntity> recordSupplier, + StreamChunkReader chunkReader, + ShareGroupIndexTaskIOConfig ioConfig, + KafkaIndexTaskTuningConfig tuningConfig, + TaskToolbox toolbox + ) throws Exception + { + // Share groups manage delivery state on the broker; the Committer is a + // no-op placeholder required by the appenderator driver contract. + final Supplier<Committer> committerSupplier = () -> new Committer() + { + @Override + public Object getMetadata() + { + return ImmutableMap.of("type", "share_group"); + } + + @Override + public void run() + { + } + }; + + log.info("Starting share group ingestion loop for topic[%s], group[%s].", + ioConfig.getTopic(), ioConfig.getGroupId()); + + long totalRowsIngested = 0; + boolean lockTimeoutLogged = false; + boolean appenderatorClosedNormally = false; + + try { + recordSupplier.subscribe(Collections.singleton(ioConfig.getTopic())); + driver.startJob(this::acquireLockForRestoredSegment); + while (!task.isStopRequested()) { + + final List<OrderedPartitionableRecord<KafkaTopicPartition, Long, KafkaRecordEntity>> records; + try { + records = recordSupplier.poll(ioConfig.getPollTimeout()); + } + catch (WakeupException e) { + if (task.isStopRequested()) { + log.info("Wake-up received while stopGracefully is in progress; exiting ingestion loop."); + break; + } + throw e; + } + + if (!lockTimeoutLogged) { + final Optional<Integer> effectiveLockMs = recordSupplier.acquisitionLockTimeoutMs(); + log.info( + "Effective broker acquisition lock timeout for share-group[%s]: %s ms", + ioConfig.getGroupId(), + effectiveLockMs.map(String::valueOf).orElse("<unknown>") + ); + lockTimeoutLogged = true; + } + + if (records.isEmpty()) { + continue; + } + + final Map<KafkaTopicPartition, List<Long>> batchOffsets = new HashMap<>(); + boolean midBatchPersistNeeded = false; + boolean pushThresholdLogged = false; + + for (OrderedPartitionableRecord<KafkaTopicPartition, Long, KafkaRecordEntity> record : records) { + batchOffsets.computeIfAbsent(record.getPartitionId(), k -> new ArrayList<>()) + .add(record.getSequenceNumber()); + + final List<InputRow> rows = chunkReader.parse(record.getData(), false); + + for (InputRow row : rows) { + final AppenderatorDriverAddResult addResult = driver.add( + row, + sequenceName, + committerSupplier, + true, + false + ); + + if (!addResult.isOk()) { + // Throw without acknowledging so the broker redelivers this batch. + throw new ISE( + "Could not allocate segment for row with timestamp[%s]", + row.getTimestamp() + ); + } + + totalRowsIngested++; + midBatchPersistNeeded |= addResult.isPersistRequired(); + + if (!pushThresholdLogged && addResult.isPushRequired( + tuningConfig.getPartitionsSpec().getMaxRowsPerSegment(), + tuningConfig.getPartitionsSpec() + .getMaxTotalRowsOr(DynamicPartitionsSpec.DEFAULT_MAX_TOTAL_ROWS) + )) { + pushThresholdLogged = true; + log.info( + "Segment row threshold reached during batch (current segment[%s], rows[%d]); " + + "will publish at end-of-batch boundary.", + addResult.getSegmentIdentifier(), + addResult.getNumRowsInSegment() + ); + } + } + } + + if (midBatchPersistNeeded) { + driver.persist(committerSupplier.get()); + } + + driver.persist(committerSupplier.get()); + + final TransactionalSegmentPublisher publisher = new TransactionalSegmentPublisher() + { + @Override + public SegmentPublishResult publishAnnotatedSegments( + @Nullable Set<DataSegment> mustBeNullOrEmptyOverwriteSegments, + Set<DataSegment> segmentsToPush, + @Nullable Object commitMetadata, + SegmentSchemaMapping segmentSchemaMapping + ) throws IOException + { + if (segmentsToPush.isEmpty()) { + return SegmentPublishResult.ok(segmentsToPush); + } + return toolbox.getTaskActionClient().submit( + SegmentTransactionalAppendAction.forSegments(segmentsToPush, segmentSchemaMapping) + ); + } + + @Override + public boolean supportsEmptyPublish() + { + return false; + } + }; + + final SegmentsAndCommitMetadata published = driver.publish( + publisher, + committerSupplier.get(), + Collections.singletonList(sequenceName) + ).get(); + + final int segmentCount = published.getSegments().size(); + if (segmentCount > 0) { + log.info("Published %d segments with %d total rows.", segmentCount, totalRowsIngested); + } + + // Acknowledge only after the segment is published; any earlier failure + // leaves records unacknowledged so the broker redelivers them. + for (Map.Entry<KafkaTopicPartition, List<Long>> entry : batchOffsets.entrySet()) { + for (Long offset : entry.getValue()) { + recordSupplier.acknowledge(entry.getKey(), offset, AcknowledgeType.ACCEPT); + } + } + + final Map<KafkaTopicPartition, Optional<Exception>> commitResult = recordSupplier.commitSync(); + long commitFailures = 0L; + for (Map.Entry<KafkaTopicPartition, Optional<Exception>> entry : commitResult.entrySet()) { + if (entry.getValue().isPresent()) { + commitFailures++; + log.warn( + entry.getValue().get(), + "Commit failed for partition[%s]. Records may be redelivered.", + entry.getKey() + ); + } + } + if (commitFailures > 0) { + task.emitMetric(toolbox.getEmitter(), METRIC_COMMIT_FAILURES, commitFailures); + } + } + + driver.persist(committerSupplier.get()); + log.info("Share group ingestion complete. Total rows ingested: %d", totalRowsIngested); + appenderator.close(); + appenderatorClosedNormally = true; + } + finally { + if (!appenderatorClosedNormally) { + try { + appenderator.closeNow(); + } + catch (Exception e) { + log.warn(e, "Exception during emergency closeNow() of Appenderator; continuing teardown."); + } + } + } + + return TaskStatus.success(task.getId()); + } + + private boolean acquireLockForRestoredSegment(SegmentIdWithShardSpec segmentId) + { + try { + if (lockGranularity == LockGranularity.SEGMENT) { + return toolbox.getTaskActionClient().submit( + new SegmentLockAcquireAction( + lockType, Review Comment: [P2] Avoid APPEND locks with segment granularity **Finding:** When forceTimeChunkLock is false, this task still defaults useConcurrentLocks to true, making lockType APPEND. The segment restore branch passes that type to SegmentLockAcquireAction even though Druid permits APPEND only with time-chunk locking and TaskLockbox rejects an APPEND segment request when conflicting locks exist; SegmentAllocateAction receives the same unsupported pairing. A segment-granularity task can therefore fail to restore a persisted segment, whereas the existing SeekableStreamIndexTaskRunner uses EXCLUSIVE for this branch. **Suggestion:** Use a supported lock-type/granularity pairing for segment mode and keep allocation, restoration, and publishing consistent, or explicitly reject forceTimeChunkLock=false; cover the configuration with a lockbox-backed test. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
