FrankChen021 commented on code in PR #19311: URL: https://github.com/apache/druid/pull/19311#discussion_r4133565073
########## extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/DurableShareGroupBatchStager.java: ########## @@ -0,0 +1,129 @@ +/* + * 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.google.common.base.Preconditions; +import org.apache.druid.data.input.kafka.KafkaRecordEntity; +import org.apache.druid.data.input.kafka.KafkaTopicPartition; +import org.apache.druid.indexing.common.actions.StageShareInboxBatchAction; +import org.apache.druid.indexing.common.actions.TaskActionClient; +import org.apache.druid.indexing.overlord.ShareInboxBatch; +import org.apache.druid.indexing.overlord.ShareInboxStageResult; +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 java.io.IOException; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.UUID; + +final class DurableShareGroupBatchStager implements ShareGroupBatchStager +{ + private static final Logger log = new Logger(DurableShareGroupBatchStager.class); + + private final ShareGroupBatchStore batchStore; + private final TaskActionClient taskActionClient; + private final String dataSource; + private final String inboxId; + private final String groupId; + private final ShareGroupSourceIdentity sourceIdentity; + private final String specFingerprint; + private final int receiptPageSize; + + DurableShareGroupBatchStager( + ShareGroupBatchStore batchStore, + TaskActionClient taskActionClient, + String dataSource, + String inboxId, + String groupId, + ShareGroupSourceIdentity sourceIdentity, + String specFingerprint, + int receiptPageSize + ) + { + this.batchStore = Preconditions.checkNotNull(batchStore, "batchStore"); + this.taskActionClient = Preconditions.checkNotNull(taskActionClient, "taskActionClient"); + this.dataSource = Preconditions.checkNotNull(dataSource, "dataSource"); + this.inboxId = Preconditions.checkNotNull(inboxId, "inboxId"); + this.groupId = Preconditions.checkNotNull(groupId, "groupId"); + this.sourceIdentity = Preconditions.checkNotNull(sourceIdentity, "sourceIdentity"); + this.specFingerprint = Preconditions.checkNotNull(specFingerprint, "specFingerprint"); + Preconditions.checkArgument(receiptPageSize > 0, "receiptPageSize must be positive"); + this.receiptPageSize = receiptPageSize; + } + + @Override + public void stage( + List<OrderedPartitionableRecord<KafkaTopicPartition, Long, KafkaRecordEntity>> records + ) throws IOException + { + final ShareGroupBatchStore.StoredBatch storedBatch = batchStore.store(records); + final ShareGroupRawBatch rawBatch = storedBatch.getBatch(); + final List<Long> offsets = rawBatch.getRecords().stream() + .map(ShareGroupRawRecord::getOffset) + .toList(); + final ShareInboxBatch inboxBatch = new ShareInboxBatch( + dataSource, + inboxId, + groupId, + sourceIdentity.getClusterId(), + sourceIdentity.getTopicId(), + sourceIdentity.getTopicName(), + rawBatch.getPartition(), + specFingerprint, + UUID.randomUUID().toString(), + storedBatch.getPath(), + storedBatch.getSha256(), + storedBatch.getSize(), + offsets, + receiptPageSize + ); + final ShareInboxStageResult result = taskActionClient.submit(new StageShareInboxBatchAction(inboxBatch)); + validateResult(result, offsets); + if (result.getStatus() == ShareInboxStageResult.Status.ALREADY_STAGED) { + try { + batchStore.delete(storedBatch); Review Comment: [P1] Preserve the referenced object when a stage request is retried **Finding:** If the Overlord commits StageShareInboxBatchAction but its response is lost, RemoteTaskActionClient can automatically retry the same request. stageShareInboxBatch then finds the receipts and returns ALREADY_STAGED even though the existing manifest references this exact storedBatch path (the coordinator test explicitly expects this result for a repeated batch). This branch deletes that referenced object and staging reports success, allowing ACCEPT to be sent to Kafka. If inbox processing has not read the object yet, the acknowledged records can no longer be recovered. The lost-response unit test only models an exception escaping submit, so it misses the successful internal retry. **Suggestion:** Make staging retries idempotent by manifest ID and distinguish an already referenced object from a redundant new upload before deleting anything; cover commit followed by a lost response and successful retry. ########## extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/AppenderatorShareInboxBatchHandler.java: ########## @@ -0,0 +1,178 @@ +/* + * 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.google.common.base.Supplier; +import org.apache.druid.data.input.Committer; +import org.apache.druid.data.input.InputRow; +import org.apache.druid.data.input.kafka.KafkaRecordEntity; +import org.apache.druid.data.input.kafka.KafkaTopicPartition; +import org.apache.druid.indexing.common.actions.SegmentTransactionalAppendFromShareInboxAction; +import org.apache.druid.indexing.common.actions.TaskActionClient; +import org.apache.druid.indexing.overlord.SegmentPublishResult; +import org.apache.druid.indexing.overlord.ShareInboxCompletionRequest; +import org.apache.druid.indexing.overlord.ShareInboxManifest; +import org.apache.druid.indexing.seekablestream.StreamChunkReader; +import org.apache.druid.indexing.seekablestream.common.OrderedPartitionableRecord; +import org.apache.druid.java.util.common.ISE; +import org.apache.druid.segment.SegmentSchemaMapping; +import org.apache.druid.segment.realtime.appenderator.AppenderatorDriverAddResult; +import org.apache.druid.segment.realtime.appenderator.StreamAppenderatorDriver; +import org.apache.druid.segment.realtime.appenderator.TransactionalSegmentPublisher; +import org.apache.druid.timeline.DataSegment; + +import javax.annotation.Nullable; +import java.io.IOException; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.TreeMap; +import java.util.UUID; + +final class AppenderatorShareInboxBatchHandler implements ShareInboxBatchHandler +{ + private final StreamAppenderatorDriver driver; + private final StreamChunkReader<KafkaRecordEntity> chunkReader; + private final TaskActionClient taskActionClient; + private final String taskId; + + AppenderatorShareInboxBatchHandler( + StreamAppenderatorDriver driver, + StreamChunkReader<KafkaRecordEntity> chunkReader, + TaskActionClient taskActionClient, + String taskId + ) + { + this.driver = driver; + this.chunkReader = chunkReader; + this.taskActionClient = taskActionClient; + this.taskId = taskId; + } + + @Override + public void process( + List<ShareInboxManifest> manifests, + List<OrderedPartitionableRecord<KafkaTopicPartition, Long, KafkaRecordEntity>> records, + ShareInboxClaimLease claimLease + ) throws Exception + { + final String completionId = UUID.randomUUID().toString(); + final String sequenceName = taskId + "-share-inbox-" + completionId; + final Supplier<Committer> committerSupplier = () -> new Committer() + { + @Override + public Object getMetadata() + { + return Map.of("shareInboxCompletionId", completionId); + } + + @Override + public void run() + { + } + }; + + boolean rowsAdded = false; + for (OrderedPartitionableRecord<KafkaTopicPartition, Long, KafkaRecordEntity> record : records) { + claimLease.assertValid(); + final List<InputRow> rows = chunkReader.parse(kafkaEntities(record), false); + for (InputRow row : rows) { + claimLease.assertValid(); + final AppenderatorDriverAddResult addResult = driver.add( + row, + sequenceName, + committerSupplier, + true, + false Review Comment: [P2] Honor the appenderator persistence threshold inside the row loop **Finding:** driver.add is called with allowIncrementalPersists=false, but the returned isPersistRequired signal is never handled. The only persist happens after every record in the claim has been parsed and added. A valid claim can contain up to 100,000 Kafka records by default, and each record can produce multiple rows, so maxRowsInMemory/maxBytesInMemory can be exceeded for the entire batch. Large claims or a deliberately small memory setting can exhaust the task heap and fail repeatedly on the same durable inbox batch. **Suggestion:** Enable incremental persists or persist as soon as addResult.isPersistRequired() becomes true, and verify a multi-row batch that crosses the configured memory threshold. ########## extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/AppenderatorShareInboxBatchHandler.java: ########## @@ -0,0 +1,178 @@ +/* + * 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.google.common.base.Supplier; +import org.apache.druid.data.input.Committer; +import org.apache.druid.data.input.InputRow; +import org.apache.druid.data.input.kafka.KafkaRecordEntity; +import org.apache.druid.data.input.kafka.KafkaTopicPartition; +import org.apache.druid.indexing.common.actions.SegmentTransactionalAppendFromShareInboxAction; +import org.apache.druid.indexing.common.actions.TaskActionClient; +import org.apache.druid.indexing.overlord.SegmentPublishResult; +import org.apache.druid.indexing.overlord.ShareInboxCompletionRequest; +import org.apache.druid.indexing.overlord.ShareInboxManifest; +import org.apache.druid.indexing.seekablestream.StreamChunkReader; +import org.apache.druid.indexing.seekablestream.common.OrderedPartitionableRecord; +import org.apache.druid.java.util.common.ISE; +import org.apache.druid.segment.SegmentSchemaMapping; +import org.apache.druid.segment.realtime.appenderator.AppenderatorDriverAddResult; +import org.apache.druid.segment.realtime.appenderator.StreamAppenderatorDriver; +import org.apache.druid.segment.realtime.appenderator.TransactionalSegmentPublisher; +import org.apache.druid.timeline.DataSegment; + +import javax.annotation.Nullable; +import java.io.IOException; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.TreeMap; +import java.util.UUID; + +final class AppenderatorShareInboxBatchHandler implements ShareInboxBatchHandler +{ + private final StreamAppenderatorDriver driver; + private final StreamChunkReader<KafkaRecordEntity> chunkReader; + private final TaskActionClient taskActionClient; + private final String taskId; + + AppenderatorShareInboxBatchHandler( + StreamAppenderatorDriver driver, + StreamChunkReader<KafkaRecordEntity> chunkReader, + TaskActionClient taskActionClient, + String taskId + ) + { + this.driver = driver; + this.chunkReader = chunkReader; + this.taskActionClient = taskActionClient; + this.taskId = taskId; + } + + @Override + public void process( + List<ShareInboxManifest> manifests, + List<OrderedPartitionableRecord<KafkaTopicPartition, Long, KafkaRecordEntity>> records, + ShareInboxClaimLease claimLease + ) throws Exception + { + final String completionId = UUID.randomUUID().toString(); + final String sequenceName = taskId + "-share-inbox-" + completionId; + final Supplier<Committer> committerSupplier = () -> new Committer() + { + @Override + public Object getMetadata() + { + return Map.of("shareInboxCompletionId", completionId); + } + + @Override + public void run() + { + } + }; + + boolean rowsAdded = false; + for (OrderedPartitionableRecord<KafkaTopicPartition, Long, KafkaRecordEntity> record : records) { + claimLease.assertValid(); + final List<InputRow> rows = chunkReader.parse(kafkaEntities(record), false); + for (InputRow row : rows) { + claimLease.assertValid(); + final AppenderatorDriverAddResult addResult = driver.add( + row, + sequenceName, + committerSupplier, + true, + false + ); + if (!addResult.isOk()) { + throw new ISE("Could not allocate segment for share inbox row with timestamp[%s]", row.getTimestamp()); + } + rowsAdded = true; + } + } + if (rowsAdded) { + claimLease.assertValid(); + driver.persist(committerSupplier.get()); + } + + final ShareInboxCompletionRequest completionRequest = completionRequest(manifests, completionId); + final TransactionalSegmentPublisher publisher = new TransactionalSegmentPublisher() + { + @Override + public SegmentPublishResult publishAnnotatedSegments( + @Nullable Set<DataSegment> mustBeNullOrEmptyOverwriteSegments, + Set<DataSegment> segmentsToPush, + @Nullable Object commitMetadata, + SegmentSchemaMapping segmentSchemaMapping + ) throws IOException + { + claimLease.assertValid(); + return taskActionClient.submit( + new SegmentTransactionalAppendFromShareInboxAction( + segmentsToPush, + completionRequest, + segmentSchemaMapping + ) + ); + } + + @Override + public boolean supportsEmptyPublish() + { + return true; + } + }; + + claimLease.assertValid(); + driver.publish(publisher, committerSupplier.get(), List.of(sequenceName)).get(); Review Comment: [P2] Hand published segments off before retaining the next batches **Finding:** The publish future is awaited and discarded, and no share-group path calls registerHandoff or drops the published segments. StreamAppenderatorDriver.publish removes sequence tracking, but it does not remove the underlying StreamAppenderator sinks; registerHandoff is the path that drops them once a Historical serves them. Since this task loops indefinitely and creates a fresh sequence for each claim, completed batches retain their persisted files, mapped indexes, and sink bookkeeping until task termination. A normally running ingestion task therefore grows local resource usage without bound even after the segments are available on Historicals. **Suggestion:** Register handoff using the publication result and manage the resulting futures so completed segments are dropped and handoff failures are observed. ########## indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentTransactionalAppendFromShareInboxAction.java: ########## @@ -0,0 +1,178 @@ +/* + * 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.common.actions; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.core.type.TypeReference; +import org.apache.druid.error.DruidException; +import org.apache.druid.error.InvalidInput; +import org.apache.druid.indexing.common.TaskLock; +import org.apache.druid.indexing.common.TaskLockType; +import org.apache.druid.indexing.common.task.IndexTaskUtils; +import org.apache.druid.indexing.common.task.PendingSegmentAllocatingTask; +import org.apache.druid.indexing.common.task.Task; +import org.apache.druid.indexing.overlord.CriticalAction; +import org.apache.druid.indexing.overlord.SegmentPublishResult; +import org.apache.druid.indexing.overlord.ShareInboxCompletionRequest; +import org.apache.druid.metadata.ReplaceTaskLock; +import org.apache.druid.segment.SegmentSchemaMapping; +import org.apache.druid.timeline.DataSegment; + +import javax.annotation.Nullable; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Set; +import java.util.stream.Collectors; + +public class SegmentTransactionalAppendFromShareInboxAction implements TaskAction<SegmentPublishResult> +{ + public static final String TYPE = "segmentTransactionalAppendFromShareInbox"; + + private final Set<DataSegment> segments; + private final ShareInboxCompletionRequest completionRequest; + @Nullable + private final SegmentSchemaMapping segmentSchemaMapping; + + @JsonCreator + public SegmentTransactionalAppendFromShareInboxAction( + @JsonProperty("segments") Set<DataSegment> segments, + @JsonProperty("completionRequest") ShareInboxCompletionRequest completionRequest, + @JsonProperty("segmentSchemaMapping") @Nullable SegmentSchemaMapping segmentSchemaMapping + ) + { + this.segments = Set.copyOf(Objects.requireNonNull(segments, "segments")); + this.completionRequest = Objects.requireNonNull(completionRequest, "completionRequest"); + this.segmentSchemaMapping = segmentSchemaMapping; + } + + @JsonProperty + public Set<DataSegment> getSegments() + { + return segments; + } + + @JsonProperty + public ShareInboxCompletionRequest getCompletionRequest() + { + return completionRequest; + } + + @Nullable + @JsonProperty + public SegmentSchemaMapping getSegmentSchemaMapping() + { + return segmentSchemaMapping; + } + + @Override + public TypeReference<SegmentPublishResult> getReturnTypeReference() + { + return new TypeReference<>() {}; + } + + @Override + public SegmentPublishResult perform(Task task, TaskActionToolbox toolbox) + { + if (!(task instanceof PendingSegmentAllocatingTask)) { + throw DruidException.defensive( + "Task[%s] of type[%s] cannot append share inbox segments because it does not implement PendingSegmentAllocatingTask.", + task.getId(), + task.getType() + ); + } + ClaimShareInboxManifestsAction.validateTask( + task, + completionRequest.getDataSource(), + completionRequest.getTaskId() + ); + + final List<TaskLock> locks = toolbox.getTaskLockbox().findLocksForTask(task); + for (TaskLock lock : locks) { + if (lock.getType() != TaskLockType.APPEND) { Review Comment: [P2] Keep the accepted task lock context compatible with publication **Finding:** The prior non-APPEND lock finding remains in the new inbox publisher. ShareGroupIndexTask preserves an explicit useConcurrentLocks=false context, and the runner then allocates with EXCLUSIVE or SHARED locks through determineLockTypeForAppend. This action rejects those locks unconditionally. The accepted task consequently stages and acknowledges Kafka records but fails at its first non-empty segment publication; retrying the same configuration cannot drain the inbox. **Suggestion:** Reject or normalize non-APPEND contexts before acquisition starts, or implement a completion publisher compatible with the effective lock type and test useConcurrentLocks=false. ########## server/src/main/java/org/apache/druid/metadata/IndexerSQLMetadataStorageCoordinator.java: ########## @@ -163,6 +179,590 @@ public Set<String> retrieveAllDatasourceNames() ); } + @Override + public ShareInboxStageResult stageShareInboxBatch(ShareInboxBatch batch) + { + return connector.retryTransaction( + (handle, status) -> { + try { + return stageShareInboxBatch(handle, batch); + } + catch (RuntimeException e) { + if (connector.isUniqueConstraintViolation(e)) { + throw new RetryTransactionException("Concurrent share inbox receipt update"); + } + throw e; + } + }, + SQLMetadataConnector.QUIET_RETRIES, + SQLMetadataConnector.DEFAULT_MAX_TRIES + ); + } + + @Override + public ShareInboxClaimResult claimShareInboxManifests(ShareInboxClaimRequest request) + { + return connector.retryTransaction( + (handle, status) -> claimShareInboxManifests(handle, request), + SQLMetadataConnector.QUIET_RETRIES, + SQLMetadataConnector.DEFAULT_MAX_TRIES + ); + } + + private ShareInboxClaimResult claimShareInboxManifests(Handle handle, ShareInboxClaimRequest request) + { + validateShareInboxGeneration( + handle, + request.getDataSource(), + request.getInboxId(), + request.getSpecFingerprint() + ); + final DateTime now = DateTimes.nowUtc(); + final String nowString = now.toString(); + final String claimExpiresAt = now.plus(request.getClaimDurationMillis()).toString(); + final List<ShareInboxCandidate> candidates = handle.createQuery( + StringUtils.format( + "SELECT manifest_id, object_size, record_count, claim_epoch FROM %s " + + "WHERE data_source = :dataSource AND inbox_id = :inboxId " + + "AND spec_fingerprint = :specFingerprint " + + "AND (state = 'STAGED' OR (state = 'CLAIMED' AND claim_expires_at <= :now)) " + + "ORDER BY created_date, manifest_id %s", + dbTables.getShareInboxTable(), + connector.limitClause(request.getMaxManifests()) + ) + ) + .bind("dataSource", request.getDataSource()) + .bind("inboxId", request.getInboxId()) + .bind("specFingerprint", request.getSpecFingerprint()) + .bind("now", nowString) + .map( + (index, resultSet, context) -> + new ShareInboxCandidate( + resultSet.getString("manifest_id"), + resultSet.getLong("object_size"), + resultSet.getInt("record_count"), + resultSet.getLong("claim_epoch") + ) + ) + .list(); + + final List<ShareInboxManifest> claimed = new ArrayList<>(); + long claimedBytes = 0; + int claimedRecords = 0; + for (ShareInboxCandidate candidate : candidates) { + if (claimedBytes + candidate.objectSize > request.getMaxBytes() Review Comment: [P2] Ensure every staged manifest can fit the processing limits **Finding:** A manifest larger than maxRecords or maxBytes is skipped even when the claim is empty. Staging uses separate maxStagingRecords/maxStagingBytes limits and never splits a partition batch to the processing limits. For example, maxProcessingRecords=100 with a 500-record polled partition batch stages and acknowledges all 500 records, but that manifest is never claimable. Once the oldest maxManifests candidates are oversized, the SQL LIMIT also prevents later smaller manifests from being considered, so the inbox can stop progressing indefinitely. **Suggestion:** Split or validate staged batches against processing bounds, or allow one oversized manifest to make progress, and ensure oversized candidates cannot permanently hide later work. ########## extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/ShareGroupIndexTaskRunner.java: ########## @@ -0,0 +1,618 @@ +/* + * 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 org.apache.druid.common.config.Configs; +import org.apache.druid.data.input.InputFormat; +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.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.SegmentLockAcquireAction; +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.seekablestream.SeekableStreamIndexTask; +import org.apache.druid.indexing.seekablestream.StreamChunkReader; +import org.apache.druid.indexing.seekablestream.common.AcknowledgingRecordSupplier; +import org.apache.druid.java.util.common.ISE; +import org.apache.druid.java.util.common.concurrent.Execs; +import org.apache.druid.java.util.common.logger.Logger; +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.SegmentIdWithShardSpec; +import org.apache.druid.segment.realtime.appenderator.StreamAppenderatorDriver; +import org.apache.druid.storage.StorageConnector; +import org.apache.kafka.clients.admin.Admin; +import org.apache.kafka.clients.consumer.ConsumerConfig; + +import javax.annotation.Nullable; +import java.io.IOException; +import java.time.Duration; +import java.util.Properties; +import java.util.concurrent.ArrayBlockingQueue; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Future; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; + +/** + * Runs durable acquisition and inbox processing for {@link ShareGroupIndexTask}. + */ +public class ShareGroupIndexTaskRunner +{ + private static final Logger log = new Logger(ShareGroupIndexTaskRunner.class); + private static final long DEFAULT_MAX_POLL_INTERVAL_MS = 300_000L; + + private final ShareGroupIndexTask task; + private final TaskToolbox toolbox; + private final ObjectMapper configMapper; + private final LockGranularity lockGranularity; + private final TaskLockType lockType; + private final Function<ShareGroupIndexTaskIOConfig, + AcknowledgingRecordSupplier<KafkaTopicPartition, Long, KafkaRecordEntity>> supplierFactory; + private final Function<ShareGroupIndexTaskIOConfig, ShareGroupSourceIdentity> sourceIdentityFactory; + + @Nullable + private ShareGroupSourceIdentity sourceIdentity; + @Nullable + private String specFingerprint; + + private final AtomicReference<AcknowledgingRecordSupplier<KafkaTopicPartition, Long, KafkaRecordEntity>> + activeSupplier = new AtomicReference<>(); + private final AtomicReference<ShareGroupAcquisitionLoop> activeAcquisitionLoop = new AtomicReference<>(); + private final AtomicReference<ShareInboxProcessor> activeInboxProcessor = new AtomicReference<>(); + + ShareGroupIndexTaskRunner( + ShareGroupIndexTask task, + TaskToolbox toolbox, + ObjectMapper configMapper + ) + { + this(task, toolbox, configMapper, null, null); + } + + @VisibleForTesting + ShareGroupIndexTaskRunner( + ShareGroupIndexTask task, + TaskToolbox toolbox, + ObjectMapper configMapper, + @Nullable Function<ShareGroupIndexTaskIOConfig, + AcknowledgingRecordSupplier<KafkaTopicPartition, Long, KafkaRecordEntity>> supplierFactory + ) + { + this(task, toolbox, configMapper, supplierFactory, null); + } + + @VisibleForTesting + ShareGroupIndexTaskRunner( + ShareGroupIndexTask task, + TaskToolbox toolbox, + ObjectMapper configMapper, + @Nullable Function<ShareGroupIndexTaskIOConfig, + AcknowledgingRecordSupplier<KafkaTopicPartition, Long, KafkaRecordEntity>> supplierFactory, + @Nullable Function<ShareGroupIndexTaskIOConfig, ShareGroupSourceIdentity> sourceIdentityFactory + ) + { + this.task = task; + this.toolbox = toolbox; + this.configMapper = configMapper; + 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; Review Comment: [P2] Reject the unsupported APPEND and segment-lock combination **Finding:** The prior lock-granularity finding is still present: forceTimeChunkLock=false selects SEGMENT while the default concurrent-lock context selects APPEND. Both allocation and restored-segment acquisition receive that pair. TaskLockbox explicitly rejects APPEND/REPLACE requests with non-TIME_CHUNK granularity when an overlapping non-reusable lock exists, so running this accepted configuration alongside ordinary concurrent ingestion or replacement can prevent allocation or restoration. The restored-lock unit tests mock the action result and therefore do not exercise this lockbox restriction. **Suggestion:** Require time-chunk APPEND locks for this task or implement a supported alternative consistently across allocation, restoration, and inbox publication, 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]
