github-advanced-security[bot] commented on code in PR #20197: URL: https://github.com/apache/druid/pull/20197#discussion_r3884939006
########## server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerAcquireLifecycleTest.java: ########## @@ -0,0 +1,500 @@ +/* + * 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.segment.loading; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.annotation.JsonTypeName; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.util.concurrent.MoreExecutors; +import com.google.common.util.concurrent.Uninterruptibles; +import org.apache.druid.error.DruidException; +import org.apache.druid.jackson.DefaultObjectMapper; +import org.apache.druid.java.util.common.Intervals; +import org.apache.druid.java.util.emitter.EmittingLogger; +import org.apache.druid.segment.Segment; +import org.apache.druid.segment.TestIndex; +import org.apache.druid.segment.TestSegmentUtils; +import org.apache.druid.server.metrics.NoopServiceEmitter; +import org.apache.druid.timeline.DataSegment; +import org.apache.druid.timeline.SegmentId; +import org.apache.druid.timeline.partition.NumberedShardSpec; +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.junit.jupiter.api.io.TempDir; + +import java.io.Closeable; +import java.io.File; +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.AbstractExecutorService; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Semaphore; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; + +/** + * Lifecycle-focused tests for {@link SegmentLocalCacheManager#acquireSegment}'s hold-handoff choreography: the + * pre-placed reservation hold must be released exactly once on every path (cancel-while-queued, cancel-mid-load, + * delivery losing the race with close, empty delivery, normal release). + */ +class SegmentLocalCacheManagerAcquireLifecycleTest +{ + private static final long SEGMENT_SIZE = 1000L; + + /** + * Per-segment-name gate controlling {@link GatedLoadSpec#loadSegment}. Static because Jackson materializes fresh + * {@link GatedLoadSpec} instances from the load spec map on every acquire. + */ + static class Gate + { + final CountDownLatch entered = new CountDownLatch(1); + final CountDownLatch proceed = new CountDownLatch(1); + final CountDownLatch exited = new CountDownLatch(1); + final AtomicInteger loadCount = new AtomicInteger(); + final AtomicBoolean sawInterrupt = new AtomicBoolean(); + volatile boolean blockUninterruptibly = false; + } + + private static final Map<String, Gate> GATES = new ConcurrentHashMap<>(); + + private static Gate gate(String name) + { + return GATES.computeIfAbsent(name, ignored -> new Gate()); + } + + @JsonTypeName("gated") + public static class GatedLoadSpec implements LoadSpec + { + private final int size; + private final String name; + + @JsonCreator + public GatedLoadSpec(@JsonProperty("size") int size, @JsonProperty("name") String name) + { + this.size = size; + this.name = name; + } + + @Override + public LoadSpecResult loadSegment(File destDir) throws SegmentLoadingException + { + final Gate gate = gate(name); + gate.loadCount.incrementAndGet(); + gate.entered.countDown(); + try { + if (gate.blockUninterruptibly) { + Uninterruptibles.awaitUninterruptibly(gate.proceed); + } else { + try { + gate.proceed.await(); + } + catch (InterruptedException e) { + gate.sawInterrupt.set(true); + throw new SegmentLoadingException(e, "interrupted while loading[%s]", name); + } + } + return new TestSegmentUtils.TestLoadSpec(size, name).loadSegment(destDir); + } + finally { + gate.exited.countDown(); + } + } + } + + /** + * Executor that holds every submitted task instead of running it, until the test calls {@link #dispatchAll()}. + * Lets tests deterministically construct the submitted-but-not-yet-started state that a real pool (which starts + * tasks immediately, especially in the virtual-thread mode) cannot guarantee. + */ + private static class DeferredDispatchExecutorService extends AbstractExecutorService + { + private final List<Runnable> held = new ArrayList<>(); + + @Override + public synchronized void execute(Runnable command) + { + held.add(command); + } + + synchronized void dispatchAll() + { + // run inline: a cancelled FutureTask's run() is a no-op, so this is safe on the test thread + for (Runnable runnable : held) { + runnable.run(); + } + held.clear(); + } + + @Override + public void shutdown() + { + } + + @Override + public List<Runnable> shutdownNow() + { + synchronized (this) { + final List<Runnable> remaining = new ArrayList<>(held); + held.clear(); + return remaining; + } + } + + @Override + public boolean isShutdown() + { + return false; + } + + @Override + public boolean isTerminated() + { + return false; + } + + @Override + public boolean awaitTermination(long timeout, TimeUnit unit) + { + return true; + } + } + + @TempDir + File tempDir; + + private ObjectMapper jsonMapper; + private SegmentLocalCacheManager manager; + private StorageLocation location; + + @BeforeEach + void setUp() throws IOException + { + EmittingLogger.registerEmitter(new NoopServiceEmitter()); + GATES.clear(); + jsonMapper = new DefaultObjectMapper(); + jsonMapper.registerSubtypes(GatedLoadSpec.class); + jsonMapper.registerSubtypes(TestSegmentUtils.TestLoadSpec.class); + jsonMapper.registerSubtypes(TestSegmentUtils.TestSegmentizerFactory.class); + + final StorageLocationConfig locationConfig = new StorageLocationConfig( + new File(tempDir, "cache"), + 100_000L, + null + ); + final SegmentLoaderConfig loaderConfig = SegmentLoaderConfig.builder() + .locations(locationConfig) + .virtualStorage(true) + .virtualStorageLoadThreads(1) + .infoDir(new File(tempDir, "info")) + .build(); + final List<StorageLocation> storageLocations = loaderConfig.toStorageLocations(); + location = storageLocations.get(0); + manager = new SegmentLocalCacheManager( + storageLocations, + loaderConfig, + StorageLoadingThreadPool.createFromConfig(loaderConfig), + new LeastBytesUsedStorageLocationSelectorStrategy(storageLocations), + TestIndex.INDEX_IO, + jsonMapper + ); + manager.getCachedSegments(); + } + + @AfterEach + void tearDown() + { + // open all gates so any still-blocked load task can unwind before the executor is torn down + for (Gate gate : GATES.values()) { + gate.proceed.countDown(); + } + manager.shutdown(); + } + + private DataSegment makeSegment(String name) + { + return DataSegment.builder() Review Comment: ## CodeQL / Deprecated method or constructor invocation Invoking [DataSegment.builder](1) should be avoided because it has been deprecated. [Show more details](https://github.com/apache/druid/security/code-scanning/11932) ########## multi-stage-query/src/main/java/org/apache/druid/msq/querykit/ReadableInputQueue.java: ########## @@ -356,36 +427,69 @@ @Override public void close() { + final List<AcquireSegmentAction> handlesToClose; + final List<SegmentReferenceHolder> holdersToDrain; + final List<SettableFuture<ReadableInput>> futuresToFail = new ArrayList<>(); + + // Snapshot and clear everything under the monitor, but do the actual closing OUTSIDE it: closing a handle can + // synchronously run its canceler (a deferred acquire's canceler takes its own stage lock and may complete the + // handle), and holding the queue monitor across that work would stall every concurrent delivery and nextInput() + // call behind it. Clearing the loading map first means a late-arriving onSegmentReady finds no entry + // (future == null) and returns without touching queue state. synchronized (this) { + closed = true; readablePartitions.clear(); queryableServers.clear(); loadableSegments.clear(); - // Cancel all pending segment loads. - for (AcquireSegmentAction acquireSegmentAction : loadingSegments) { - CloseableUtils.closeAndSuppressExceptions( - acquireSegmentAction, + handlesToClose = new ArrayList<>(loadingSegments.keySet()); + futuresToFail.addAll(loadingSegments.values()); + loadingSegments.clear(); + + // Also fail deliveries that are mid-completion (removed from loadingSegments but not yet completed); + // SettableFuture's first-write-wins semantics make whichever write loses a harmless no-op. + futuresToFail.addAll(inFlightDeliveries); + + holdersToDrain = new ArrayList<>(loadedSegments); + loadedSegments.clear(); - // AcquireSegmentAction currently doesn't have a meaningful toString method, so if this message - // ever actually gets logged, it won't mention the specific segment that had a problem. Perhaps - // one day this will change. - e -> log.warn(e, "Failed to close loadingSegment[%s]", acquireSegmentAction) + // Drop loadahead futures that were never handed out: their loads are covered by the closing/failing here, and + // handing them out after close would deliver holders this close() is draining. Also keeps remaining() + // reporting 0 after close. + pendingNextInputs.clear(); + } + + // Cancel all pending segment loads: closing a NEW handle runs its canceler (aborting the load); closing a + // READY-but-unclaimed handle (its ready callback hasn't reached the queue monitor yet) closes the delivered + // result. + for (final AcquireSegmentAction acquireSegmentAction : handlesToClose) { + CloseableUtils.closeAndSuppressExceptions( + acquireSegmentAction, + + // AcquireSegmentAction currently doesn't have a meaningful toString method, so if this message + // ever actually gets logged, it won't mention the specific segment that had a problem. Perhaps + // one day this will change. + e -> log.warn(e, "Failed to close loadingSegment[%s]", acquireSegmentAction) Review Comment: ## CodeQL / Use of default toString() Default toString(): AcquireSegmentAction inherits toString() from Object, and so is not suitable for printing. [Show more details](https://github.com/apache/druid/security/code-scanning/11931) ########## server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerAcquireLifecycleTest.java: ########## @@ -0,0 +1,500 @@ +/* + * 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.segment.loading; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.annotation.JsonTypeName; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.util.concurrent.MoreExecutors; +import com.google.common.util.concurrent.Uninterruptibles; +import org.apache.druid.error.DruidException; +import org.apache.druid.jackson.DefaultObjectMapper; +import org.apache.druid.java.util.common.Intervals; +import org.apache.druid.java.util.emitter.EmittingLogger; +import org.apache.druid.segment.Segment; +import org.apache.druid.segment.TestIndex; +import org.apache.druid.segment.TestSegmentUtils; +import org.apache.druid.server.metrics.NoopServiceEmitter; +import org.apache.druid.timeline.DataSegment; +import org.apache.druid.timeline.SegmentId; +import org.apache.druid.timeline.partition.NumberedShardSpec; +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.junit.jupiter.api.io.TempDir; + +import java.io.Closeable; +import java.io.File; +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.AbstractExecutorService; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Semaphore; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; + +/** + * Lifecycle-focused tests for {@link SegmentLocalCacheManager#acquireSegment}'s hold-handoff choreography: the + * pre-placed reservation hold must be released exactly once on every path (cancel-while-queued, cancel-mid-load, + * delivery losing the race with close, empty delivery, normal release). + */ +class SegmentLocalCacheManagerAcquireLifecycleTest +{ + private static final long SEGMENT_SIZE = 1000L; + + /** + * Per-segment-name gate controlling {@link GatedLoadSpec#loadSegment}. Static because Jackson materializes fresh + * {@link GatedLoadSpec} instances from the load spec map on every acquire. + */ + static class Gate + { + final CountDownLatch entered = new CountDownLatch(1); + final CountDownLatch proceed = new CountDownLatch(1); + final CountDownLatch exited = new CountDownLatch(1); + final AtomicInteger loadCount = new AtomicInteger(); + final AtomicBoolean sawInterrupt = new AtomicBoolean(); + volatile boolean blockUninterruptibly = false; + } + + private static final Map<String, Gate> GATES = new ConcurrentHashMap<>(); + + private static Gate gate(String name) + { + return GATES.computeIfAbsent(name, ignored -> new Gate()); + } + + @JsonTypeName("gated") + public static class GatedLoadSpec implements LoadSpec + { + private final int size; + private final String name; + + @JsonCreator + public GatedLoadSpec(@JsonProperty("size") int size, @JsonProperty("name") String name) + { + this.size = size; + this.name = name; + } + + @Override + public LoadSpecResult loadSegment(File destDir) throws SegmentLoadingException + { + final Gate gate = gate(name); + gate.loadCount.incrementAndGet(); + gate.entered.countDown(); + try { + if (gate.blockUninterruptibly) { + Uninterruptibles.awaitUninterruptibly(gate.proceed); + } else { + try { + gate.proceed.await(); + } + catch (InterruptedException e) { + gate.sawInterrupt.set(true); + throw new SegmentLoadingException(e, "interrupted while loading[%s]", name); + } + } + return new TestSegmentUtils.TestLoadSpec(size, name).loadSegment(destDir); + } + finally { + gate.exited.countDown(); + } + } + } + + /** + * Executor that holds every submitted task instead of running it, until the test calls {@link #dispatchAll()}. + * Lets tests deterministically construct the submitted-but-not-yet-started state that a real pool (which starts + * tasks immediately, especially in the virtual-thread mode) cannot guarantee. + */ + private static class DeferredDispatchExecutorService extends AbstractExecutorService + { + private final List<Runnable> held = new ArrayList<>(); + + @Override + public synchronized void execute(Runnable command) + { + held.add(command); + } + + synchronized void dispatchAll() + { + // run inline: a cancelled FutureTask's run() is a no-op, so this is safe on the test thread + for (Runnable runnable : held) { + runnable.run(); + } + held.clear(); + } + + @Override + public void shutdown() + { + } + + @Override + public List<Runnable> shutdownNow() + { + synchronized (this) { + final List<Runnable> remaining = new ArrayList<>(held); + held.clear(); + return remaining; + } + } + + @Override + public boolean isShutdown() + { + return false; + } + + @Override + public boolean isTerminated() + { + return false; + } + + @Override + public boolean awaitTermination(long timeout, TimeUnit unit) + { + return true; + } + } + + @TempDir + File tempDir; + + private ObjectMapper jsonMapper; + private SegmentLocalCacheManager manager; + private StorageLocation location; + + @BeforeEach + void setUp() throws IOException + { + EmittingLogger.registerEmitter(new NoopServiceEmitter()); + GATES.clear(); + jsonMapper = new DefaultObjectMapper(); + jsonMapper.registerSubtypes(GatedLoadSpec.class); + jsonMapper.registerSubtypes(TestSegmentUtils.TestLoadSpec.class); + jsonMapper.registerSubtypes(TestSegmentUtils.TestSegmentizerFactory.class); + + final StorageLocationConfig locationConfig = new StorageLocationConfig( + new File(tempDir, "cache"), + 100_000L, + null + ); + final SegmentLoaderConfig loaderConfig = SegmentLoaderConfig.builder() + .locations(locationConfig) + .virtualStorage(true) + .virtualStorageLoadThreads(1) + .infoDir(new File(tempDir, "info")) + .build(); + final List<StorageLocation> storageLocations = loaderConfig.toStorageLocations(); + location = storageLocations.get(0); + manager = new SegmentLocalCacheManager( + storageLocations, + loaderConfig, + StorageLoadingThreadPool.createFromConfig(loaderConfig), + new LeastBytesUsedStorageLocationSelectorStrategy(storageLocations), + TestIndex.INDEX_IO, + jsonMapper + ); + manager.getCachedSegments(); + } + + @AfterEach + void tearDown() + { + // open all gates so any still-blocked load task can unwind before the executor is torn down + for (Gate gate : GATES.values()) { + gate.proceed.countDown(); + } + manager.shutdown(); + } + + private DataSegment makeSegment(String name) + { + return DataSegment.builder() + .dataSource("test_ds") + .interval(Intervals.of("2024-01-01/2024-01-02")) + .version("v1") + .loadSpec(ImmutableMap.of("type", "gated", "size", (int) SEGMENT_SIZE, "name", name)) + .dimensions(ImmutableList.of()) + .metrics(ImmutableList.of()) + .shardSpec(new NumberedShardSpec(Integer.parseInt(name.substring(name.length() - 1)), 0)) Review Comment: ## CodeQL / Missing catch of NumberFormatException Potential uncaught 'java.lang.NumberFormatException'. [Show more details](https://github.com/apache/druid/security/code-scanning/11930) -- 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]
