github-advanced-security[bot] commented on code in PR #20805:
URL: https://github.com/apache/druid/pull/20805#discussion_r4448052930
##########
multi-stage-query/src/main/java/org/apache/druid/msq/querykit/ReadableInputQueue.java:
##########
@@ -471,20 +597,137 @@
}
/**
- * Calls {@link #nextSegmentInput()} and adds the future to {@link
#pendingNextInputs}. Returns whether a future
- * was added.
+ * Acquires the next segment that {@link #inputPruner} does not rule out and
adds its future to
+ * {@link #pendingNextInputs}. Returns whether a future was added.
*/
private boolean addLoadaheadFuture()
{
- final ListenableFuture<NumberedReadableInput> nextFuture =
loadNextSegment();
- if (nextFuture != null) {
- synchronized (this) {
- pendingNextInputs.add(nextFuture);
- }
- return true;
- } else {
+ if (canSkipRemaining()) {
+ return false;
+ }
+
+ final NumberedItem<LoadableSegment> nextItem = pollUnprunedSegment();
+ if (nextItem == null) {
return false;
}
+
+ final ListenableFuture<NumberedReadableInput> nextFuture =
acquireSegment(nextItem);
+ synchronized (this) {
+ pendingNextInputs.put(nextFuture, nextItem.item());
+ }
+ return true;
+ }
+
+ /**
+ * Releases pending inputs that will never be handed out for processing.
Loads still in progress are canceled, and
+ * loaded inputs are released. An input that is mid-delivery (claimed from
{@link #loadingSegments} by
+ * {@link #onSegmentReady} but not yet completed) is left for {@link
#close()}; this is a conscious decision to
+ * avoid complexity for what is a sliver of a race.
+ * <p>
+ * Called outside the queue monitor, because canceling a load can run its
canceler and releasing a segment can take
+ * cache locks.
+ */
+ private void releaseSkippedInputs(final
List<ListenableFuture<NumberedReadableInput>> skippedFutures)
+ {
+ if (skippedFutures.isEmpty()) {
+ return;
+ }
+
+ final Set<ListenableFuture<NumberedReadableInput>> pendingSkips =
Sets.newIdentityHashSet();
+ pendingSkips.addAll(skippedFutures);
+ final Map<AcquireSegmentAction, SettableFuture<NumberedReadableInput>>
loadsToCancel = new LinkedHashMap<>();
+
+ synchronized (this) {
+ final Iterator<Map.Entry<AcquireSegmentAction,
SettableFuture<NumberedReadableInput>>> iterator =
+ loadingSegments.entrySet().iterator();
+ while (iterator.hasNext()) {
+ final Map.Entry<AcquireSegmentAction,
SettableFuture<NumberedReadableInput>> entry = iterator.next();
+ if (pendingSkips.remove(entry.getValue())) {
+ loadsToCancel.put(entry.getKey(), entry.getValue());
+ iterator.remove();
+ }
+ }
+ }
+
+ // Cancel all pending segment loads in segments that have been skipped.
+ for (final Map.Entry<AcquireSegmentAction,
SettableFuture<NumberedReadableInput>> entry : loadsToCancel.entrySet()) {
+ closeLoadingSegment(entry.getKey());
+ entry.getValue().setException(DruidException.defensive("Segment load
canceled because its input was skipped"));
+ }
+
+ // Lastly, release completed loads that were skipped if their future is
done.
+ for (final ListenableFuture<NumberedReadableInput> future : pendingSkips) {
+ if (future.isDone()) {
+ final NumberedReadableInput input;
+ try {
+ input = Futures.getDone(future);
+ }
+ catch (ExecutionException | CancellationException e) {
+ // A failed load has already closed its handle in onSegmentReady, so
there is nothing to release.
+ continue;
+ }
+ releaseLoadedSegment(input.input().getSegment());
+ }
+ }
+ }
+
+ /**
+ * Removes {@code holder} from {@link #loadedSegments} and closes its
segment, unless {@link #close()} claimed it
+ * first.
+ */
+ private void releaseLoadedSegment(final SegmentReferenceHolder holder)
+ {
+ synchronized (this) {
+ loadedSegments.remove(holder);
+ }
+ closeLoadedSegment(holder);
+ }
+
+ /**
+ * Closes a load that has not been delivered: closing a NEW handle runs its
canceler (aborting the load), and closing
+ * a ready but unclaimed handle (its ready callback hasn't reached the queue
monitor yet) closes the delivered result.
+ * The caller must have claimed the handle's entry from {@link
#loadingSegments}, and is responsible for failing its
+ * future. Called outside the queue monitor.
+ */
+ private static void closeLoadingSegment(final AcquireSegmentAction
acquireSegmentAction)
+ {
+ 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/12141)
--
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]
