github-advanced-security[bot] commented on code in PR #20197: URL: https://github.com/apache/druid/pull/20197#discussion_r4088856216
########## server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerAcquireLifecycleTest.java: ########## @@ -0,0 +1,504 @@ +/* + * 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 int SEGMENT_SIZE = 1000; + + /** + * 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(SegmentId.of( + "test_ds", + Intervals.of("2024-01-01/2024-01-02"), + "v1", + new NumberedShardSpec( + Integer.parseInt(name.substring(name.length() - 1)), Review Comment: ## CodeQL / Missing catch of NumberFormatException Potential uncaught 'java.lang.NumberFormatException'. [Show more details](https://github.com/apache/druid/security/code-scanning/12005) -- 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]
