github-advanced-security[bot] commented on code in PR #19311: URL: https://github.com/apache/druid/pull/19311#discussion_r4124329095
########## extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/ShareGroupBatchStore.java: ########## @@ -0,0 +1,231 @@ +/* + * 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.seekablestream.common.OrderedPartitionableRecord; +import org.apache.druid.storage.StorageConnector; + +import java.io.FilterOutputStream; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.nio.charset.StandardCharsets; +import java.security.DigestInputStream; +import java.security.DigestOutputStream; +import java.security.MessageDigest; +import java.util.Arrays; +import java.util.Base64; +import java.util.HexFormat; +import java.util.List; +import java.util.UUID; + +final class ShareGroupBatchStore +{ + static final class StoredBatch + { + private final String path; + private final String sha256; + private final long size; + private final ShareGroupRawBatch batch; + + private StoredBatch(String path, String sha256, long size, ShareGroupRawBatch batch) + { + this.path = path; + this.sha256 = sha256; + this.size = size; + this.batch = batch; + } + + String getPath() + { + return path; + } + + String getSha256() + { + return sha256; + } + + long getSize() + { + return size; + } + + ShareGroupRawBatch getBatch() + { + return batch; + } + } + + private static final class CountingOutputStream extends FilterOutputStream + { + private long count; + + private CountingOutputStream(OutputStream outputStream) + { + super(outputStream); + } + + @Override + public void write(int value) throws IOException + { + out.write(value); + count++; + } + + @Override + public void write(byte[] bytes, int offset, int length) throws IOException + { + out.write(bytes, offset, length); + count += length; + } + + private long getCount() + { + return count; + } + } + + private final StorageConnector storageConnector; + private final String taskId; + private final ShareGroupRawBatchWriter writer = new ShareGroupRawBatchWriter(); + private final ShareGroupRawBatchReader reader = new ShareGroupRawBatchReader(); + + ShareGroupBatchStore(StorageConnector storageConnector, String taskId) + { + this.storageConnector = Preconditions.checkNotNull(storageConnector, "storageConnector"); + this.taskId = Preconditions.checkNotNull(taskId, "taskId"); + } + + StoredBatch store( + List<OrderedPartitionableRecord<KafkaTopicPartition, Long, KafkaRecordEntity>> records + ) throws IOException + { + return store(ShareGroupRawBatch.fromOrderedRecords(records)); + } + + StoredBatch store(ShareGroupRawBatch batch) throws IOException + { + final String path = objectPath(batch); + final MessageDigest writeDigest = ShareGroupRawBatchWriter.newDigest(); + final CountingOutputStream countingOutput; + try (OutputStream storageOutput = storageConnector.write(path)) { + countingOutput = new CountingOutputStream(storageOutput); + writer.write(batch, new DigestOutputStream(countingOutput, writeDigest)); Review Comment: ## CodeQL / Potential output resource leak This DigestOutputStream is not always closed on method exit. [Show more details](https://github.com/apache/druid/security/code-scanning/12022) ########## server/src/test/java/org/apache/druid/metadata/IndexerSQLMetadataStorageCoordinatorTest.java: ########## @@ -4849,6 +5225,135 @@ .build(); } + private static ShareInboxBatch shareInboxBatch(String manifestId, String fingerprint, List<Long> offsets) + { + return new ShareInboxBatch( + "datasource", + "inbox", + "group", + "cluster", + "topic-id", + "topic", + 0, + fingerprint, + manifestId, + "path/" + manifestId, + "hash-" + manifestId, + 100, + offsets, + 4_096 + ); + } + + private static ShareInboxClaimRequest shareInboxClaimRequest(String owner, long maxBytes) + { + return new ShareInboxClaimRequest( + "datasource", + "inbox", + "fingerprint", + owner, + 10, + 10, + maxBytes, + 60_000 + ); + } + + private static ShareInboxRenewRequest shareInboxRenewRequest(String owner, Map<String, Long> claims) + { + return new ShareInboxRenewRequest( + "datasource", + "inbox", + "fingerprint", + owner, + claims, + 60_000 + ); + } + + private static ShareInboxCompletionRequest shareInboxCompletionRequest( + ShareInboxManifest manifest, + String completionId + ) + { + return new ShareInboxCompletionRequest( + manifest.getDataSource(), + manifest.getInboxId(), + manifest.getSpecFingerprint(), + manifest.getClaimOwner(), + Map.of(manifest.getManifestId(), manifest.getClaimEpoch()), + completionId + ); + } + + private static DataSegment shareInboxSegment(String version) + { + return new DataSegment( + "datasource", + Intervals.of("2026-01-01/2026-01-02"), + version, + Map.of(), + List.of("dim"), + List.of("metric"), + new LinearShardSpec(0), + 9, + 100 + ); Review Comment: ## CodeQL / Deprecated method or constructor invocation Invoking [DataSegment.DataSegment](1) should be avoided because it has been deprecated. [Show more details](https://github.com/apache/druid/security/code-scanning/12031) ########## indexing-service/src/main/java/org/apache/druid/indexing/common/actions/StageShareInboxBatchAction.java: ########## @@ -0,0 +1,76 @@ +/* + * 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.InvalidInput; +import org.apache.druid.indexing.common.task.Task; +import org.apache.druid.indexing.overlord.ShareInboxBatch; +import org.apache.druid.indexing.overlord.ShareInboxStageResult; + +import java.util.Objects; + +public class StageShareInboxBatchAction implements TaskAction<ShareInboxStageResult> +{ + public static final String TYPE = "stageShareInboxBatch"; + + private final ShareInboxBatch batch; + + @JsonCreator + public StageShareInboxBatchAction(@JsonProperty("batch") ShareInboxBatch batch) + { + this.batch = Objects.requireNonNull(batch, "batch"); + } + + @JsonProperty + public ShareInboxBatch getBatch() + { + return batch; + } + + @Override + public TypeReference<ShareInboxStageResult> getReturnTypeReference() + { + return new TypeReference<>() {}; + } + + @Override + public ShareInboxStageResult perform(Task task, TaskActionToolbox toolbox) + { + if (!task.getDataSource().equals(batch.getDataSource())) { + throw InvalidInput.exception( + "Task[%s] cannot stage share inbox records for datasource[%s]", + task.getId(), + batch.getDataSource() + ); + } + return toolbox.getIndexerMetadataStorageCoordinator().stageShareInboxBatch(batch); + } + + @Override + public String toString() + { + return "StageShareInboxBatchAction{" + + "batch=" + batch + Review Comment: ## CodeQL / Use of default toString() Default toString(): ShareInboxBatch inherits toString() from Object, and so is not suitable for printing. [Show more details](https://github.com/apache/druid/security/code-scanning/12028) ########## 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) { + throw InvalidInput.exception( + "Cannot use action[%s] for task[%s] as it is holding a lock of type[%s] instead of [APPEND].", + TYPE, + task.getId(), + lock.getType() + ); + } + } + + final PendingSegmentAllocatingTask allocatingTask = (PendingSegmentAllocatingTask) task; + final SegmentPublishResult result; + if (segments.isEmpty()) { + result = commit(toolbox, allocatingTask.getTaskAllocatorId(), Map.of()); + } else { + TaskLocks.checkLockCoversSegments(task, toolbox.getTaskLockbox(), segments); + final Map<DataSegment, ReplaceTaskLock> segmentToReplaceLock = TaskLocks.findReplaceLocksCoveringSegments( + task.getDataSource(), + toolbox.getTaskLockbox(), + segments + ); + try { + result = toolbox.getTaskLockbox().doInCriticalSection( + task, + segments.stream().map(DataSegment::getInterval).collect(Collectors.toSet()), + CriticalAction.<SegmentPublishResult>builder() + .onValidLocks(() -> commit(toolbox, allocatingTask.getTaskAllocatorId(), segmentToReplaceLock)) + .onInvalidLocks( + () -> SegmentPublishResult.fail( + "Invalid task locks. Maybe they are revoked by a higher priority task." + + " Please check the overlord log for details." + ) + ) + .build() + ); + } + catch (Exception e) { + throw new RuntimeException(e); + } + } + + IndexTaskUtils.emitSegmentPublishMetrics(result, task, toolbox); + return result; + } + + private SegmentPublishResult commit( + TaskActionToolbox toolbox, + String taskAllocatorId, + Map<DataSegment, ReplaceTaskLock> segmentToReplaceLock + ) + { + return toolbox.getIndexerMetadataStorageCoordinator().commitAppendSegmentsAndShareInbox( + segments, + segmentToReplaceLock, + taskAllocatorId, + segmentSchemaMapping, + completionRequest + ); + } + + @Override + public String toString() + { + return "SegmentTransactionalAppendFromShareInboxAction{" + + "segments=" + segments + + ", completionRequest=" + completionRequest + Review Comment: ## CodeQL / Use of default toString() Default toString(): ShareInboxCompletionRequest inherits toString() from Object, and so is not suitable for printing. [Show more details](https://github.com/apache/druid/security/code-scanning/12027) ########## indexing-service/src/main/java/org/apache/druid/indexing/common/actions/ClaimShareInboxManifestsAction.java: ########## @@ -0,0 +1,82 @@ +/* + * 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.InvalidInput; +import org.apache.druid.indexing.common.task.Task; +import org.apache.druid.indexing.overlord.ShareInboxClaimRequest; +import org.apache.druid.indexing.overlord.ShareInboxClaimResult; + +import java.util.Objects; + +public class ClaimShareInboxManifestsAction implements TaskAction<ShareInboxClaimResult> +{ + public static final String TYPE = "claimShareInboxManifests"; + + private final ShareInboxClaimRequest request; + + @JsonCreator + public ClaimShareInboxManifestsAction(@JsonProperty("request") ShareInboxClaimRequest request) + { + this.request = Objects.requireNonNull(request, "request"); + } + + @JsonProperty + public ShareInboxClaimRequest getRequest() + { + return request; + } + + @Override + public TypeReference<ShareInboxClaimResult> getReturnTypeReference() + { + return new TypeReference<>() {}; + } + + @Override + public ShareInboxClaimResult perform(Task task, TaskActionToolbox toolbox) + { + validateTask(task, request.getDataSource(), request.getClaimOwner()); + return toolbox.getIndexerMetadataStorageCoordinator().claimShareInboxManifests(request); + } + + static void validateTask(Task task, String dataSource, String claimOwner) + { + if (!task.getDataSource().equals(dataSource) || !task.getId().equals(claimOwner)) { + throw InvalidInput.exception( + "Task[%s] cannot manage share inbox claims for datasource[%s] as owner[%s]", + task.getId(), + dataSource, + claimOwner + ); + } + } + + @Override + public String toString() + { + return "ClaimShareInboxManifestsAction{" + + "request=" + request + Review Comment: ## CodeQL / Use of default toString() Default toString(): ShareInboxClaimRequest inherits toString() from Object, and so is not suitable for printing. [Show more details](https://github.com/apache/druid/security/code-scanning/12025) ########## extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/ShareGroupAcquisitionLoop.java: ########## @@ -0,0 +1,264 @@ +/* + * 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.seekablestream.common.AcknowledgeType; +import org.apache.druid.indexing.seekablestream.common.AcknowledgingRecordSupplier; +import org.apache.druid.indexing.seekablestream.common.OrderedPartitionableRecord; +import org.apache.kafka.common.errors.WakeupException; + +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.Executor; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.ToLongFunction; + +final class ShareGroupAcquisitionLoop +{ + private final String topic; + private final AcknowledgingRecordSupplier<KafkaTopicPartition, Long, KafkaRecordEntity> recordSupplier; + private final ShareGroupBatchStager batchStager; + private final Executor stagingExecutor; + private final ShareGroupAcquisitionRegistry registry; + private final ToLongFunction<OrderedPartitionableRecord<KafkaTopicPartition, Long, KafkaRecordEntity>> byteEstimator; + private final long pollTimeoutMs; + private final long maxPollIntervalMs; + private final BlockingQueue<ShareGroupStagingCompletion> completions = new LinkedBlockingQueue<>(); + private final AtomicBoolean stopRequested = new AtomicBoolean(); + private Thread ownerThread; + + ShareGroupAcquisitionLoop( + String topic, + AcknowledgingRecordSupplier<KafkaTopicPartition, Long, KafkaRecordEntity> recordSupplier, + ShareGroupBatchStager batchStager, + Executor stagingExecutor, + ShareGroupAcquisitionRegistry registry, + ToLongFunction<OrderedPartitionableRecord<KafkaTopicPartition, Long, KafkaRecordEntity>> byteEstimator, + long pollTimeoutMs, + long maxPollIntervalMs + ) + { + Preconditions.checkArgument(pollTimeoutMs >= 0, "pollTimeoutMs must not be negative"); + Preconditions.checkArgument(maxPollIntervalMs > 0, "maxPollIntervalMs must be positive"); + this.topic = Preconditions.checkNotNull(topic, "topic"); + this.recordSupplier = Preconditions.checkNotNull(recordSupplier, "recordSupplier"); + this.batchStager = Preconditions.checkNotNull(batchStager, "batchStager"); + this.stagingExecutor = Preconditions.checkNotNull(stagingExecutor, "stagingExecutor"); + this.registry = Preconditions.checkNotNull(registry, "registry"); + this.byteEstimator = Preconditions.checkNotNull(byteEstimator, "byteEstimator"); + this.pollTimeoutMs = pollTimeoutMs; + this.maxPollIntervalMs = maxPollIntervalMs; + } + + void run() throws Exception + { + checkOwnerThread(); + recordSupplier.subscribe(Set.of(topic)); + while (!stopRequested.get()) { + try { + runOnce(); + } + catch (WakeupException e) { + if (!stopRequested.get()) { + throw e; + } + } + } + } + + void runOnce() throws InterruptedException + { + checkOwnerThread(); + drainCompletions(); + final List<OrderedPartitionableRecord<KafkaTopicPartition, Long, KafkaRecordEntity>> records = + recordSupplier.poll(pollTimeoutMs); + drainCompletions(); + if (records.isEmpty()) { + return; + } + + final long acquisitionLockTimeoutMs = recordSupplier.acquisitionLockTimeoutMs().orElse(-1); + final Map<ShareGroupRecordIdentity, + OrderedPartitionableRecord<KafkaTopicPartition, Long, KafkaRecordEntity>> tracked = new LinkedHashMap<>(); + final List<OrderedPartitionableRecord<KafkaTopicPartition, Long, KafkaRecordEntity>> saturated = + new ArrayList<>(); + final Map<KafkaTopicPartition, + List<OrderedPartitionableRecord<KafkaTopicPartition, Long, KafkaRecordEntity>>> newRecordsByPartition = + new LinkedHashMap<>(); + + for (OrderedPartitionableRecord<KafkaTopicPartition, Long, KafkaRecordEntity> record : records) { + final ShareGroupRecordIdentity identity = ShareGroupRecordIdentity.from(record); + final ShareGroupAcquisitionRegistry.Observation observation = registry.observe( + record, + byteEstimator.applyAsLong(record), + acquisitionLockTimeoutMs, + maxPollIntervalMs + ); + if (observation == ShareGroupAcquisitionRegistry.Observation.SATURATED) { + saturated.add(record); + } else { + tracked.put(identity, record); + if (observation == ShareGroupAcquisitionRegistry.Observation.ADDED) { + newRecordsByPartition.computeIfAbsent(record.getPartitionId(), ignored -> new ArrayList<>()).add(record); + } + } + } + + for (List<OrderedPartitionableRecord<KafkaTopicPartition, Long, KafkaRecordEntity>> partitionRecords + : newRecordsByPartition.values()) { + submit(partitionRecords); + } + + waitForAcknowledgementDecisions(tracked.keySet()); + + boolean terminalAcknowledgementAssigned = false; + for (OrderedPartitionableRecord<KafkaTopicPartition, Long, KafkaRecordEntity> record : saturated) { + recordSupplier.acknowledge( + record.getPartitionId(), + record.getSequenceNumber(), + AcknowledgeType.RELEASE + ); + terminalAcknowledgementAssigned = true; + } + for (Map.Entry<ShareGroupRecordIdentity, + OrderedPartitionableRecord<KafkaTopicPartition, Long, KafkaRecordEntity>> trackedRecord : tracked.entrySet()) { + final AcknowledgeType type = registry.acknowledgementFor( + trackedRecord.getKey(), + stopRequested.get() + ).orElseThrow(() -> new IllegalStateException("No acknowledgement decision for " + trackedRecord.getKey())); + final OrderedPartitionableRecord<KafkaTopicPartition, Long, KafkaRecordEntity> record = trackedRecord.getValue(); + recordSupplier.acknowledge(record.getPartitionId(), record.getSequenceNumber(), type); + registry.acknowledgementAssigned(trackedRecord.getKey(), type); + terminalAcknowledgementAssigned |= type != AcknowledgeType.RENEW; + } + + if (terminalAcknowledgementAssigned) { + registry.acknowledgementsFlushed(recordSupplier.flushAcknowledgementsSync()); + } + } + + void requestStop() + { + stopRequested.set(true); + completions.offer(ShareGroupStagingCompletion.stop()); Review Comment: ## CodeQL / Ignored error status of call Method requestStop ignores exceptional return value of BlockingQueue<ShareGroupStagingCompletion>.offer. [Show more details](https://github.com/apache/druid/security/code-scanning/12024) ########## indexing-service/src/test/java/org/apache/druid/indexing/common/actions/SegmentTransactionalAppendFromShareInboxActionTest.java: ########## @@ -0,0 +1,225 @@ +/* + * 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 org.apache.druid.error.DruidException; +import org.apache.druid.indexing.common.TaskLockType; +import org.apache.druid.indexing.common.task.NoopTask; +import org.apache.druid.indexing.common.task.Task; +import org.apache.druid.indexing.overlord.GlobalTaskLockbox; +import org.apache.druid.indexing.overlord.IndexerMetadataStorageCoordinator; +import org.apache.druid.indexing.overlord.LockResult; +import org.apache.druid.indexing.overlord.SegmentPublishResult; +import org.apache.druid.indexing.overlord.ShareInboxBatch; +import org.apache.druid.indexing.overlord.ShareInboxClaimRequest; +import org.apache.druid.indexing.overlord.ShareInboxCompletionRequest; +import org.apache.druid.indexing.overlord.ShareInboxManifest; +import org.apache.druid.indexing.overlord.TimeChunkLockRequest; +import org.apache.druid.java.util.common.Intervals; +import org.apache.druid.java.util.metrics.StubServiceEmitter; +import org.apache.druid.segment.TestHelper; +import org.apache.druid.timeline.DataSegment; +import org.apache.druid.timeline.partition.LinearShardSpec; +import org.joda.time.Interval; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; + +import java.util.List; +import java.util.Map; +import java.util.Set; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class SegmentTransactionalAppendFromShareInboxActionTest +{ + private static final String DATA_SOURCE = "datasource"; + private static final String TASK_ID = "task-1"; + private static final Interval INTERVAL = Intervals.of("2026-01-01/2026-01-02"); + + @RegisterExtension + public final TaskActionTestKit actionTestKit = new TaskActionTestKit(); + + @Test + public void testSerde() throws Exception + { + final SegmentTransactionalAppendFromShareInboxAction original = new SegmentTransactionalAppendFromShareInboxAction( + Set.of(), + completionRequest("manifest-1", 2L), + null + ); + + final TaskAction<?> deserialized = TestHelper.JSON_MAPPER.readValue( + TestHelper.JSON_MAPPER.writeValueAsBytes(original), + TaskAction.class + ); + + Assertions.assertInstanceOf(SegmentTransactionalAppendFromShareInboxAction.class, deserialized); + final SegmentTransactionalAppendFromShareInboxAction action = + (SegmentTransactionalAppendFromShareInboxAction) deserialized; + Assertions.assertEquals(original.getSegments(), action.getSegments()); + Assertions.assertEquals(original.getCompletionRequest().getClaims(), action.getCompletionRequest().getClaims()); + } + + @Test + public void testPerformCompletesZeroRowManifest() + { + final NoopTask task = task(); + final ShareInboxCompletionRequest request = completionRequest("manifest-1", 2L); + final SegmentTransactionalAppendFromShareInboxAction action = + new SegmentTransactionalAppendFromShareInboxAction(Set.of(), request, null); + final TaskActionToolbox toolbox = mock(TaskActionToolbox.class); + final GlobalTaskLockbox lockbox = mock(GlobalTaskLockbox.class); + final IndexerMetadataStorageCoordinator coordinator = mock(IndexerMetadataStorageCoordinator.class); + final SegmentPublishResult expected = SegmentPublishResult.ok(Set.of()); + when(toolbox.getTaskLockbox()).thenReturn(lockbox); + when(toolbox.getIndexerMetadataStorageCoordinator()).thenReturn(coordinator); + when(toolbox.getEmitter()).thenReturn(new StubServiceEmitter()); + when(lockbox.findLocksForTask(task)).thenReturn(List.of()); + when(coordinator.commitAppendSegmentsAndShareInbox(Set.of(), Map.of(), TASK_ID, null, request)) + .thenReturn(expected); + + Assertions.assertSame(expected, action.perform(task, toolbox)); + verify(coordinator).commitAppendSegmentsAndShareInbox(Set.of(), Map.of(), TASK_ID, null, request); + } + + @Test + public void testPerformRejectsDifferentClaimOwner() + { + final NoopTask task = task(); + final ShareInboxCompletionRequest request = new ShareInboxCompletionRequest( + DATA_SOURCE, + "inbox", + "fingerprint", + "other-task", + Map.of("manifest-1", 2L), + "completion-1" + ); + final SegmentTransactionalAppendFromShareInboxAction action = + new SegmentTransactionalAppendFromShareInboxAction(Set.of(), request, null); + final TaskActionToolbox toolbox = mock(TaskActionToolbox.class); + final IndexerMetadataStorageCoordinator coordinator = mock(IndexerMetadataStorageCoordinator.class); + when(toolbox.getIndexerMetadataStorageCoordinator()).thenReturn(coordinator); + + Assertions.assertThrows(DruidException.class, () -> action.perform(task, toolbox)); + verify(coordinator, never()).commitAppendSegmentsAndShareInbox(any(), any(), any(), any(), any()); + } + + @Test + public void testPerformRejectsTaskWithoutPendingSegmentAllocation() + { + final Task task = mock(Task.class); + final SegmentTransactionalAppendFromShareInboxAction action = + new SegmentTransactionalAppendFromShareInboxAction(Set.of(), completionRequest("manifest-1", 2L), null); + when(task.getId()).thenReturn(TASK_ID); + when(task.getType()).thenReturn("test"); + + Assertions.assertThrows(DruidException.class, () -> action.perform(task, mock(TaskActionToolbox.class))); + } + + @Test + public void testPerformPublishesSegmentAndCompletesInbox() throws Exception + { + actionTestKit.getTestDerbyConnector().createShareReceiptsTable(); + actionTestKit.getTestDerbyConnector().createShareInboxTable(); + final NoopTask task = task(); + actionTestKit.getTaskLockbox().add(task); + final LockResult lockResult = actionTestKit.getTaskLockbox().lock( + task, + new TimeChunkLockRequest(TaskLockType.APPEND, task, INTERVAL, null), + 5_000 + ); + Assertions.assertTrue(lockResult.isOk()); + final DataSegment segment = segment(lockResult.getTaskLock().getVersion()); + final IndexerMetadataStorageCoordinator coordinator = actionTestKit.getMetadataStorageCoordinator(); + coordinator.stageShareInboxBatch(batch()); + final ShareInboxManifest manifest = coordinator.claimShareInboxManifests( + new ShareInboxClaimRequest(DATA_SOURCE, "inbox", "fingerprint", TASK_ID, 1, 10, 10_000, 60_000) + ).getManifests().get(0); + final ShareInboxCompletionRequest request = completionRequest(manifest.getManifestId(), manifest.getClaimEpoch()); + final SegmentTransactionalAppendFromShareInboxAction action = + new SegmentTransactionalAppendFromShareInboxAction(Set.of(segment), request, null); + + final SegmentPublishResult first = action.perform(task, actionTestKit.getTaskActionToolbox()); + final SegmentPublishResult retry = action.perform(task, actionTestKit.getTaskActionToolbox()); + + Assertions.assertTrue(first.isSuccess()); + Assertions.assertEquals(Set.of(segment), first.getSegments()); + Assertions.assertTrue(retry.isSuccess()); + Assertions.assertTrue(retry.getSegments().isEmpty()); + Assertions.assertEquals(segment, coordinator.retrieveSegmentForId(segment.getId())); + } + + private static NoopTask task() + { + return new NoopTask(TASK_ID, null, DATA_SOURCE, 0, 0, null); + } + + private static ShareInboxCompletionRequest completionRequest(String manifestId, long claimEpoch) + { + return new ShareInboxCompletionRequest( + DATA_SOURCE, + "inbox", + "fingerprint", + TASK_ID, + Map.of(manifestId, claimEpoch), + "completion-1" + ); + } + + private static ShareInboxBatch batch() + { + return new ShareInboxBatch( + DATA_SOURCE, + "inbox", + "group", + "cluster", + "topic-id", + "topic", + 0, + "fingerprint", + "manifest-1", + "path/manifest-1", + "hash", + 1, + List.of(1L), + 100 + ); + } + + private static DataSegment segment(String version) + { + return new DataSegment( + DATA_SOURCE, + INTERVAL, + version, + Map.of(), + List.of("dim"), + List.of("metric"), + new LinearShardSpec(0), + 9, + 100 + ); Review Comment: ## CodeQL / Deprecated method or constructor invocation Invoking [DataSegment.DataSegment](1) should be avoided because it has been deprecated. [Show more details](https://github.com/apache/druid/security/code-scanning/12030) ########## extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/ShareGroupAcquisitionLoopTest.java: ########## @@ -0,0 +1,483 @@ +/* + * 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 org.apache.druid.data.input.kafka.KafkaRecordEntity; +import org.apache.druid.data.input.kafka.KafkaTopicPartition; +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.kafka.clients.consumer.ConsumerRecord; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Collection; +import java.util.Deque; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; + +public class ShareGroupAcquisitionLoopTest +{ + @Test + public void testRepeatedRenewalsStartOneStagingOperation() throws Exception + { + final ScriptedSupplier supplier = new ScriptedSupplier(Optional.empty()); + for (int i = 0; i < 11; i++) { + supplier.addPoll(record("topic", 0, 1)); + } + final RecordingExecutor executor = new RecordingExecutor(); + final AtomicInteger stageCount = new AtomicInteger(); + final ShareGroupAcquisitionLoop loop = loop( + supplier, + records -> stageCount.incrementAndGet(), + executor, + registry(10, 1_000) + ); + + for (int i = 0; i < 10; i++) { + loop.runOnce(); + } + Assertions.assertEquals(1, executor.size()); + Assertions.assertEquals(0, stageCount.get()); + Assertions.assertEquals(10, supplier.acknowledgements.size()); + Assertions.assertTrue( + supplier.acknowledgements.stream().allMatch(ack -> ack.type == AcknowledgeType.RENEW) + ); + + executor.runNext(); + loop.runOnce(); + + Assertions.assertEquals(1, stageCount.get()); + Assertions.assertEquals(AcknowledgeType.ACCEPT, supplier.acknowledgements.get(10).type); + Assertions.assertEquals(1, supplier.flushCount); + } + + @Test + public void testSaturatedRecordIsReleasedWithoutSecondSubmission() throws Exception + { + final ScriptedSupplier supplier = new ScriptedSupplier(Optional.empty()); + supplier.addPoll(record("topic", 0, 1), record("topic", 0, 2)); + final RecordingExecutor executor = new RecordingExecutor(); + final ShareGroupAcquisitionLoop loop = loop( + supplier, + records -> {}, + executor, + registry(1, 1_000) + ); + + loop.runOnce(); + + Assertions.assertEquals(1, executor.size()); + Assertions.assertEquals(AcknowledgeType.RELEASE, supplier.acknowledgement(2).type); + Assertions.assertEquals(AcknowledgeType.RENEW, supplier.acknowledgement(1).type); + Assertions.assertEquals(1, supplier.flushCount); + } + + @Test + public void testByteSaturationReleasesRecord() throws Exception + { + final ScriptedSupplier supplier = new ScriptedSupplier(Optional.empty()); + supplier.addPoll(record("topic", 0, 1), record("topic", 0, 2)); + final ShareGroupAcquisitionRegistry registry = registry(2, 1); + final RecordingExecutor executor = new RecordingExecutor(); + final ShareGroupAcquisitionLoop loop = new ShareGroupAcquisitionLoop( + "topic", + supplier, + records -> {}, + executor, + registry, + ignored -> 1, + 1, + 60_000 + ); + + loop.runOnce(); + + Assertions.assertEquals(1, executor.size()); + Assertions.assertEquals(AcknowledgeType.RELEASE, supplier.acknowledgement(2).type); + } + + @Test + public void testExecutorRejectionReleasesRecord() throws Exception + { + final ScriptedSupplier supplier = new ScriptedSupplier(Optional.of(60_000)); + supplier.addPoll(record("topic", 0, 1)); + final Executor rejectingExecutor = command -> { + throw new RejectedExecutionException("full"); + }; + final ShareGroupAcquisitionRegistry registry = registry(1, 1_000); + final ShareGroupAcquisitionLoop loop = loop(supplier, records -> {}, rejectingExecutor, registry); + + loop.runOnce(); + + Assertions.assertEquals(AcknowledgeType.RELEASE, supplier.acknowledgement(1).type); + Assertions.assertEquals(1, supplier.flushCount); + Assertions.assertEquals(0, registry.size()); + } + + @Test + public void testWorkerSuccessIsAcceptedOnOwnerThread() throws Exception + { + final ScriptedSupplier supplier = new ScriptedSupplier(Optional.of(60_000)); + supplier.addPoll(record("topic", 0, 1)); + final AtomicReference<Thread> stagingThread = new AtomicReference<>(); + final ExecutorService stagingExecutor = Executors.newSingleThreadExecutor(); + try { + final ShareGroupAcquisitionLoop loop = loop( + supplier, + records -> stagingThread.set(Thread.currentThread()), + stagingExecutor, + registry(1, 1_000) + ); + + loop.runOnce(); + + Assertions.assertEquals(AcknowledgeType.ACCEPT, supplier.acknowledgement(1).type); + Assertions.assertNotSame(stagingThread.get(), supplier.acknowledgement(1).thread); + Assertions.assertEquals(Set.of(Thread.currentThread()), supplier.consumerThreads); + } + finally { + stagingExecutor.shutdownNow(); + } + } + + @Test + public void testWorkerFailureIsReleasedOnOwnerThread() throws Exception + { + final ScriptedSupplier supplier = new ScriptedSupplier(Optional.of(60_000)); + supplier.addPoll(record("topic", 0, 1)); + final ExecutorService stagingExecutor = Executors.newSingleThreadExecutor(); + try { + final ShareGroupAcquisitionLoop loop = loop( + supplier, + records -> { + throw new IllegalStateException("upload failed"); + }, + stagingExecutor, + registry(1, 1_000) + ); + + loop.runOnce(); + + Assertions.assertEquals(AcknowledgeType.RELEASE, supplier.acknowledgement(1).type); + Assertions.assertEquals(Set.of(Thread.currentThread()), supplier.consumerThreads); + } + finally { + stagingExecutor.shutdownNow(); + } + } + + @Test + public void testStopDuringUploadReleasesAndFlushes() throws Exception + { + final ScriptedSupplier supplier = new ScriptedSupplier(Optional.of(60_000)); + supplier.addPoll(record("topic", 0, 1)); + final RecordingExecutor stagingExecutor = new RecordingExecutor(); + final ShareGroupAcquisitionLoop loop = loop( + supplier, + records -> {}, + stagingExecutor, + registry(1, 1_000) + ); + final ExecutorService ownerExecutor = Executors.newSingleThreadExecutor(); + try { + final Future<?> run = ownerExecutor.submit(() -> { + try { + loop.runOnce(); + } + catch (InterruptedException e) { + throw new RuntimeException(e); + } + }); + Assertions.assertTrue(supplier.pollReturned.await(10, TimeUnit.SECONDS)); + + loop.requestStop(); + run.get(10, TimeUnit.SECONDS); + + Assertions.assertEquals(AcknowledgeType.RELEASE, supplier.acknowledgement(1).type); + Assertions.assertEquals(1, supplier.flushCount); + Assertions.assertEquals(1, supplier.wakeupCount); + } + finally { + ownerExecutor.shutdownNow(); + } + } + + @Test + public void testFailedFlushRetainsTerminalEntry() throws Exception + { + final ScriptedSupplier supplier = new ScriptedSupplier(Optional.of(60_000)); + supplier.flushFailure = new IllegalStateException("flush failed"); + supplier.addPoll(record("topic", 0, 1)); + final ShareGroupAcquisitionRegistry registry = registry(1, 1_000); + final ShareGroupAcquisitionLoop loop = loop(supplier, records -> {}, Runnable::run, registry); + + loop.runOnce(); + + Assertions.assertEquals(1, registry.size()); + Assertions.assertEquals( + ShareGroupAcquisitionRegistry.State.ACCEPT_PENDING, + registry.get(new ShareGroupRecordIdentity("topic", 0, 1)).orElseThrow().getState() + ); + } + + private static ShareGroupAcquisitionLoop loop( + ScriptedSupplier supplier, + ShareGroupBatchStager stager, + Executor executor, + ShareGroupAcquisitionRegistry registry + ) + { + return new ShareGroupAcquisitionLoop( + "topic", + supplier, + stager, + executor, + registry, + ignored -> 1, + 1, + 60_000 + ); + } + + private static ShareGroupAcquisitionRegistry registry(int maxRecords, long maxBytes) + { + final AtomicLong nanoTime = new AtomicLong(); + return new ShareGroupAcquisitionRegistry(maxRecords, maxBytes, 0.5, nanoTime::get); + } + + @SafeVarargs + private static List<OrderedPartitionableRecord<KafkaTopicPartition, Long, KafkaRecordEntity>> records( + OrderedPartitionableRecord<KafkaTopicPartition, Long, KafkaRecordEntity>... records + ) + { + return List.of(records); + } + + private static OrderedPartitionableRecord<KafkaTopicPartition, Long, KafkaRecordEntity> record( + String topic, + int partition, + long offset + ) + { + final ConsumerRecord<byte[], byte[]> consumerRecord = new ConsumerRecord<>( + topic, + partition, + offset, + null, + new byte[]{1} + ); + return new OrderedPartitionableRecord<>( + topic, + new KafkaTopicPartition(true, topic, partition), + offset, + List.of(new KafkaRecordEntity(consumerRecord)) + ); + } + + private static final class RecordingExecutor implements Executor + { + private final Deque<Runnable> commands = new ArrayDeque<>(); + + @Override + public void execute(Runnable command) + { + commands.add(command); + } + + int size() + { + return commands.size(); + } + + void runNext() + { + commands.remove().run(); + } + } + + private static final class RecordedAcknowledgement + { + private final long offset; + private final AcknowledgeType type; + private final Thread thread; + + private RecordedAcknowledgement(long offset, AcknowledgeType type, Thread thread) + { + this.offset = offset; + this.type = type; + this.thread = thread; + } + } + + private static final class ScriptedSupplier + implements AcknowledgingRecordSupplier<KafkaTopicPartition, Long, KafkaRecordEntity> + { + private final Optional<Integer> lockTimeoutMs; + private final Deque<List<OrderedPartitionableRecord<KafkaTopicPartition, Long, KafkaRecordEntity>>> polls = + new ArrayDeque<>(); + private final List<RecordedAcknowledgement> acknowledgements = new ArrayList<>(); + private final Set<Thread> consumerThreads = new HashSet<>(); + private final CountDownLatch pollReturned = new CountDownLatch(1); + private Set<String> subscription = Set.of(); + private int flushCount; + private int wakeupCount; + private Exception flushFailure; + + private ScriptedSupplier(Optional<Integer> lockTimeoutMs) + { + this.lockTimeoutMs = lockTimeoutMs; + } + + @SafeVarargs + final void addPoll(OrderedPartitionableRecord<KafkaTopicPartition, Long, KafkaRecordEntity>... pollRecords) + { + polls.add(records(pollRecords)); + } + + RecordedAcknowledgement acknowledgement(long offset) + { + return acknowledgements.stream() + .filter(acknowledgement -> acknowledgement.offset == offset) + .findFirst() + .orElseThrow(); + } + + @Override + public void subscribe(Set<String> topics) + { + consumerThreads.add(Thread.currentThread()); + subscription = Set.copyOf(topics); + } + + @Override + public void unsubscribe() + { + consumerThreads.add(Thread.currentThread()); + subscription = Set.of(); + } + + @Override + public Set<String> subscription() + { + consumerThreads.add(Thread.currentThread()); + return subscription; + } + + @Override + public List<OrderedPartitionableRecord<KafkaTopicPartition, Long, KafkaRecordEntity>> poll(long timeoutMs) + { + consumerThreads.add(Thread.currentThread()); + final List<OrderedPartitionableRecord<KafkaTopicPartition, Long, KafkaRecordEntity>> records = + polls.isEmpty() ? List.of() : polls.remove(); + pollReturned.countDown(); + return records; + } + + @Override + public void acknowledge(KafkaTopicPartition partitionId, Long offset) + { + acknowledge(partitionId, offset, AcknowledgeType.ACCEPT); + } + + @Override + public void acknowledge(KafkaTopicPartition partitionId, Long offset, AcknowledgeType type) + { + consumerThreads.add(Thread.currentThread()); + acknowledgements.add(new RecordedAcknowledgement(offset, type, Thread.currentThread())); + } + + @Override + public void acknowledge( + Map<KafkaTopicPartition, Collection<Long>> offsets, + AcknowledgeType type + ) + { + for (Map.Entry<KafkaTopicPartition, Collection<Long>> entry : offsets.entrySet()) { + for (Long offset : entry.getValue()) { + acknowledge(entry.getKey(), offset, type); + } + } + } + + @Override + public Map<KafkaTopicPartition, Optional<Exception>> flushAcknowledgementsSync() + { + consumerThreads.add(Thread.currentThread()); + flushCount++; + final Map<KafkaTopicPartition, Optional<Exception>> result = new HashMap<>(); + for (RecordedAcknowledgement acknowledgement : acknowledgements) { Review Comment: ## CodeQL / Unread local variable Variable 'RecordedAcknowledgement acknowledgement' is never read. [Show more details](https://github.com/apache/druid/security/code-scanning/12023) ########## extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaShareGroupRecordSupplierTest.java: ########## @@ -0,0 +1,461 @@ +/* + * 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 org.apache.druid.data.input.kafka.KafkaRecordEntity; +import org.apache.druid.data.input.kafka.KafkaTopicPartition; +import org.apache.druid.indexing.seekablestream.common.AcknowledgeType; +import org.apache.druid.indexing.seekablestream.common.OrderedPartitionableRecord; +import org.apache.kafka.clients.consumer.ConsumerRecord; +import org.apache.kafka.clients.consumer.ConsumerRecords; +import org.apache.kafka.clients.consumer.KafkaShareConsumer; +import org.apache.kafka.common.KafkaException; +import org.apache.kafka.common.TopicIdPartition; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.Uuid; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.Mockito; + +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class KafkaShareGroupRecordSupplierTest +{ + private KafkaShareConsumer<byte[], byte[]> mockConsumer; + private KafkaShareGroupRecordSupplier supplier; + + @SuppressWarnings("unchecked") + @BeforeEach + public void setUp() + { + mockConsumer = mock(KafkaShareConsumer.class); + supplier = new KafkaShareGroupRecordSupplier(mockConsumer); + } + + @AfterEach + public void tearDown() + { + supplier.close(); + } + + @Test + public void testSubscribeAndSubscription() + { + final Set<String> topics = Set.of("topic-a", "topic-b"); + when(mockConsumer.subscription()).thenReturn(topics); + + supplier.subscribe(topics); + verify(mockConsumer).subscribe(topics); + + Assertions.assertEquals(topics, supplier.subscription()); + } + + @Test + public void testUnsubscribe() + { + supplier.unsubscribe(); + verify(mockConsumer).unsubscribe(); + } + + @Test + public void testPollWrapsRecords() + { + final ConsumerRecord<byte[], byte[]> record1 = new ConsumerRecord<>( + "test-topic", 0, 100L, "key1".getBytes(StandardCharsets.UTF_8), "value1".getBytes(StandardCharsets.UTF_8) + ); + final ConsumerRecord<byte[], byte[]> record2 = new ConsumerRecord<>( + "test-topic", 1, 200L, "key2".getBytes(StandardCharsets.UTF_8), "value2".getBytes(StandardCharsets.UTF_8) + ); + + final Map<TopicPartition, List<ConsumerRecord<byte[], byte[]>>> recordMap = new HashMap<>(); + recordMap.put(new TopicPartition("test-topic", 0), List.of(record1)); + recordMap.put(new TopicPartition("test-topic", 1), List.of(record2)); + final ConsumerRecords<byte[], byte[]> consumerRecords = new ConsumerRecords<>(recordMap); + + when(mockConsumer.poll(any(Duration.class))).thenReturn(consumerRecords); + + final List<OrderedPartitionableRecord<KafkaTopicPartition, Long, KafkaRecordEntity>> result = + supplier.poll(1000); + + Assertions.assertEquals(2, result.size()); + + // Verify first record + final OrderedPartitionableRecord<KafkaTopicPartition, Long, KafkaRecordEntity> polled1 = + result.stream().filter(r -> r.getSequenceNumber() == 100L).findFirst().orElse(null); + Assertions.assertNotNull(polled1); + Assertions.assertEquals("test-topic", polled1.getStream()); + Assertions.assertEquals(0, polled1.getPartitionId().partition()); + Assertions.assertNotNull(polled1.getData()); + Assertions.assertEquals(1, polled1.getData().size()); + + // Verify second record + final OrderedPartitionableRecord<KafkaTopicPartition, Long, KafkaRecordEntity> polled2 = + result.stream().filter(r -> r.getSequenceNumber() == 200L).findFirst().orElse(null); + Assertions.assertNotNull(polled2); + Assertions.assertEquals(1, polled2.getPartitionId().partition()); + } + + @Test + public void testPollWithNullValue() + { + final ConsumerRecord<byte[], byte[]> record = new ConsumerRecord<>( + "test-topic", 0, 50L, "key".getBytes(StandardCharsets.UTF_8), null + ); + final Map<TopicPartition, List<ConsumerRecord<byte[], byte[]>>> recordMap = new HashMap<>(); + recordMap.put(new TopicPartition("test-topic", 0), List.of(record)); + when(mockConsumer.poll(any(Duration.class))).thenReturn(new ConsumerRecords<>(recordMap)); + + final List<OrderedPartitionableRecord<KafkaTopicPartition, Long, KafkaRecordEntity>> result = + supplier.poll(1000); + + Assertions.assertEquals(1, result.size()); + Assertions.assertTrue(result.get(0).getData().isEmpty()); + Assertions.assertSame(record, ((KafkaShareGroupRecord) result.get(0)).getConsumerRecord()); + } + + @Test + public void testPollReturnsEmptyOnTimeout() + { + when(mockConsumer.poll(any(Duration.class))).thenReturn(ConsumerRecords.empty()); + final List<OrderedPartitionableRecord<KafkaTopicPartition, Long, KafkaRecordEntity>> result = + supplier.poll(100); + Assertions.assertTrue(result.isEmpty()); + } + + @Test + public void testAcknowledgeDefaultAccept() + { + final ConsumerRecord<byte[], byte[]> record = new ConsumerRecord<>( + "test-topic", 0, 42L, "key".getBytes(StandardCharsets.UTF_8), "value".getBytes(StandardCharsets.UTF_8) + ); + pollSingleRecord(record); + + final KafkaTopicPartition partition = new KafkaTopicPartition(true, "test-topic", 0); + supplier.acknowledge(partition, 42L); + + Mockito.verify(mockConsumer).acknowledge( + Mockito.same(record), + Mockito.eq(org.apache.kafka.clients.consumer.AcknowledgeType.ACCEPT) + ); + } + + @Test + public void testAcknowledgeWithRelease() + { + final ConsumerRecord<byte[], byte[]> record = new ConsumerRecord<>( + "test-topic", 0, 10L, "key".getBytes(StandardCharsets.UTF_8), "value".getBytes(StandardCharsets.UTF_8) + ); + pollSingleRecord(record); + + final KafkaTopicPartition partition = new KafkaTopicPartition(true, "test-topic", 0); + supplier.acknowledge(partition, 10L, AcknowledgeType.RELEASE); + + Mockito.verify(mockConsumer).acknowledge( + Mockito.same(record), + Mockito.eq(org.apache.kafka.clients.consumer.AcknowledgeType.RELEASE) + ); + } + + @Test + public void testAcknowledgeWithReject() + { + final ConsumerRecord<byte[], byte[]> record = new ConsumerRecord<>( + "test-topic", 0, 10L, "key".getBytes(StandardCharsets.UTF_8), "value".getBytes(StandardCharsets.UTF_8) + ); + pollSingleRecord(record); + + final KafkaTopicPartition partition = new KafkaTopicPartition(true, "test-topic", 0); + supplier.acknowledge(partition, 10L, AcknowledgeType.REJECT); + + Mockito.verify(mockConsumer).acknowledge( + Mockito.same(record), + Mockito.eq(org.apache.kafka.clients.consumer.AcknowledgeType.REJECT) + ); + } + + @Test + public void testAcknowledgeWithRenew() + { + final ConsumerRecord<byte[], byte[]> record = new ConsumerRecord<>( + "test-topic", 0, 99L, "key".getBytes(StandardCharsets.UTF_8), "value".getBytes(StandardCharsets.UTF_8) + ); + pollSingleRecord(record); + + final KafkaTopicPartition partition = new KafkaTopicPartition(true, "test-topic", 0); + supplier.acknowledge(partition, 99L, AcknowledgeType.RENEW); + + Mockito.verify(mockConsumer).acknowledge( + Mockito.same(record), + Mockito.eq(org.apache.kafka.clients.consumer.AcknowledgeType.RENEW) + ); + } + + @Test + public void testAcknowledgeBatch() + { + final List<ConsumerRecord<byte[], byte[]>> recordsP0 = Arrays.asList( + new ConsumerRecord<>("test-topic", 0, 1L, "k".getBytes(StandardCharsets.UTF_8), "v".getBytes(StandardCharsets.UTF_8)), + new ConsumerRecord<>("test-topic", 0, 2L, "k".getBytes(StandardCharsets.UTF_8), "v".getBytes(StandardCharsets.UTF_8)), + new ConsumerRecord<>("test-topic", 0, 3L, "k".getBytes(StandardCharsets.UTF_8), "v".getBytes(StandardCharsets.UTF_8)) + ); + final List<ConsumerRecord<byte[], byte[]>> recordsP1 = Arrays.asList( + new ConsumerRecord<>("test-topic", 1, 10L, "k".getBytes(StandardCharsets.UTF_8), "v".getBytes(StandardCharsets.UTF_8)), + new ConsumerRecord<>("test-topic", 1, 11L, "k".getBytes(StandardCharsets.UTF_8), "v".getBytes(StandardCharsets.UTF_8)) + ); + final Map<TopicPartition, List<ConsumerRecord<byte[], byte[]>>> recordMap = new HashMap<>(); + recordMap.put(new TopicPartition("test-topic", 0), recordsP0); + recordMap.put(new TopicPartition("test-topic", 1), recordsP1); + Mockito.when(mockConsumer.poll(Mockito.any(Duration.class))) + .thenReturn(new ConsumerRecords<>(recordMap)); + supplier.poll(1000); + + final KafkaTopicPartition p0 = new KafkaTopicPartition(true, "test-topic", 0); + final KafkaTopicPartition p1 = new KafkaTopicPartition(true, "test-topic", 1); + + final Map<KafkaTopicPartition, java.util.Collection<Long>> offsets = new HashMap<>(); + offsets.put(p0, Arrays.asList(1L, 2L, 3L)); + offsets.put(p1, Arrays.asList(10L, 11L)); + + supplier.acknowledge(offsets, AcknowledgeType.ACCEPT); + + Mockito.verify(mockConsumer, Mockito.times(5)).acknowledge( + Mockito.<ConsumerRecord<byte[], byte[]>>any(), + Mockito.eq(org.apache.kafka.clients.consumer.AcknowledgeType.ACCEPT) + ); + } + + private void pollSingleRecord(ConsumerRecord<byte[], byte[]> record) + { + final Map<TopicPartition, List<ConsumerRecord<byte[], byte[]>>> recordMap = new HashMap<>(); + recordMap.put(new TopicPartition(record.topic(), record.partition()), List.of(record)); + Mockito.when(mockConsumer.poll(Mockito.any(Duration.class))) + .thenReturn(new ConsumerRecords<>(recordMap)); Review Comment: ## CodeQL / Deprecated method or constructor invocation Invoking [ConsumerRecords.ConsumerRecords](1) should be avoided because it has been deprecated. [Show more details](https://github.com/apache/druid/security/code-scanning/12029) ########## indexing-service/src/main/java/org/apache/druid/indexing/common/actions/RenewShareInboxClaimsAction.java: ########## @@ -0,0 +1,69 @@ +/* + * 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.indexing.common.task.Task; +import org.apache.druid.indexing.overlord.ShareInboxRenewRequest; +import org.apache.druid.indexing.overlord.ShareInboxRenewResult; + +import java.util.Objects; + +public class RenewShareInboxClaimsAction implements TaskAction<ShareInboxRenewResult> +{ + public static final String TYPE = "renewShareInboxClaims"; + + private final ShareInboxRenewRequest request; + + @JsonCreator + public RenewShareInboxClaimsAction(@JsonProperty("request") ShareInboxRenewRequest request) + { + this.request = Objects.requireNonNull(request, "request"); + } + + @JsonProperty + public ShareInboxRenewRequest getRequest() + { + return request; + } + + @Override + public TypeReference<ShareInboxRenewResult> getReturnTypeReference() + { + return new TypeReference<>() {}; + } + + @Override + public ShareInboxRenewResult perform(Task task, TaskActionToolbox toolbox) + { + ClaimShareInboxManifestsAction.validateTask(task, request.getDataSource(), request.getClaimOwner()); + return toolbox.getIndexerMetadataStorageCoordinator().renewShareInboxClaims(request); + } + + @Override + public String toString() + { + return "RenewShareInboxClaimsAction{" + + "request=" + request + Review Comment: ## CodeQL / Use of default toString() Default toString(): ShareInboxRenewRequest inherits toString() from Object, and so is not suitable for printing. [Show more details](https://github.com/apache/druid/security/code-scanning/12026) -- 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]
