FrankChen021 commented on code in PR #20197: URL: https://github.com/apache/druid/pull/20197#discussion_r4082573403
########## multi-stage-query/src/main/java/org/apache/druid/msq/input/AcquireSegmentHandles.java: ########## @@ -0,0 +1,161 @@ +/* + * 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.msq.input; + +import org.apache.druid.common.asyncresource.AsyncResource; +import org.apache.druid.error.DruidException; +import org.apache.druid.java.util.common.logger.Logger; +import org.apache.druid.segment.ReferenceCountedSegmentProvider; +import org.apache.druid.segment.Segment; +import org.apache.druid.segment.loading.AcquireSegmentAction; +import org.apache.druid.segment.loading.AcquireSegmentResult; +import org.apache.druid.utils.CloseableUtils; + +import java.io.Closeable; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.Supplier; + +/** + * Shared machinery for bridging {@link AsyncResource}-producing segment sources into the releasable + * {@link AcquireSegmentAction} consumer handle. + */ +final class AcquireSegmentHandles +{ + private static final Logger log = new Logger(AcquireSegmentHandles.class); + + private AcquireSegmentHandles() + { + // No instantiation. + } + + /** + * Wraps {@code target} with a CAS guard so it is closed at most once. {@link AcquireSegmentAction#close()} is not + * idempotent, and a producer's error path can race its canceler. + */ + static Runnable closeOnce(Closeable target) + { + final AtomicBoolean closed = new AtomicBoolean(false); + return () -> { + if (closed.compareAndSet(false, true)) { + CloseableUtils.closeAndSuppressExceptions( + target, + e -> log.warn(e, "Failed to close acquire resource of class[%s]", target.getClass().getName()) + ); + } + }; + } + + /** + * Bridges a NON-releasable {@code inner} resource (e.g. one produced by {@code AsyncResources.collect/transform/ + * recover} combinators) into a releasable {@link AcquireSegmentAction}. The lifecycle of {@code inner} is folded + * into the delivered result: when the delivered result carries a segment, that segment's close also closes + * {@code inner} (releasing whatever the combinator chain owns, e.g. cached source files, which therefore strictly + * outlive the segment); when it is empty, {@code inner} is closed at delivery. Closing the returned action before + * readiness closes {@code inner} (cancelling in-flight work). + * <p> + * Uses {@code inner.get()}, not {@code release()}: combinator-produced resources are not releasable; inner keeps + * nominal ownership and the fold transfers the close responsibility onto the delivered segment. + */ + static AcquireSegmentAction fromResource(AsyncResource<AcquireSegmentResult> inner) + { + final Runnable closeInnerOnce = closeOnce(inner); + final AcquireSegmentAction outer = new AcquireSegmentAction(closeInnerOnce); + deliverOnReady(inner, outer, closeInnerOnce, () -> foldInnerClose(inner.get(), closeInnerOnce)); + return outer; + } + + /** + * Wires a releasable {@code inner} handle into {@code outer}: on readiness, ownership of the result transfers + * inner to outer. + * + * @param inner the source handle + * @param outer the handle handed to the consumer + * @param closeInnerOnce close-once guard for {@code inner}, shared with the canceler that may race this transfer + */ + static void transferOnReady(AcquireSegmentAction inner, AcquireSegmentAction outer, Runnable closeInnerOnce) + { + deliverOnReady(inner, outer, closeInnerOnce, inner::release); + } + + /** + * Shared delivery skeleton for {@link #fromResource} and {@link #transferOnReady}. On {@code inner}'s readiness, + * obtain the result and deliver it to {@code outer}; on failure report to {@code outer} (silently absorbed if outer + * was closed first) and release {@code inner} exactly once; if delivery loses the race with {@code outer}'s close, + * close the orphaned result. + */ + private static void deliverOnReady( + AsyncResource<AcquireSegmentResult> inner, + AcquireSegmentAction outer, + Runnable closeInnerOnce, + Supplier<AcquireSegmentResult> obtainResult + ) + { + try { + inner.addReadyCallback(() -> { + final AcquireSegmentResult delivered; + try { + delivered = obtainResult.get(); + } + catch (Throwable t) { + outer.setException(t); + closeInnerOnce.run(); + return; + } + if (!outer.set(delivered)) { + // the action was closed while delivering; we own the orphaned result (closing it releases inner too) + CloseableUtils.closeAndSuppressExceptions( + delivered, + e -> log.warn(e, "Failed to close orphaned acquire result after losing the delivery race") + ); + } + }); + } + catch (DruidException e) { Review Comment: [P2] Do not swallow ready-callback failures **Finding:** SettableAsyncResource.addReadyCallback invokes the callback inline when inner is already READY. In that case foldInnerClose can throw its own defensive DruidException at lines 147-153 (for a non-LeafReference result), but this outer catch treats it as the expected concurrent-close race and leaves outer in NEW without closing inner. The consumer then waits forever and the underlying resource can leak instead of receiving a failed action. **Suggestion:** Limit this catch to the expected closed-before-registration case, and route exceptions raised while obtaining or folding an already-ready result through outer.setException plus closeInnerOnce. Add a regression test where a ready inner resource contains a non-Leaf segment and verify the returned action becomes failed rather than hanging. -- 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]
