diff --git a/core/trino-main/src/main/java/io/trino/execution/StageInfo.java b/core/trino-main/src/main/java/io/trino/execution/StageInfo.java index 24b162e64a6b..cbe8977ac92c 100644 --- a/core/trino-main/src/main/java/io/trino/execution/StageInfo.java +++ b/core/trino-main/src/main/java/io/trino/execution/StageInfo.java @@ -48,6 +48,7 @@ public record StageInfo( requireNonNull(subStages, "subStages is null"); requireNonNull(tables, "tables is null"); tasks = ImmutableList.copyOf(tasks); + subStages = ImmutableList.copyOf(subStages); tables = ImmutableMap.copyOf(tables); } diff --git a/core/trino-main/src/main/java/io/trino/execution/executor/dedicated/TaskEntry.java b/core/trino-main/src/main/java/io/trino/execution/executor/dedicated/TaskEntry.java index 8b9d5ba8e3af..0a16e395bc7d 100644 --- a/core/trino-main/src/main/java/io/trino/execution/executor/dedicated/TaskEntry.java +++ b/core/trino-main/src/main/java/io/trino/execution/executor/dedicated/TaskEntry.java @@ -25,14 +25,19 @@ import io.trino.execution.executor.scheduler.Schedulable; import io.trino.execution.executor.scheduler.SchedulerContext; import io.trino.spi.VersionEmbedder; +import jakarta.annotation.Nullable; +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Deque; import java.util.HashSet; -import java.util.LinkedList; -import java.util.Queue; +import java.util.List; import java.util.Set; import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Consumer; import java.util.function.DoubleSupplier; +import static com.google.common.base.Throwables.throwIfUnchecked; import static com.google.common.util.concurrent.MoreExecutors.directExecutor; import static java.util.Objects.requireNonNull; @@ -53,13 +58,13 @@ class TaskEntry private volatile boolean destroyed; @GuardedBy("this") - private int runningLeafSplits; + private final Deque pending = new ArrayDeque<>(); @GuardedBy("this") - private final Queue pending = new LinkedList<>(); + private final Set running = new HashSet<>(); @GuardedBy("this") - private final Set running = new HashSet<>(); + private final Set runningLeafSplits = new HashSet<>(); public TaskEntry(TaskId taskId, FairScheduler scheduler, VersionEmbedder versionEmbedder, Tracer tracer, int initialConcurrency, DoubleSupplier utilization) { @@ -78,79 +83,197 @@ public TaskId taskId() return taskId; } - public synchronized void destroy() + public void destroy() { - if (destroyed) { - return; + List runningSplits; + List pendingSplits; + List claimedSplits; + synchronized (this) { + if (destroyed) { + return; + } + + destroyed = true; + runningSplits = new ArrayList<>(running); + pendingSplits = new ArrayList<>(pending); + claimedSplits = new ArrayList<>(runningLeafSplits); + running.clear(); + pending.clear(); + runningLeafSplits.clear(); } - scheduler.removeGroup(group); + // Complete futures before closing. Driver.close() can rethrow operator close failures, and + // no such failure may strand another split or prevent its task from finishing. + pendingSplits.forEach(split -> split.done().set(null)); + claimedSplits.forEach(split -> split.done().set(null)); - destroyed = true; + Throwable failure = null; + try { + scheduler.removeGroup(group); + } + catch (Throwable t) { + failure = t; + } - for (SplitRunner split : running) { - split.close(); + for (SplitRunner split : runningSplits) { + failure = closeSplit(split, failure); + } + for (QueuedSplit split : pendingSplits) { + failure = closeSplit(split.split(), failure); } - running.clear(); - for (QueuedSplit split : pending) { - split.split().close(); - split.done.set(null); + if (failure != null) { + throwIfUnchecked(failure); + throw new RuntimeException(failure); } - pending.clear(); } - public synchronized ListenableFuture enqueueLeafSplit(SplitRunner split) + public ListenableFuture enqueueLeafSplit(SplitRunner split) { SettableFuture done = SettableFuture.create(); - pending.add(new QueuedSplit(split, done)); + synchronized (this) { + if (!destroyed) { + pending.addLast(new QueuedSplit(split, done)); + return done; + } + } + + // The task was removed concurrently with enqueueSplits. There is nobody left to claim the + // split, so finish it immediately rather than leaving its future in a destroyed queue. + done.set(null); + split.close(); return done; } - /** - * @return true if a split was scheduled; false if no splits are pending - */ - public synchronized boolean dequeueAndRunLeafSplit(Runnable doneCallback) + /// Claim the next pending leaf split and account for it as running. The claimed split must be + /// handed to [#startLeafSplit] to actually run; the two steps are separate so the caller can + /// start the split without holding any lock. + /// + /// The task can be destroyed between claiming and starting. A production `SplitRunner` + /// tolerates being started after close, and the scheduler rejects work for the removed group. + /// + /// @return null if no splits are pending or the task has been destroyed + @Nullable + public synchronized QueuedSplit claimLeafSplit() { + if (destroyed) { + return null; + } + QueuedSplit split = pending.poll(); if (split == null) { - return false; + return null; } - runSplit(split.split()) + runningLeafSplits.add(split); + running.add(split.split()); + + return split; + } + + /// Start a split claimed via [#claimLeafSplit()]. Must be called without holding a lock. + /// If this throws, the claim was not consumed and must be requeued via [#releaseLeafSplit]. + public void startLeafSplit(QueuedSplit split, Consumer doneCallback) + { + submit(split.split()) .addListener(() -> { - leafSplitDone(split); - doneCallback.run(); + finishLeafSplit(split, doneCallback); }, directExecutor()); + } - runningLeafSplits++; + /// Requeue a claimed split that could not be started. Thread creation can fail temporarily on + /// an overloaded worker, and the periodic scheduling pass will retry it after capacity clears. + public void releaseLeafSplit(QueuedSplit split) + { + boolean destroyed; + synchronized (this) { + runningLeafSplits.remove(split); + running.remove(split.split()); + destroyed = this.destroyed; + if (!destroyed) { + pending.addFirst(split); + } + } - return true; + if (destroyed) { + split.done().set(null); + } } - private synchronized void leafSplitDone(QueuedSplit split) + private void finishLeafSplit(QueuedSplit split, Consumer doneCallback) { - runningLeafSplits--; + boolean close; + synchronized (this) { + runningLeafSplits.remove(split); + close = running.remove(split.split()); + } + + // Complete the future and notify the executor even when Driver.close() fails. split.done().set(null); + try { + if (close) { + split.split().close(); + } + } + finally { + doneCallback.accept(this); + } + } + + public ListenableFuture runSplit(SplitRunner split) + { + boolean destroyed; + synchronized (this) { + destroyed = this.destroyed; + if (!destroyed) { + running.add(split); + } + } + + if (destroyed) { + SettableFuture done = SettableFuture.create(); + done.set(null); + split.close(); + return done; + } + + try { + ListenableFuture done = submit(split); + done.addListener(() -> splitDone(split), directExecutor()); + return done; + } + catch (Throwable t) { + try { + splitDone(split); + } + catch (Throwable closeFailure) { + t.addSuppressed(closeFailure); + } + throw t; + } } - public synchronized ListenableFuture runSplit(SplitRunner split) + /// Hand a split that is already accounted for in `running` to the scheduler. Must be called + /// without holding a lock: the scheduler creates the thread that runs the split, which on a + /// loaded worker is slow enough to stall every other operation on the lock. + private ListenableFuture submit(SplitRunner split) { int splitId = nextSplitId(); - ListenableFuture done = scheduler.submit( + return scheduler.submit( group, splitId, new VersionEmbedderBridge(versionEmbedder, new SplitProcessor(taskId, splitId, split, tracer))); - done.addListener(() -> splitDone(split), directExecutor()); - running.add(split); - - return done; } - private synchronized void splitDone(SplitRunner split) + private void splitDone(SplitRunner split) { - split.close(); - running.remove(split); + boolean close; + synchronized (this) { + close = running.remove(split); + } + if (close) { + split.close(); + } } private int nextSplitId() @@ -160,7 +283,7 @@ private int nextSplitId() public synchronized int runningLeafSplits() { - return runningLeafSplits; + return runningLeafSplits.size(); } @Override @@ -171,7 +294,7 @@ public boolean isDestroyed() public synchronized void updateConcurrency() { - concurrency.update(utilization.getAsDouble(), runningLeafSplits); + concurrency.update(utilization.getAsDouble(), runningLeafSplits.size()); } public synchronized int pendingLeafSplitCount() @@ -194,7 +317,21 @@ public synchronized int targetConcurrency() return concurrency.targetConcurrency(); } - private record QueuedSplit(SplitRunner split, SettableFuture done) {} + record QueuedSplit(SplitRunner split, SettableFuture done) {} + + private static Throwable closeSplit(SplitRunner split, @Nullable Throwable failure) + { + try { + split.close(); + } + catch (Throwable t) { + if (failure == null) { + return t; + } + failure.addSuppressed(t); + } + return failure; + } private record VersionEmbedderBridge(VersionEmbedder versionEmbedder, Schedulable delegate) implements Schedulable diff --git a/core/trino-main/src/main/java/io/trino/execution/executor/dedicated/ThreadPerDriverTaskExecutor.java b/core/trino-main/src/main/java/io/trino/execution/executor/dedicated/ThreadPerDriverTaskExecutor.java index 7281f768dabe..b2edabb918a6 100644 --- a/core/trino-main/src/main/java/io/trino/execution/executor/dedicated/ThreadPerDriverTaskExecutor.java +++ b/core/trino-main/src/main/java/io/trino/execution/executor/dedicated/ThreadPerDriverTaskExecutor.java @@ -15,6 +15,7 @@ import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Ticker; +import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; import com.google.common.util.concurrent.ListenableFuture; import com.google.errorprone.annotations.ThreadSafe; @@ -30,27 +31,31 @@ import io.trino.execution.executor.RunningSplitInfo; import io.trino.execution.executor.TaskExecutor; import io.trino.execution.executor.TaskHandle; +import io.trino.execution.executor.dedicated.TaskEntry.QueuedSplit; import io.trino.execution.executor.scheduler.FairScheduler; import io.trino.spi.VersionEmbedder; +import jakarta.annotation.Nullable; import jakarta.annotation.PostConstruct; import jakarta.annotation.PreDestroy; import org.weakref.jmx.Managed; import org.weakref.jmx.Nested; -import java.util.ArrayDeque; import java.util.ArrayList; -import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.OptionalInt; -import java.util.Queue; import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ScheduledThreadPoolExecutor; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; +import java.util.function.Consumer; import java.util.function.DoubleSupplier; import java.util.function.Predicate; import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Throwables.throwIfUnchecked; import static io.airlift.concurrent.Threads.daemonThreadsNamed; import static java.lang.Math.max; import static java.lang.Math.min; @@ -61,6 +66,7 @@ public class ThreadPerDriverTaskExecutor implements TaskExecutor { private static final Logger LOG = Logger.get(ThreadPerDriverTaskExecutor.class); + private static final long FAILURE_LOG_INTERVAL_NANOS = TimeUnit.MINUTES.toNanos(1); private final FairScheduler scheduler; private final Tracer tracer; @@ -70,17 +76,20 @@ public class ThreadPerDriverTaskExecutor private final int maxDriversPerTask; private final ScheduledThreadPoolExecutor backgroundTasks = new ScheduledThreadPoolExecutor(2, daemonThreadsNamed("task-executor-scheduler-%s")); - @GuardedBy("this") - private final Map tasks = new HashMap<>(); + private final Map tasks = new ConcurrentHashMap<>(); - @GuardedBy("this") - private boolean closed; + private volatile boolean closed; @GuardedBy("this") private int runningLeafDrivers; + private final AtomicBoolean schedulingLeafSplits = new AtomicBoolean(); + private final AtomicBoolean rescheduleLeafSplits = new AtomicBoolean(); + // Do not inline this field to avoid creating lambdas that cannot be cached by JVM. - private final Runnable leafSplitDoneCallback = this::leafSplitDone; + private final Consumer leafSplitDoneCallback = this::leafSplitDone; + private final FailureLogger schedulingFailureLogger = new FailureLogger("Error scheduling leaf splits"); + private final Runnable scheduleMoreLeafSplitsQuietly = maintenance(this::scheduleMoreLeafSplits, schedulingFailureLogger); @Inject public ThreadPerDriverTaskExecutor(TaskManagerConfig config, Tracer tracer, VersionEmbedder versionEmbedder) @@ -109,19 +118,41 @@ public ThreadPerDriverTaskExecutor(Tracer tracer, VersionEmbedder versionEmbedde public synchronized void start() { scheduler.start(); - backgroundTasks.scheduleWithFixedDelay(this::scheduleMoreLeafSplits, 0, 100, TimeUnit.MILLISECONDS); - backgroundTasks.scheduleWithFixedDelay(this::adjustConcurrency, 0, 10, TimeUnit.MILLISECONDS); - backgroundTasks.scheduleWithFixedDelay(this::logDiagnostics, 0, 30, TimeUnit.SECONDS); + backgroundTasks.scheduleWithFixedDelay(scheduleMoreLeafSplitsQuietly, 0, 100, TimeUnit.MILLISECONDS); + backgroundTasks.scheduleWithFixedDelay(maintenance(this::adjustConcurrency, "Error adjusting task concurrency"), 0, 10, TimeUnit.MILLISECONDS); + backgroundTasks.scheduleWithFixedDelay(maintenance(this::logDiagnostics, "Error logging diagnostics"), 0, 30, TimeUnit.SECONDS); } @PreDestroy @Override public synchronized void stop() { + if (closed) { + return; + } closed = true; - tasks.values().forEach(TaskEntry::destroy); + + Throwable failure = null; + for (TaskEntry task : tasks.values()) { + try { + task.destroy(); + } + catch (Throwable t) { + failure = addFailure(failure, t); + } + } backgroundTasks.shutdownNow(); - scheduler.close(); + try { + scheduler.close(); + } + catch (Throwable t) { + failure = addFailure(failure, t); + } + + if (failure != null) { + throwIfUnchecked(failure); + throw new RuntimeException(failure); + } } @Override @@ -148,22 +179,18 @@ public synchronized TaskHandle addTask( public void removeTask(TaskHandle handle) { TaskEntry entry = (TaskEntry) handle; - synchronized (this) { - tasks.remove(entry.taskId()); - } - if (!entry.isDestroyed()) { - entry.destroy(); - } + tasks.remove(entry.taskId(), entry); + entry.destroy(); } @Override - public synchronized List> enqueueSplits(TaskHandle handle, boolean intermediate, List splits) + public List> enqueueSplits(TaskHandle handle, boolean intermediate, List splits) { checkArgument(!closed, "Executor is already closed"); TaskEntry entry = (TaskEntry) handle; - List> futures = new ArrayList<>(); + List> futures = new ArrayList<>(splits.size()); for (SplitRunner split : splits) { if (intermediate) { futures.add(entry.runSplit(split)); @@ -173,48 +200,229 @@ public synchronized List> enqueueSplits(TaskHandle handle } } - scheduleMoreLeafSplits(); + scheduleMoreLeafSplitsQuietly.run(); return futures; } - private boolean scheduleLeafSplit(TaskEntry task) + private void leafSplitDone(TaskEntry task) { - boolean scheduled = task.dequeueAndRunLeafSplit(leafSplitDoneCallback); - if (scheduled) { - runningLeafDrivers++; + ClaimedLeafSplit replacement; + synchronized (this) { + runningLeafDrivers--; + replacement = claimLeafSplitForTask(task); + } + + if (replacement != null) { + startClaimedSplits(ImmutableList.of(replacement)); + return; + } + + // When the task drains, let another task use the freed slot immediately. This runs the + // global pass once per task drain instead of once per split completion. + if (!task.hasPendingLeafSplits()) { + scheduleMoreLeafSplitsQuietly.run(); } + } - return scheduled; + @GuardedBy("this") + @Nullable + private ClaimedLeafSplit claimLeafSplitForTask(TaskEntry task) + { + if (closed) { + return null; + } + + int taskRunning = task.runningLeafSplits(); + if (taskRunning >= minDriversPerTask && + (runningLeafDrivers >= targetGlobalLeafDrivers || taskRunning >= min(task.targetConcurrency(), maxDriversPerTask))) { + return null; + } + + List claimed = new ArrayList<>(1); + if (!claimLeafSplit(task, claimed)) { + return null; + } + return claimed.getFirst(); + } + + @VisibleForTesting + void scheduleMoreLeafSplits() + { + rescheduleLeafSplits.set(true); + boolean retry = true; + while (true) { + if (!schedulingLeafSplits.compareAndSet(false, true)) { + return; + } + + try { + do { + rescheduleLeafSplits.set(false); + retry = startClaimedSplits(claimMoreLeafSplits()); + } + while (retry && rescheduleLeafSplits.get()); + } + finally { + schedulingLeafSplits.set(false); + } + + // Close the race where another caller requested a pass after the last flag check but + // before the gate was released. + if (!retry || !rescheduleLeafSplits.get()) { + return; + } + } + } + + private boolean startClaimedSplits(List claimed) + { + // Start the splits outside all locks. Thread creation is slow on an overloaded worker. + for (int i = 0; i < claimed.size(); i++) { + ClaimedLeafSplit split = claimed.get(i); + try { + split.task().startLeafSplit(split.split(), leafSplitDoneCallback); + } + catch (Throwable e) { + // Claims before this one have listeners and remain running. Requeue this claim and + // every unstarted claim behind it; the periodic pass retries them later. + releaseClaims(claimed.subList(i, claimed.size())); + schedulingFailureLogger.log(e); + return false; + } + } + return true; } - private synchronized void leafSplitDone() + private void releaseClaims(List claimed) { - runningLeafDrivers--; - scheduleMoreLeafSplits(); + synchronized (this) { + runningLeafDrivers -= claimed.size(); + } + + // Each task requeues at the head, so release in reverse to preserve claim order. + for (int i = claimed.size() - 1; i >= 0; i--) { + ClaimedLeafSplit split = claimed.get(i); + split.task().releaseLeafSplit(split.split()); + } } - private synchronized void scheduleMoreLeafSplits() + private synchronized List claimMoreLeafSplits() { - // schedule minimum guaranteed leaf drivers for each task + if (closed) { + return ImmutableList.of(); + } + + List claimed = new ArrayList<>(); + + // claim minimum guaranteed leaf drivers for each task for (TaskEntry task : tasks.values()) { int target = max(0, minDriversPerTask - task.runningLeafSplits()); for (int i = 0; i < target; i++) { - if (!scheduleLeafSplit(task)) { + if (!claimLeafSplit(task, claimed)) { break; } } } - // schedule additional drivers up to the target global leaf drivers - Queue queue = new ArrayDeque<>(tasks.values()); + // Claim additional drivers up to the target global leaf drivers. Iterate in rounds to + // retain the previous round-robin behavior without copying the task map into a queue. int target = targetGlobalLeafDrivers - runningLeafDrivers; - for (int i = 0; i < target && !queue.isEmpty(); i++) { - TaskEntry task = queue.poll(); - if (task.runningLeafSplits() < min(task.targetConcurrency(), maxDriversPerTask)) { - scheduleLeafSplit(task); - if (task.hasPendingLeafSplits()) { - queue.add(task); + boolean progress = true; + while (target > 0 && progress) { + progress = false; + for (TaskEntry task : tasks.values()) { + if (target == 0) { + break; } + if (task.runningLeafSplits() < min(task.targetConcurrency(), maxDriversPerTask) && claimLeafSplit(task, claimed)) { + target--; + progress = true; + } + } + } + + return claimed; + } + + @GuardedBy("this") + private boolean claimLeafSplit(TaskEntry task, List claimed) + { + QueuedSplit split = task.claimLeafSplit(); + if (split == null) { + return false; + } + + runningLeafDrivers++; + claimed.add(new ClaimedLeafSplit(task, split)); + + return true; + } + + private record ClaimedLeafSplit(TaskEntry task, QueuedSplit split) {} + + /// Wrap a task so that a failure does not take its caller down with it. + /// [ScheduledThreadPoolExecutor#scheduleWithFixedDelay] silently stops rescheduling a task + /// that throws, which would leave the worker permanently without leaf split scheduling or + /// concurrency adjustment. Failures here are typically symptoms of an overloaded worker, + /// such as being unable to create a thread, and it is expected to recover once load subsides. + @VisibleForTesting + static Runnable maintenance(Runnable task, String errorMessage) + { + return maintenance(task, new FailureLogger(errorMessage)); + } + + private static Runnable maintenance(Runnable task, FailureLogger failureLogger) + { + return () -> { + try { + task.run(); + } + catch (Throwable e) { + failureLogger.log(e); + } + }; + } + + private static Throwable addFailure(Throwable failure, Throwable newFailure) + { + if (failure == null) { + return newFailure; + } + failure.addSuppressed(newFailure); + return failure; + } + + private static final class FailureLogger + { + private final String message; + private final AtomicLong lastLogNanos = new AtomicLong(Long.MIN_VALUE); + private final AtomicLong suppressedFailures = new AtomicLong(); + + private FailureLogger(String message) + { + this.message = requireNonNull(message, "message is null"); + } + + public void log(Throwable failure) + { + long now = System.nanoTime(); + while (true) { + long last = lastLogNanos.get(); + if (last != Long.MIN_VALUE && now - last < FAILURE_LOG_INTERVAL_NANOS) { + suppressedFailures.incrementAndGet(); + return; + } + if (lastLogNanos.compareAndSet(last, now)) { + break; + } + } + + long suppressed = suppressedFailures.getAndSet(0); + if (suppressed == 0) { + LOG.warn(failure, "%s", message); + } + else { + LOG.warn(failure, "%s (%s similar failures suppressed)", message, suppressed); } } } @@ -255,13 +463,13 @@ public Set getStuckSplitTaskIds(Duration processingDurationThreshold, Pr } @Managed - public synchronized int getTasks() + public int getTasks() { return tasks.size(); } @Managed - public synchronized int getTotalRunningSplits() + public int getTotalRunningSplits() { return tasks.values().stream() .mapToInt(TaskEntry::totalRunningSplits) @@ -269,7 +477,7 @@ public synchronized int getTotalRunningSplits() } @Managed - public synchronized int getTotalRunningLeafSplits() + public int getTotalRunningLeafSplits() { return tasks.values().stream() .mapToInt(TaskEntry::runningLeafSplits) @@ -277,7 +485,7 @@ public synchronized int getTotalRunningLeafSplits() } @Managed - public synchronized int getTotalPendingLeafSplits() + public int getTotalPendingLeafSplits() { return tasks.values().stream() .mapToInt(TaskEntry::pendingLeafSplitCount) diff --git a/core/trino-main/src/main/java/io/trino/execution/executor/scheduler/FairScheduler.java b/core/trino-main/src/main/java/io/trino/execution/executor/scheduler/FairScheduler.java index a904c301019d..ad1192a96796 100644 --- a/core/trino-main/src/main/java/io/trino/execution/executor/scheduler/FairScheduler.java +++ b/core/trino-main/src/main/java/io/trino/execution/executor/scheduler/FairScheduler.java @@ -13,12 +13,12 @@ */ package io.trino.execution.executor.scheduler; +import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Ticker; import com.google.common.util.concurrent.ListenableFuture; import com.google.common.util.concurrent.ListeningExecutorService; import com.google.common.util.concurrent.MoreExecutors; import com.google.errorprone.annotations.ThreadSafe; -import com.google.errorprone.annotations.concurrent.GuardedBy; import io.airlift.concurrent.ThreadPoolExecutorMBean; import io.airlift.log.Logger; @@ -27,8 +27,11 @@ import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.SynchronousQueue; +import java.util.concurrent.ThreadFactory; import java.util.concurrent.ThreadPoolExecutor; import java.util.concurrent.TimeUnit; +import java.util.concurrent.locks.ReadWriteLock; +import java.util.concurrent.locks.ReentrantReadWriteLock; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Preconditions.checkState; @@ -65,10 +68,20 @@ public final class FairScheduler private final Gate paused = new Gate(true); - @GuardedBy("this") - private boolean closed; + /// Prevents group creation and removal from racing with [#close()]. Submission only needs the + /// volatile `closed` check: a task accepted while close is in progress either observes that its + /// group was removed, or is rejected by the executor after shutdown. + private final ReadWriteLock lifecycleLock = new ReentrantReadWriteLock(); + + private volatile boolean closed; public FairScheduler(int maxConcurrentTasks, String threadNameFormat, Ticker ticker) + { + this(maxConcurrentTasks, daemonThreadsNamed(threadNameFormat), ticker); + } + + @VisibleForTesting + public FairScheduler(int maxConcurrentTasks, ThreadFactory threadFactory, Ticker ticker) { this.ticker = requireNonNull(ticker, "ticker is null"); @@ -77,7 +90,7 @@ public FairScheduler(int maxConcurrentTasks, String threadNameFormat, Ticker tic schedulerExecutor = Executors.newCachedThreadPool(daemonThreadsNamed("fair-scheduler-%d")); schedulerExecutorMBean = new ThreadPoolExecutorMBean((ThreadPoolExecutor) schedulerExecutor); - executor = new ThreadPoolExecutor(0, Integer.MAX_VALUE, 60L, TimeUnit.SECONDS, new SynchronousQueue<>(), daemonThreadsNamed(threadNameFormat)); + executor = new ThreadPoolExecutor(0, Integer.MAX_VALUE, 60L, TimeUnit.SECONDS, new SynchronousQueue<>(), threadFactory); executorMBean = new ThreadPoolExecutorMBean(executor); taskExecutor = MoreExecutors.listeningDecorator(executor); } @@ -110,41 +123,59 @@ public void resume() } @Override - public synchronized void close() + public void close() { - if (closed) { - return; - } - closed = true; + lifecycleLock.writeLock().lock(); + try { + if (closed) { + return; + } + closed = true; - Set tasks = queue.finishAll(); + Set tasks = queue.finishAll(); - for (TaskControl task : tasks) { - task.cancel(); - } + for (TaskControl task : tasks) { + task.cancel(); + } - taskExecutor.shutdownNow(); - schedulerExecutor.shutdownNow(); + taskExecutor.shutdownNow(); + schedulerExecutor.shutdownNow(); + } + finally { + lifecycleLock.writeLock().unlock(); + } } - public synchronized Group createGroup(String name) + public Group createGroup(String name) { - checkArgument(!closed, "Already closed"); + lifecycleLock.readLock().lock(); + try { + checkArgument(!closed, "Already closed"); - Group group = new Group(name); - queue.startGroup(group); + Group group = new Group(name); + queue.startGroup(group); - return group; + return group; + } + finally { + lifecycleLock.readLock().unlock(); + } } - public synchronized void removeGroup(Group group) + public void removeGroup(Group group) { - checkArgument(!closed, "Already closed"); + lifecycleLock.readLock().lock(); + try { + checkArgument(!closed, "Already closed"); - Set tasks = queue.finishGroup(group); + Set tasks = queue.finishGroup(group); - for (TaskControl task : tasks) { - task.cancel(); + for (TaskControl task : tasks) { + task.cancel(); + } + } + finally { + lifecycleLock.readLock().unlock(); } } @@ -155,7 +186,7 @@ public Set getTasks(Group group) .collect(toImmutableSet()); } - public synchronized ListenableFuture submit(Group group, int id, Schedulable runner) + public ListenableFuture submit(Group group, int id, Schedulable runner) { checkArgument(!closed, "Already closed"); diff --git a/core/trino-main/src/main/java/io/trino/operator/AssignUniqueIdOperator.java b/core/trino-main/src/main/java/io/trino/operator/AssignUniqueIdOperator.java index 391127c20b65..ae85c554f45c 100644 --- a/core/trino-main/src/main/java/io/trino/operator/AssignUniqueIdOperator.java +++ b/core/trino-main/src/main/java/io/trino/operator/AssignUniqueIdOperator.java @@ -37,9 +37,9 @@ public class AssignUniqueIdOperator private static final long ROW_IDS_PER_REQUEST = 1L << 20L; private static final long MAX_ROW_ID = 1L << 40L; - public static OperatorFactory createOperatorFactory(int operatorId, PlanNodeId planNodeId) + public static OperatorFactory createOperatorFactory(int operatorId, PlanNodeId planNodeId, AtomicLong valuePool) { - return createAdapterOperatorFactory(new Factory(operatorId, planNodeId)); + return createAdapterOperatorFactory(new Factory(operatorId, planNodeId, valuePool)); } private static class Factory @@ -48,13 +48,10 @@ private static class Factory private final int operatorId; private final PlanNodeId planNodeId; private boolean closed; + // The unique id embeds only the stage and partition of the task, so all AssignUniqueId + // operators in a task must draw row ids from a single pool for their ids to be distinct private final AtomicLong valuePool; - private Factory(int operatorId, PlanNodeId planNodeId) - { - this(operatorId, planNodeId, new AtomicLong()); - } - private Factory(int operatorId, PlanNodeId planNodeId, AtomicLong valuePool) { this.operatorId = operatorId; diff --git a/core/trino-main/src/main/java/io/trino/sql/planner/LocalExecutionPlanner.java b/core/trino-main/src/main/java/io/trino/sql/planner/LocalExecutionPlanner.java index c29299eb808b..474d7cfd9d3b 100644 --- a/core/trino-main/src/main/java/io/trino/sql/planner/LocalExecutionPlanner.java +++ b/core/trino-main/src/main/java/io/trino/sql/planner/LocalExecutionPlanner.java @@ -303,6 +303,7 @@ import java.util.OptionalInt; import java.util.Set; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; import java.util.function.Consumer; import java.util.function.Function; import java.util.function.Supplier; @@ -686,6 +687,8 @@ private static class LocalExecutionPlanContext // this is shared with all subContexts private final AtomicInteger nextPipelineId; + // this is shared with all subContexts; see AssignUniqueIdOperator.Factory for the rationale + private final AtomicLong assignUniqueIdValuePool; private int nextOperatorId; private boolean inputDriver = true; @@ -696,19 +699,22 @@ public LocalExecutionPlanContext(TaskContext taskContext) this(taskContext, new ArrayList<>(), Optional.empty(), - new AtomicInteger(0)); + new AtomicInteger(0), + new AtomicLong()); } private LocalExecutionPlanContext( TaskContext taskContext, List driverFactories, Optional indexSourceContext, - AtomicInteger nextPipelineId) + AtomicInteger nextPipelineId, + AtomicLong assignUniqueIdValuePool) { this.taskContext = taskContext; this.driverFactories = driverFactories; this.indexSourceContext = indexSourceContext; this.nextPipelineId = nextPipelineId; + this.assignUniqueIdValuePool = assignUniqueIdValuePool; } public void addDriverFactory(boolean outputDriver, PhysicalOperation physicalOperation, LocalExecutionPlanContext context) @@ -802,6 +808,11 @@ private int getNextOperatorId() return nextOperatorId++; } + private AtomicLong getAssignUniqueIdValuePool() + { + return assignUniqueIdValuePool; + } + private boolean isInputDriver() { return inputDriver; @@ -815,12 +826,12 @@ private void setInputDriver(boolean inputDriver) public LocalExecutionPlanContext createSubContext() { checkState(indexSourceContext.isEmpty(), "index build plan cannot have sub-contexts"); - return new LocalExecutionPlanContext(taskContext, driverFactories, indexSourceContext, nextPipelineId); + return new LocalExecutionPlanContext(taskContext, driverFactories, indexSourceContext, nextPipelineId, assignUniqueIdValuePool); } public LocalExecutionPlanContext createIndexSourceSubContext(IndexSourceContext indexSourceContext) { - return new LocalExecutionPlanContext(taskContext, driverFactories, Optional.of(indexSourceContext), nextPipelineId); + return new LocalExecutionPlanContext(taskContext, driverFactories, Optional.of(indexSourceContext), nextPipelineId, assignUniqueIdValuePool); } public OptionalInt getDriverInstanceCount() @@ -3645,7 +3656,8 @@ public PhysicalOperation visitAssignUniqueId(AssignUniqueId node, LocalExecution OperatorFactory operatorFactory = AssignUniqueIdOperator.createOperatorFactory( context.getNextOperatorId(), - node.getId()); + node.getId(), + context.getAssignUniqueIdValuePool()); return new PhysicalOperation(operatorFactory, makeLayout(node), source); } diff --git a/core/trino-main/src/test/java/io/trino/execution/executor/dedicated/TestThreadPerDriverTaskExecutor.java b/core/trino-main/src/test/java/io/trino/execution/executor/dedicated/TestThreadPerDriverTaskExecutor.java index 1dffac946808..3a322b08367b 100644 --- a/core/trino-main/src/test/java/io/trino/execution/executor/dedicated/TestThreadPerDriverTaskExecutor.java +++ b/core/trino-main/src/test/java/io/trino/execution/executor/dedicated/TestThreadPerDriverTaskExecutor.java @@ -13,10 +13,12 @@ */ package io.trino.execution.executor.dedicated; +import com.google.common.base.Ticker; import com.google.common.collect.ImmutableList; import com.google.common.util.concurrent.AbstractFuture; import com.google.common.util.concurrent.Futures; import com.google.common.util.concurrent.ListenableFuture; +import com.google.common.util.concurrent.SettableFuture; import io.airlift.testing.TestingTicker; import io.airlift.units.Duration; import io.opentelemetry.api.trace.Span; @@ -34,17 +36,352 @@ import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Future; import java.util.concurrent.Phaser; +import java.util.concurrent.ThreadFactory; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Function; +import static com.google.common.util.concurrent.Uninterruptibles.awaitUninterruptibly; +import static io.airlift.concurrent.Threads.daemonThreadsNamed; import static io.airlift.tracing.Tracing.noopTracer; import static io.trino.util.EmbedVersion.testingVersionEmbedder; +import static java.util.concurrent.Executors.newSingleThreadExecutor; import static java.util.concurrent.TimeUnit.MILLISECONDS; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; public class TestThreadPerDriverTaskExecutor { + @Test + @Timeout(30) + public void testSlowSplitStartDoesNotBlockTaskManagement() + throws Exception + { + // A worker under load takes a long time to create a thread. While a split is being started, + // tasks must still be able to come and go, otherwise the worker cannot shed the load that + // made thread creation slow in the first place. + AtomicBoolean stallNextThread = new AtomicBoolean(); + CountDownLatch threadCreationStarted = new CountDownLatch(1); + CountDownLatch releaseThreadCreation = new CountDownLatch(1); + + ThreadFactory threadFactory = runnable -> { + if (stallNextThread.compareAndSet(true, false)) { + threadCreationStarted.countDown(); + awaitUninterruptibly(releaseThreadCreation); + } + Thread thread = new Thread(runnable); + thread.setDaemon(true); + return thread; + }; + + FairScheduler scheduler = new FairScheduler(3, threadFactory, Ticker.systemTicker()); + ThreadPerDriverTaskExecutor executor = new ThreadPerDriverTaskExecutor(noopTracer(), testingVersionEmbedder(), scheduler, 1, 1, Integer.MAX_VALUE); + // Do not start the executor's periodic scheduler: the replacement below must be started by + // the completion callback whose progress this test verifies. + scheduler.start(); + ExecutorService submitter = newSingleThreadExecutor(daemonThreadsNamed("submitter")); + try { + TaskHandle task = executor.addTask(new TaskId(new StageId("query", 1), 1, 1), () -> 0, 10, new Duration(1, MILLISECONDS), OptionalInt.empty()); + TestFuture firstBlocked = new TestFuture(); + ListenableFuture firstDone = executor.enqueueSplits( + task, + false, + ImmutableList.of(new TestingSplitRunner(ImmutableList.of( + _ -> firstBlocked, + _ -> Futures.immediateVoidFuture())))) + .getFirst(); + firstBlocked.awaitListenerAdded(); + + CountDownLatch replacementStarted = new CountDownLatch(1); + ListenableFuture replacementDone = executor.enqueueSplits( + task, + false, + ImmutableList.of(new TestingSplitRunner(ImmutableList.of(_ -> { + replacementStarted.countDown(); + return Futures.immediateVoidFuture(); + })))) + .getFirst(); + + TaskHandle stalledTask = executor.addTask(new TaskId(new StageId("query", 1), 2, 1), () -> 0, 10, new Duration(1, MILLISECONDS), OptionalInt.empty()); + SplitRunner stalledSplit = new TestingSplitRunner(ImmutableList.of(_ -> Futures.immediateVoidFuture())); + + stallNextThread.set(true); + Future stalled = submitter.submit(() -> executor.enqueueSplits(stalledTask, false, ImmutableList.of(stalledSplit))); + threadCreationStarted.await(); + + // Registering, removing and completing a leaf split must not wait for the stalled + // thread creation. + TaskId otherTaskId = new TaskId(new StageId("query", 1), 3, 1); + TaskHandle other = executor.addTask(otherTaskId, () -> 0, 10, new Duration(1, MILLISECONDS), OptionalInt.empty()); + executor.removeTask(other); + firstBlocked.set(null); + firstDone.get(10, TimeUnit.SECONDS); + assertThat(replacementStarted.await(10, TimeUnit.SECONDS)).isTrue(); + replacementDone.get(10, TimeUnit.SECONDS); + + releaseThreadCreation.countDown(); + stalled.get(); + } + finally { + releaseThreadCreation.countDown(); + submitter.shutdownNow(); + executor.stop(); + } + } + + @Test + @Timeout(30) + public void testFailureToStartSplitRequeuesItsClaim() + throws Exception + { + // Starting a split creates its thread, which an overloaded worker fails to do with + // OutOfMemoryError. The leaf driver accounting has to survive that, otherwise the worker + // permanently loses the capacity of every split it failed to start. + AtomicBoolean failNextThread = new AtomicBoolean(); + ThreadFactory threadFactory = runnable -> { + if (failNextThread.compareAndSet(true, false)) { + throw new OutOfMemoryError("unable to create native thread"); + } + Thread thread = new Thread(runnable); + thread.setDaemon(true); + return thread; + }; + + FairScheduler scheduler = new FairScheduler(2, threadFactory, Ticker.systemTicker()); + // no minimum guarantee and a global budget of one leaf driver, so that a leaked claim is + // observable: it would leave targetGlobalLeafDrivers - runningLeafDrivers at zero and no + // later split could be scheduled + ThreadPerDriverTaskExecutor executor = new ThreadPerDriverTaskExecutor(noopTracer(), testingVersionEmbedder(), scheduler, 0, Integer.MAX_VALUE, 1); + // start the scheduler but not the executor's background tasks, so that splits are only + // scheduled by the explicit enqueueSplits calls below + scheduler.start(); + try { + TaskId taskId = new TaskId(new StageId("query", 1), 1, 1); + TaskEntry task = (TaskEntry) executor.addTask(taskId, () -> 0, 10, new Duration(1, MILLISECONDS), OptionalInt.empty()); + + // Queue the split without scheduling it, then fail its first start attempt. + ListenableFuture done = task.enqueueLeafSplit(new TestingSplitRunner(ImmutableList.of(_ -> Futures.immediateVoidFuture()))); + + failNextThread.set(true); + executor.scheduleMoreLeafSplits(); + + assertThat(done).isNotDone(); + assertThat(executor.getTotalRunningLeafSplits()).isEqualTo(0); + assertThat(executor.getTotalRunningSplits()).isEqualTo(0); + assertThat(executor.getTotalPendingLeafSplits()).isEqualTo(1); + + // A later pass retries the same split after thread capacity becomes available. + executor.scheduleMoreLeafSplits(); + done.get(); + } + finally { + executor.stop(); + } + } + + @Test + @Timeout(30) + public void testFailureToStartSecondSplitPreservesStartedClaim() + throws Exception + { + AtomicInteger threads = new AtomicInteger(); + ThreadFactory threadFactory = runnable -> { + if (threads.incrementAndGet() == 2) { + throw new OutOfMemoryError("unable to create native thread"); + } + Thread thread = new Thread(runnable); + thread.setDaemon(true); + return thread; + }; + + FairScheduler scheduler = new FairScheduler(2, threadFactory, Ticker.systemTicker()); + ThreadPerDriverTaskExecutor executor = new ThreadPerDriverTaskExecutor(noopTracer(), testingVersionEmbedder(), scheduler, 1, Integer.MAX_VALUE, 2); + scheduler.start(); + try { + TaskEntry firstTask = (TaskEntry) executor.addTask(new TaskId(new StageId("query", 1), 1, 1), () -> 0, 10, new Duration(1, MILLISECONDS), OptionalInt.empty()); + TaskEntry secondTask = (TaskEntry) executor.addTask(new TaskId(new StageId("query", 1), 2, 1), () -> 0, 10, new Duration(1, MILLISECONDS), OptionalInt.empty()); + SettableFuture blocked = SettableFuture.create(); + CountDownLatch started = new CountDownLatch(1); + ListenableFuture firstDone = firstTask.enqueueLeafSplit(new TestingSplitRunner(ImmutableList.of(_ -> { + started.countDown(); + return blocked; + }, _ -> Futures.immediateVoidFuture()))); + ListenableFuture secondDone = secondTask.enqueueLeafSplit(new TestingSplitRunner(ImmutableList.of(_ -> { + started.countDown(); + return blocked; + }, _ -> Futures.immediateVoidFuture()))); + + executor.scheduleMoreLeafSplits(); + started.await(); + + assertThat(executor.getTotalRunningLeafSplits()).isEqualTo(1); + assertThat(executor.getTotalRunningSplits()).isEqualTo(1); + assertThat(executor.getTotalPendingLeafSplits()).isEqualTo(1); + + blocked.set(null); + executor.scheduleMoreLeafSplits(); + firstDone.get(); + secondDone.get(); + } + finally { + executor.stop(); + } + } + + @Test + @Timeout(30) + public void testStartRegistersResilientSchedulingTask() + throws Exception + { + AtomicBoolean failNextThread = new AtomicBoolean(true); + ThreadFactory threadFactory = runnable -> { + if (failNextThread.compareAndSet(true, false)) { + throw new OutOfMemoryError("unable to create native thread"); + } + Thread thread = new Thread(runnable); + thread.setDaemon(true); + return thread; + }; + + FairScheduler scheduler = new FairScheduler(1, threadFactory, Ticker.systemTicker()); + ThreadPerDriverTaskExecutor executor = new ThreadPerDriverTaskExecutor(noopTracer(), testingVersionEmbedder(), scheduler, 0, Integer.MAX_VALUE, 1); + TaskEntry task = (TaskEntry) executor.addTask(new TaskId(new StageId("query", 1), 1, 1), () -> 0, 10, new Duration(1, MILLISECONDS), OptionalInt.empty()); + ListenableFuture done = task.enqueueLeafSplit(new TestingSplitRunner(ImmutableList.of(_ -> Futures.immediateVoidFuture()))); + + executor.start(); + try { + // The immediate pass fails. The wrapped fixed-delay task must remain registered so the + // 100ms retry can start the split. + done.get(10, TimeUnit.SECONDS); + } + finally { + executor.stop(); + } + } + + @Test + @Timeout(30) + public void testLeafCompletionSchedulesReplacementWhenCloseFails() + throws Exception + { + FairScheduler scheduler = FairScheduler.newInstance(1); + ThreadPerDriverTaskExecutor executor = new ThreadPerDriverTaskExecutor(noopTracer(), testingVersionEmbedder(), scheduler, 1, 1, 1); + try { + TaskEntry task = (TaskEntry) executor.addTask(new TaskId(new StageId("query", 1), 1, 1), () -> 0, 10, new Duration(1, MILLISECONDS), OptionalInt.empty()); + ListenableFuture failingDone = task.enqueueLeafSplit(new CloseFailingSplitRunner()); + CountDownLatch replacementStarted = new CountDownLatch(1); + ListenableFuture replacementDone = task.enqueueLeafSplit(new TestingSplitRunner(ImmutableList.of(_ -> { + replacementStarted.countDown(); + return Futures.immediateVoidFuture(); + }))); + + executor.scheduleMoreLeafSplits(); + + failingDone.get(10, TimeUnit.SECONDS); + assertThat(replacementStarted.await(10, TimeUnit.SECONDS)).isTrue(); + replacementDone.get(10, TimeUnit.SECONDS); + } + finally { + executor.stop(); + } + } + + @Test + @Timeout(30) + public void testLeafCompletionSchedulesAnotherTaskWhenTaskDrains() + throws Exception + { + FairScheduler scheduler = FairScheduler.newInstance(1); + ThreadPerDriverTaskExecutor executor = new ThreadPerDriverTaskExecutor(noopTracer(), testingVersionEmbedder(), scheduler, 0, 1, 1); + try { + TaskEntry firstTask = (TaskEntry) executor.addTask(new TaskId(new StageId("query", 1), 1, 1), () -> 0, 10, new Duration(1, MILLISECONDS), OptionalInt.empty()); + TestFuture firstBlocked = new TestFuture(); + ListenableFuture firstDone = firstTask.enqueueLeafSplit(new TestingSplitRunner(ImmutableList.of( + _ -> firstBlocked, + _ -> Futures.immediateVoidFuture()))); + + executor.scheduleMoreLeafSplits(); + firstBlocked.awaitListenerAdded(); + + TaskEntry secondTask = (TaskEntry) executor.addTask(new TaskId(new StageId("query", 1), 2, 1), () -> 0, 10, new Duration(1, MILLISECONDS), OptionalInt.empty()); + CountDownLatch secondStarted = new CountDownLatch(1); + ListenableFuture secondDone = secondTask.enqueueLeafSplit(new TestingSplitRunner(ImmutableList.of(_ -> { + secondStarted.countDown(); + return Futures.immediateVoidFuture(); + }))); + + firstBlocked.set(null); + firstDone.get(10, TimeUnit.SECONDS); + assertThat(secondStarted.await(10, TimeUnit.SECONDS)).isTrue(); + secondDone.get(10, TimeUnit.SECONDS); + } + finally { + executor.stop(); + } + } + + @Test + public void testDestroyCompletesEveryLeafFutureWhenCloseFails() + throws Exception + { + FairScheduler scheduler = FairScheduler.newInstance(1); + TaskEntry task = new TaskEntry( + new TaskId(new StageId("query", 1), 1, 1), + scheduler, + testingVersionEmbedder(), + noopTracer(), + 1, + () -> 0); + TestingSplitRunner failing = new CloseFailingSplitRunner(); + TestingSplitRunner pending = new TestingSplitRunner(ImmutableList.of(_ -> Futures.immediateVoidFuture())); + ListenableFuture claimedDone = task.enqueueLeafSplit(failing); + ListenableFuture pendingDone = task.enqueueLeafSplit(pending); + task.claimLeafSplit(); + + try { + assertThatThrownBy(task::destroy) + .isInstanceOf(RuntimeException.class) + .hasMessage("close failed"); + + claimedDone.get(); + pendingDone.get(); + assertThat(failing.isFinished()).isTrue(); + assertThat(pending.isFinished()).isTrue(); + + TestingSplitRunner afterDestroy = new TestingSplitRunner(ImmutableList.of(_ -> Futures.immediateVoidFuture())); + task.enqueueLeafSplit(afterDestroy).get(); + assertThat(afterDestroy.isFinished()).isTrue(); + } + finally { + scheduler.close(); + } + } + + @Test + public void testMaintenanceSurvivesFailure() + { + // scheduleWithFixedDelay stops rescheduling a task that throws, so the wrapper must swallow + // the failure. It has to catch Error too: the failure an overloaded worker produces is + // OutOfMemoryError from creating a split's thread. + AtomicInteger runs = new AtomicInteger(); + Runnable task = ThreadPerDriverTaskExecutor.maintenance( + () -> { + if (runs.incrementAndGet() == 1) { + throw new OutOfMemoryError("unable to create native thread"); + } + }, + "Error in test task"); + + task.run(); + task.run(); + + assertThat(runs.get()).isEqualTo(2); + } + @Test @Timeout(10) public void testCancellationWhileProcessing() @@ -244,7 +581,7 @@ public final String getInfo() } @Override - public final void close() + public void close() { finished = true; @@ -255,4 +592,20 @@ public final void close() } } } + + private static class CloseFailingSplitRunner + extends TestingSplitRunner + { + public CloseFailingSplitRunner() + { + super(ImmutableList.of(_ -> Futures.immediateVoidFuture())); + } + + @Override + public void close() + { + super.close(); + throw new RuntimeException("close failed"); + } + } } diff --git a/core/trino-main/src/test/java/io/trino/execution/executor/scheduler/TestFairScheduler.java b/core/trino-main/src/test/java/io/trino/execution/executor/scheduler/TestFairScheduler.java index bf6810618f4a..82f6fe3a6ec7 100644 --- a/core/trino-main/src/test/java/io/trino/execution/executor/scheduler/TestFairScheduler.java +++ b/core/trino-main/src/test/java/io/trino/execution/executor/scheduler/TestFairScheduler.java @@ -13,6 +13,7 @@ */ package io.trino.execution.executor.scheduler; +import com.google.common.base.Ticker; import com.google.common.util.concurrent.AbstractFuture; import com.google.common.util.concurrent.ListenableFuture; import com.google.common.util.concurrent.SettableFuture; @@ -23,14 +24,119 @@ import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Future; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.ThreadFactory; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; +import static com.google.common.util.concurrent.Uninterruptibles.awaitUninterruptibly; +import static io.airlift.concurrent.Threads.daemonThreadsNamed; +import static java.util.concurrent.Executors.newSingleThreadExecutor; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; public class TestFairScheduler { + @Test + @Timeout(30) + public void testSubmitDoesNotBlockOtherOperations() + throws Exception + { + // Creating the thread that runs a task is slow on a loaded worker. Simulate that and verify + // that it does not hold up group management or the submission of unrelated tasks. + AtomicBoolean stallNextThread = new AtomicBoolean(); + CountDownLatch threadCreationStarted = new CountDownLatch(1); + CountDownLatch releaseThreadCreation = new CountDownLatch(1); + + ThreadFactory threadFactory = runnable -> { + if (stallNextThread.compareAndSet(true, false)) { + threadCreationStarted.countDown(); + awaitUninterruptibly(releaseThreadCreation); + } + Thread thread = new Thread(runnable); + thread.setDaemon(true); + return thread; + }; + + FairScheduler scheduler = new FairScheduler(2, threadFactory, Ticker.systemTicker()); + scheduler.start(); + ExecutorService submitter = newSingleThreadExecutor(daemonThreadsNamed("submitter")); + try { + Group group = scheduler.createGroup("G1"); + + stallNextThread.set(true); + Future stalled = submitter.submit(() -> scheduler.submit(group, 1, _ -> {})); + threadCreationStarted.await(); + + // these must not wait for the stalled thread creation to complete + Group other = scheduler.createGroup("G2"); + scheduler.submit(other, 1, _ -> {}).get(); + scheduler.removeGroup(other); + + releaseThreadCreation.countDown(); + stalled.get(); + } + finally { + releaseThreadCreation.countDown(); + submitter.shutdownNow(); + scheduler.close(); + } + } + + @Test + @Timeout(30) + public void testCloseDoesNotWaitForThreadCreation() + throws Exception + { + CountDownLatch threadCreationStarted = new CountDownLatch(1); + CountDownLatch releaseThreadCreation = new CountDownLatch(1); + ThreadFactory threadFactory = runnable -> { + threadCreationStarted.countDown(); + awaitUninterruptibly(releaseThreadCreation); + Thread thread = new Thread(runnable); + thread.setDaemon(true); + return thread; + }; + + FairScheduler scheduler = new FairScheduler(1, threadFactory, Ticker.systemTicker()); + scheduler.start(); + Group group = scheduler.createGroup("G"); + ExecutorService executor = newSingleThreadExecutor(daemonThreadsNamed("submitter")); + try { + Future submit = executor.submit(() -> scheduler.submit(group, 1, _ -> {})); + threadCreationStarted.await(); + + // Shutdown is safe without waiting for the submit call to finish creating its thread. + scheduler.close(); + + releaseThreadCreation.countDown(); + assertThatThrownBy(submit::get) + .isInstanceOf(ExecutionException.class) + .cause() + .isInstanceOf(RejectedExecutionException.class); + } + finally { + releaseThreadCreation.countDown(); + executor.shutdownNow(); + scheduler.close(); + } + } + + @Test + public void testSubmitAfterClose() + { + FairScheduler scheduler = FairScheduler.newInstance(1); + Group group = scheduler.createGroup("G"); + scheduler.close(); + + assertThatThrownBy(() -> scheduler.submit(group, 1, _ -> {})) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("Already closed"); + } + @Test public void testBasic() throws ExecutionException, InterruptedException diff --git a/core/trino-main/src/test/java/io/trino/operator/TestAssignUniqueIdOperator.java b/core/trino-main/src/test/java/io/trino/operator/TestAssignUniqueIdOperator.java new file mode 100644 index 000000000000..d5f97981a154 --- /dev/null +++ b/core/trino-main/src/test/java/io/trino/operator/TestAssignUniqueIdOperator.java @@ -0,0 +1,133 @@ +/* + * Licensed 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 io.trino.operator; + +import io.trino.RowPagesBuilder; +import io.trino.spi.Page; +import io.trino.spi.block.Block; +import io.trino.sql.planner.plan.PlanNodeId; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestInstance; +import org.junit.jupiter.api.parallel.Execution; + +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.atomic.AtomicLong; + +import static io.airlift.concurrent.Threads.daemonThreadsNamed; +import static io.trino.RowPagesBuilder.rowPagesBuilder; +import static io.trino.SessionTestUtils.TEST_SESSION; +import static io.trino.operator.OperatorAssertion.toPages; +import static io.trino.spi.type.BigintType.BIGINT; +import static io.trino.testing.TestingTaskContext.createTaskContext; +import static java.util.concurrent.Executors.newCachedThreadPool; +import static java.util.concurrent.Executors.newScheduledThreadPool; +import static org.assertj.core.api.Assertions.assertThat; +import static org.junit.jupiter.api.TestInstance.Lifecycle.PER_CLASS; +import static org.junit.jupiter.api.parallel.ExecutionMode.CONCURRENT; + +@TestInstance(PER_CLASS) +@Execution(CONCURRENT) +final class TestAssignUniqueIdOperator +{ + private static final int ROWS_PER_PAGE = 1000; + private static final int PAGE_COUNT = 10; + private static final int TOTAL_ROWS = ROWS_PER_PAGE * PAGE_COUNT; + + private final ExecutorService executor = newCachedThreadPool(daemonThreadsNamed(getClass().getSimpleName() + "-%s")); + private final ScheduledExecutorService scheduledExecutor = newScheduledThreadPool(2, daemonThreadsNamed(getClass().getSimpleName() + "-scheduledExecutor-%s")); + + @AfterAll + void tearDown() + { + executor.shutdownNow(); + scheduledExecutor.shutdownNow(); + } + + @Test + void testAssignUniqueIds() + { + OperatorFactory operatorFactory = AssignUniqueIdOperator.createOperatorFactory(0, new PlanNodeId("test"), new AtomicLong()); + TaskContext taskContext = createTaskContext(executor, scheduledExecutor, TEST_SESSION); + + Set ids = new HashSet<>(); + collectIds(operatorFactory, newDriverContext(taskContext, 0), ids); + + assertThat(ids).hasSize(TOTAL_ROWS); + } + + @Test + void testFactoriesSharingValuePoolProduceDistinctIds() + { + // Two AssignUniqueId nodes executing in the same task (e.g. the MERGE target and source + // sides of a colocated join) must not produce the same id for different rows + AtomicLong valuePool = new AtomicLong(); + OperatorFactory targetFactory = AssignUniqueIdOperator.createOperatorFactory(0, new PlanNodeId("target"), valuePool); + OperatorFactory sourceFactory = AssignUniqueIdOperator.createOperatorFactory(1, new PlanNodeId("source"), valuePool); + + // Both operators run in the same task, so the id mask is identical for the two sides + // and only the shared pool keeps their ids apart + TaskContext taskContext = createTaskContext(executor, scheduledExecutor, TEST_SESSION); + + Set ids = new HashSet<>(); + collectIds(targetFactory, newDriverContext(taskContext, 0), ids); + collectIds(sourceFactory, newDriverContext(taskContext, 1), ids); + + assertThat(ids).hasSize(2 * TOTAL_ROWS); + } + + @Test + void testDuplicatedFactoriesProduceDistinctIds() + { + // Factories duplicated for additional pipelines (e.g. lookup outer drivers) share the pool + OperatorFactory operatorFactory = AssignUniqueIdOperator.createOperatorFactory(0, new PlanNodeId("test"), new AtomicLong()); + OperatorFactory duplicateFactory = operatorFactory.duplicate(); + TaskContext taskContext = createTaskContext(executor, scheduledExecutor, TEST_SESSION); + + Set ids = new HashSet<>(); + collectIds(operatorFactory, newDriverContext(taskContext, 0), ids); + collectIds(duplicateFactory, newDriverContext(taskContext, 1), ids); + + assertThat(ids).hasSize(2 * TOTAL_ROWS); + } + + private static void collectIds(OperatorFactory operatorFactory, DriverContext driverContext, Set ids) + { + RowPagesBuilder pagesBuilder = rowPagesBuilder(BIGINT); + for (int page = 0; page < PAGE_COUNT; page++) { + pagesBuilder.addSequencePage(ROWS_PER_PAGE, 0); + } + List input = pagesBuilder.build(); + + List output = toPages(operatorFactory, driverContext, input); + for (Page page : output) { + assertThat(page.getChannelCount()).isEqualTo(2); + Block idBlock = page.getBlock(1); + for (int position = 0; position < page.getPositionCount(); position++) { + assertThat(ids.add(BIGINT.getLong(idBlock, position))).isTrue(); + } + } + } + + private static DriverContext newDriverContext(TaskContext taskContext, int pipelineId) + { + return taskContext + .addPipelineContext(pipelineId, true, true, false) + .addDriverContext(); + } +} diff --git a/core/trino-spi/src/main/java/io/trino/spi/connector/JoinCondition.java b/core/trino-spi/src/main/java/io/trino/spi/connector/JoinCondition.java index 3d303f64cfec..fa28dbac2308 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/connector/JoinCondition.java +++ b/core/trino-spi/src/main/java/io/trino/spi/connector/JoinCondition.java @@ -46,7 +46,7 @@ public enum Operator LESS_THAN_OR_EQUAL("<=", StandardFunctions.LESS_THAN_OR_EQUAL_OPERATOR_FUNCTION_NAME), GREATER_THAN(">", StandardFunctions.GREATER_THAN_OPERATOR_FUNCTION_NAME), GREATER_THAN_OR_EQUAL(">=", StandardFunctions.GREATER_THAN_OR_EQUAL_OPERATOR_FUNCTION_NAME), - IDENTICAL("≡", StandardFunctions.IDENTICAL_OPERATOR_FUNCTION_NAME), + IDENTICAL("IS NOT DISTINCT FROM", StandardFunctions.IDENTICAL_OPERATOR_FUNCTION_NAME), /**/; private static final Map byFunctionName = Stream.of(values()) diff --git a/lib/trino-record-decoder/src/main/java/io/trino/decoder/csv/CsvColumnDecoder.java b/lib/trino-record-decoder/src/main/java/io/trino/decoder/csv/CsvColumnDecoder.java index ed7f90f3da90..a106cc5ce541 100644 --- a/lib/trino-record-decoder/src/main/java/io/trino/decoder/csv/CsvColumnDecoder.java +++ b/lib/trino-record-decoder/src/main/java/io/trino/decoder/csv/CsvColumnDecoder.java @@ -73,10 +73,7 @@ private static boolean isSupportedType(Type type) if (type instanceof VarcharType) { return true; } - if (ImmutableList.of(BIGINT, INTEGER, SMALLINT, TINYINT, BOOLEAN, DOUBLE).contains(type)) { - return true; - } - return false; + return ImmutableList.of(BIGINT, INTEGER, SMALLINT, TINYINT, BOOLEAN, DOUBLE).contains(type); } public FieldValueProvider decodeField(String[] tokens) diff --git a/lib/trino-record-decoder/src/main/java/io/trino/decoder/raw/RawColumnDecoder.java b/lib/trino-record-decoder/src/main/java/io/trino/decoder/raw/RawColumnDecoder.java index f0fbe93bdb25..fc3e22d99eb5 100644 --- a/lib/trino-record-decoder/src/main/java/io/trino/decoder/raw/RawColumnDecoder.java +++ b/lib/trino-record-decoder/src/main/java/io/trino/decoder/raw/RawColumnDecoder.java @@ -164,10 +164,7 @@ private static boolean isSupportedType(Type type) if (type instanceof VarcharType) { return true; } - if (ImmutableList.of(BIGINT, INTEGER, SMALLINT, TINYINT, BOOLEAN, DOUBLE).contains(type)) { - return true; - } - return false; + return ImmutableList.of(BIGINT, INTEGER, SMALLINT, TINYINT, BOOLEAN, DOUBLE).contains(type); } private void checkFieldTypeOneOf(FieldType declaredFieldType, String columnName, FieldType... allowedFieldTypes) diff --git a/plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/BaseJdbcConnectorTest.java b/plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/BaseJdbcConnectorTest.java index 36ec79e4c2c6..c74b681312ee 100644 --- a/plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/BaseJdbcConnectorTest.java +++ b/plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/BaseJdbcConnectorTest.java @@ -1189,9 +1189,9 @@ public void testJoinPushdown() List nonEqualities = Stream.concat( Stream.of(JoinCondition.Operator.values()) - .filter(operator -> operator != JoinCondition.Operator.EQUAL && operator != JoinCondition.Operator.IDENTICAL) + .filter(operator -> operator != JoinCondition.Operator.EQUAL) .map(JoinCondition.Operator::getValue), - Stream.of("IS DISTINCT FROM", "IS NOT DISTINCT FROM")) + Stream.of("IS DISTINCT FROM")) .collect(toImmutableList()); // basic case @@ -1413,9 +1413,6 @@ private boolean expectVarcharJoinPushdown(String operator) private JoinCondition.Operator toJoinConditionOperator(String operator) { - if (operator.equals("IS NOT DISTINCT FROM")) { - return JoinCondition.Operator.IDENTICAL; - } return Stream.of(JoinCondition.Operator.values()) .filter(joinOperator -> joinOperator.getValue().equals(operator)) .collect(toOptional()) diff --git a/plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/TestDefaultJdbcQueryBuilder.java b/plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/TestDefaultJdbcQueryBuilder.java index 3e6c78b9d38f..96b1bef8e4b3 100644 --- a/plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/TestDefaultJdbcQueryBuilder.java +++ b/plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/TestDefaultJdbcQueryBuilder.java @@ -543,6 +543,13 @@ public void testBuildJoinSql() @Test public void testBuildJoinSqlLegacy() throws SQLException + { + testBuildJoinSqlLegacy(JoinCondition.Operator.EQUAL, "="); + testBuildJoinSqlLegacy(JoinCondition.Operator.IDENTICAL, "IS NOT DISTINCT FROM"); + } + + private void testBuildJoinSqlLegacy(JoinCondition.Operator joinOperator, String sqlOperator) + throws SQLException { Connection connection = database.getConnection(); @@ -553,7 +560,7 @@ public void testBuildJoinSqlLegacy() JoinType.INNER, new PreparedQuery("SELECT * FROM \"test_table\"", List.of()), new PreparedQuery("SELECT * FROM \"test_table\"", List.of()), - List.of(new JdbcJoinCondition(columns.get(7), JoinCondition.Operator.EQUAL, columns.get(8))), + List.of(new JdbcJoinCondition(columns.get(7), joinOperator, columns.get(8))), Map.of(columns.get(2), "name1"), Map.of(columns.get(3), "name2")); try (PreparedStatement preparedStatement = queryBuilder.prepareStatement(jdbcClient, SESSION, connection, preparedQuery, Optional.empty())) { @@ -562,7 +569,7 @@ public void testBuildJoinSqlLegacy() "(SELECT * FROM \"test_table\") l " + "INNER JOIN " + "(SELECT * FROM \"test_table\") r " + - "ON l.\"col_7\" = r.\"col_8\""); + "ON l.\"col_7\" " + sqlOperator + " r.\"col_8\""); long count = 0; try (ResultSet resultSet = preparedStatement.executeQuery()) { while (resultSet.next()) { diff --git a/plugin/trino-blob-cache-alluxio/src/test/java/io/trino/blob/cache/alluxio/TestAlluxioCacheFileSystemAccessOperations.java b/plugin/trino-blob-cache-alluxio/src/test/java/io/trino/blob/cache/alluxio/TestAlluxioCacheFileSystemAccessOperations.java index 430851b2f862..c4b1fa9db5f5 100644 --- a/plugin/trino-blob-cache-alluxio/src/test/java/io/trino/blob/cache/alluxio/TestAlluxioCacheFileSystemAccessOperations.java +++ b/plugin/trino-blob-cache-alluxio/src/test/java/io/trino/blob/cache/alluxio/TestAlluxioCacheFileSystemAccessOperations.java @@ -135,6 +135,8 @@ public void testCache() int readTimes = 3; assertCacheOperations(0, location, content, readTimes, ImmutableMultiset.builder() + .add(new CacheOperationSpan("InputFile.length", location.toString())) + .add(new CacheOperationSpan("InputFile.newInput", location.toString())) .addCopies(new CacheOperationSpan("Alluxio.readCached", location.toString(), 11), readTimes) .addCopies(new CacheOperationSpan("AlluxioCacheManager.get", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 0, 11), readTimes) .add(new CacheOperationSpan("AlluxioCacheManager.put", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 0, 11)) @@ -151,6 +153,8 @@ public void testCache() readTimes = 7; assertCacheOperations(0, location, modifiedContent, readTimes, ImmutableMultiset.builder() + .add(new CacheOperationSpan("InputFile.length", location.toString())) + .add(new CacheOperationSpan("InputFile.newInput", location.toString())) .add(new CacheOperationSpan("Input.readFully", location.toString(), 16)) .add(new CacheOperationSpan("Alluxio.writeCache", location.toString(), 16)) .add(new CacheOperationSpan("AlluxioCacheManager.put", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 0, 16)) @@ -174,6 +178,8 @@ public void testPartialCacheHits() assertCacheOperations(location, Arrays.copyOf(content, PAGE_SIZE), ImmutableMultiset.builder() + .add(new CacheOperationSpan("InputFile.length", location.toString())) + .add(new CacheOperationSpan("InputFile.newInput", location.toString())) .add(new CacheOperationSpan("Alluxio.readCached", "memory:///partial", 0, PAGE_SIZE)) .add(new CacheOperationSpan("AlluxioCacheManager.get", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 0, PAGE_SIZE)) .add(new CacheOperationSpan("Input.readFully", location.toString(), 0, PAGE_SIZE)) @@ -183,6 +189,8 @@ public void testPartialCacheHits() assertCacheOperations(location, Arrays.copyOf(content, PAGE_SIZE + 10), ImmutableMultiset.builder() + .add(new CacheOperationSpan("InputFile.length", location.toString())) + .add(new CacheOperationSpan("InputFile.newInput", location.toString())) .add(new CacheOperationSpan("Alluxio.readCached", location.toString(), 0, PAGE_SIZE + 10)) .add(new CacheOperationSpan("AlluxioCacheManager.get", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 0, PAGE_SIZE)) .add(new CacheOperationSpan("AlluxioCacheManager.get", cacheKey(location, cacheKeyProvider.currentCacheVersion()), PAGE_SIZE, 10)) @@ -193,6 +201,7 @@ public void testPartialCacheHits() assertCacheOperations(location, Arrays.copyOf(content, PAGE_SIZE + 10), ImmutableMultiset.builder() + .add(new CacheOperationSpan("InputFile.length", location.toString())) .add(new CacheOperationSpan("AlluxioCacheManager.get", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 0, PAGE_SIZE)) .add(new CacheOperationSpan("AlluxioCacheManager.get", cacheKey(location, cacheKeyProvider.currentCacheVersion()), PAGE_SIZE, 10)) .add(new CacheOperationSpan("Alluxio.readCached", location.toString(), PAGE_SIZE + 10)) @@ -200,6 +209,7 @@ public void testPartialCacheHits() assertCacheOperations(location, content, ImmutableMultiset.builder() + .add(new CacheOperationSpan("InputFile.length", location.toString())) .add(new CacheOperationSpan("AlluxioCacheManager.get", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 0, PAGE_SIZE)) .add(new CacheOperationSpan("AlluxioCacheManager.get", cacheKey(location, cacheKeyProvider.currentCacheVersion()), PAGE_SIZE, PAGE_SIZE)) .add(new CacheOperationSpan("Alluxio.readCached", location.toString(), 0, PAGE_SIZE * 2)) @@ -221,6 +231,8 @@ public void testMultiPageExternalsReads() assertCacheOperations(location, Arrays.copyOf(content, PAGE_SIZE + 1), ImmutableMultiset.builder() + .add(new CacheOperationSpan("InputFile.length", location.toString())) + .add(new CacheOperationSpan("InputFile.newInput", location.toString())) .add(new CacheOperationSpan("AlluxioCacheManager.get", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 0, PAGE_SIZE)) .add(new CacheOperationSpan("AlluxioCacheManager.put", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 0, PAGE_SIZE)) .add(new CacheOperationSpan("AlluxioCacheManager.put", cacheKey(location, cacheKeyProvider.currentCacheVersion()), PAGE_SIZE, PAGE_SIZE)) @@ -231,6 +243,8 @@ public void testMultiPageExternalsReads() cacheKeyProvider.increaseCacheVersion(); assertCacheOperations(location, Arrays.copyOf(content, 2 * PAGE_SIZE), ImmutableMultiset.builder() + .add(new CacheOperationSpan("InputFile.length", location.toString())) + .add(new CacheOperationSpan("InputFile.newInput", location.toString())) .add(new CacheOperationSpan("AlluxioCacheManager.get", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 0, PAGE_SIZE)) .add(new CacheOperationSpan("AlluxioCacheManager.put", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 0, PAGE_SIZE)) .add(new CacheOperationSpan("AlluxioCacheManager.put", cacheKey(location, cacheKeyProvider.currentCacheVersion()), PAGE_SIZE, PAGE_SIZE)) @@ -283,6 +297,8 @@ public void testCacheWithMissingPage() int readTimes = 3; assertCacheOperations(0, location, content, readTimes, ImmutableMultiset.builder() + .add(new CacheOperationSpan("InputFile.length", location.toString())) + .add(new CacheOperationSpan("InputFile.newInput", location.toString())) .addCopies(new CacheOperationSpan("Alluxio.readCached", location.toString(), 12), readTimes) .addCopies(new CacheOperationSpan("AlluxioCacheManager.get", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 0, 12), readTimes) .add(new CacheOperationSpan("AlluxioCacheManager.put", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 0, 12)) @@ -298,6 +314,8 @@ public void testCacheWithMissingPage() assertCacheOperations(location, content, ImmutableMultiset.builder() + .add(new CacheOperationSpan("InputFile.length", location.toString())) + .add(new CacheOperationSpan("InputFile.newInput", location.toString())) .add(new CacheOperationSpan("Alluxio.readCached", location.toString(), 12)) .add(new CacheOperationSpan("AlluxioCacheManager.get", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 12)) .add(new CacheOperationSpan("Input.readFully", location.toString(), 12)) @@ -320,6 +338,8 @@ public void testCacheWithCorruptedPage() int readTimes = 3; assertCacheOperations(0, location, content, readTimes, ImmutableMultiset.builder() + .add(new CacheOperationSpan("InputFile.length", location.toString())) + .add(new CacheOperationSpan("InputFile.newInput", location.toString())) .addCopies(new CacheOperationSpan("Alluxio.readCached", location.toString(), 14), readTimes) .addCopies(new CacheOperationSpan("AlluxioCacheManager.get", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 0, 14), readTimes) .add(new CacheOperationSpan("AlluxioCacheManager.put", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 0, 14)) @@ -335,6 +355,8 @@ public void testCacheWithCorruptedPage() assertCacheOperations(location, content, ImmutableMultiset.builder() + .add(new CacheOperationSpan("InputFile.length", location.toString())) + .add(new CacheOperationSpan("InputFile.newInput", location.toString())) .add(new CacheOperationSpan("Alluxio.readCached", location.toString(), 14)) .add(new CacheOperationSpan("AlluxioCacheManager.get", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 14)) .add(new CacheOperationSpan("Input.readFully", location.toString(), 14)) @@ -357,6 +379,8 @@ public void testCacheHitAfterReadFromNoneZeroPosition() int readTimes = 5; assertCacheOperations(8, location, readContent, readTimes, ImmutableMultiset.builder() + .add(new CacheOperationSpan("InputFile.length", location.toString())) + .add(new CacheOperationSpan("InputFile.newInput", location.toString())) .addCopies(new CacheOperationSpan("Alluxio.readCached", location.toString(), 8, 2), readTimes) .addCopies(new CacheOperationSpan("AlluxioCacheManager.get", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 8, 2), readTimes) .add(new CacheOperationSpan("Input.readFully", location.toString(), 0, 11)) @@ -384,6 +408,7 @@ private void assertCachedRead(Location location, int fileSize) throws IOException { ImmutableMultiset.Builder builder = ImmutableMultiset.builder() + .add(new CacheOperationSpan("InputFile.length", location.toString())) .add(new CacheOperationSpan("Alluxio.readCached", location.toString(), 0, fileSize)); for (int offset = 0; offset < fileSize; offset = offset + PAGE_SIZE) { @@ -397,6 +422,8 @@ private void assertUnCachedRead(Location location, int fileSize) throws IOException { ImmutableMultiset.Builder builder = ImmutableMultiset.builder() + .add(new CacheOperationSpan("InputFile.length", location.toString())) + .add(new CacheOperationSpan("InputFile.newInput", location.toString())) .add(new CacheOperationSpan("Alluxio.readCached", location.toString(), fileSize)) .add(new CacheOperationSpan("Alluxio.writeCache", location.toString(), fileSize)) .add(new CacheOperationSpan("Input.readFully", location.toString(), fileSize)); @@ -446,20 +473,30 @@ private void assertCacheOperations(int position, Location location, byte[] conte private Multiset getCacheOperations(List spans) { return spans.stream() - .filter(span -> span.getName().startsWith("Input.") || span.getName().startsWith("Alluxio")) + .filter(span -> span.getName().startsWith("Input.") || span.getName().startsWith("InputFile.") || span.getName().startsWith("Alluxio")) .map(CacheOperationSpan::create) .collect(toCollection(HashMultiset::create)); } private record CacheOperationSpan(String spanName, String location, long position, long length) { - public CacheOperationSpan(String spanName, String location, long length) + CacheOperationSpan(String spanName, String location, long length) { this(spanName, location, 0, length); } - public static CacheOperationSpan create(SpanData span) + CacheOperationSpan(String spanName, String location) { + this(spanName, location, 0, 0); + } + + static CacheOperationSpan create(SpanData span) + { + // Delegate calls made by the cache itself, which cost a round trip to remote storage + if (span.getName().startsWith("InputFile.")) { + return new CacheOperationSpan(span.getName(), getLocation(span)); + } + Attributes attributes = span.getAttributes(); long length = switch (span.getName()) { @@ -493,7 +530,7 @@ private static String cacheKey(Location location, int cacheVersion) private static String getLocation(SpanData span) { - if (span.getName().startsWith("Input.")) { + if (span.getName().startsWith("Input.") || span.getName().startsWith("InputFile.")) { return requireNonNull(span.getAttributes().get(FILE_LOCATION)); } return requireNonNullElse(span.getAttributes().get(CACHE_FILE_LOCATION), span.getAttributes().get(CACHE_KEY)); diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergConnectorTest.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergConnectorTest.java index b1b7c5f84fbe..b10014f560d4 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergConnectorTest.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergConnectorTest.java @@ -53,8 +53,10 @@ import io.trino.spi.statistics.TableStatistics; import io.trino.sql.planner.Plan; import io.trino.sql.planner.optimizations.PlanNodeSearcher; +import io.trino.sql.planner.plan.AssignUniqueId; import io.trino.sql.planner.plan.ExchangeNode; import io.trino.sql.planner.plan.FilterNode; +import io.trino.sql.planner.plan.JoinNode; import io.trino.sql.planner.plan.OutputNode; import io.trino.sql.planner.plan.TableScanNode; import io.trino.sql.planner.plan.TableWriterNode; @@ -126,6 +128,7 @@ import static io.trino.SystemSessionProperties.ENABLE_DYNAMIC_FILTERING; import static io.trino.SystemSessionProperties.IGNORE_STATS_CALCULATOR_FAILURES; import static io.trino.SystemSessionProperties.ITERATIVE_OPTIMIZER_TIMEOUT; +import static io.trino.SystemSessionProperties.JOIN_DISTRIBUTION_TYPE; import static io.trino.SystemSessionProperties.MAX_HASH_PARTITION_COUNT; import static io.trino.SystemSessionProperties.MAX_WRITER_TASK_COUNT; import static io.trino.SystemSessionProperties.SCALE_WRITERS; @@ -162,6 +165,7 @@ import static io.trino.spi.type.VarcharType.VARCHAR; import static io.trino.sql.planner.assertions.PlanMatchPattern.node; import static io.trino.sql.planner.optimizations.PlanNodeSearcher.searchFrom; +import static io.trino.sql.planner.plan.ExchangeNode.Scope.REMOTE; import static io.trino.testing.MaterializedResult.resultBuilder; import static io.trino.testing.QueryAssertions.assertEqualsIgnoreOrder; import static io.trino.testing.TestingConnectorSession.SESSION; @@ -1284,6 +1288,55 @@ public void testMergeWithNestedFieldPartitionedTable() assertUpdate("DROP TABLE " + targetTable); } + @Test // regression test for https://github.com/trinodb/trino/issues/30639 + public void testMergeWithPartitionedJoinAndUnmodifiedRows() + { + // The join with a bucket-partitioned target is colocated with the table partitioning, so + // the target and source AssignUniqueId nodes execute in the same stage and must not + // assign the same id to different rows + try (TestTable target = newTrinoTable( + "test_merge_unique_id_target_", + "WITH (partitioning = ARRAY['bucket(k1, 16)', 'bucket(k2, 8)']) AS " + + "SELECT 'a-' || CAST(i AS varchar) AS k1, 'b-' || CAST(i AS varchar) AS k2, " + + "'UPDATED' AS value, TIMESTAMP '2026-01-01 00:00:00.000000' AS updated_at " + + "FROM UNNEST(sequence(1, 20)) t(i)"); + // Every fifth source row matches a target row, so matched and unmatched rows are + // interleaved and each AssignUniqueId driver assigns its first ids to both kinds + TestTable source = newTrinoTable( + "test_merge_unique_id_source_", + "AS SELECT " + + "IF(i % 5 = 0, 'a-' || CAST(i / 5 AS varchar), 'x-' || CAST(i AS varchar)) AS k1, " + + "IF(i % 5 = 0, 'b-' || CAST(i / 5 AS varchar), 'y-' || CAST(i AS varchar)) AS k2, " + + "'DELETED' AS value, TIMESTAMP '2026-01-01 00:00:00.000000' AS updated_at " + + "FROM UNNEST(sequence(1, 100)) t(i)")) { + Session session = Session.builder(getSession()) + .setSystemProperty(JOIN_DISTRIBUTION_TYPE, "PARTITIONED") + .setCatalogSessionProperty(ICEBERG_CATALOG, BUCKET_EXECUTION_ENABLED, "true") + .build(); + + // Every row is excluded by a WHEN condition, so the merge must change nothing + // instead of failing with MERGE_TARGET_ROW_MULTIPLE_MATCHES + assertUpdate( + session, + "MERGE INTO %s t USING %s s ".formatted(target.getName(), source.getName()) + + "ON t.k1 = s.k1 AND t.k2 = s.k2 " + + "WHEN MATCHED AND s.updated_at > t.updated_at THEN UPDATE SET value = s.value, updated_at = s.updated_at " + + "WHEN NOT MATCHED AND s.value <> 'DELETED' THEN INSERT (k1, k2, value, updated_at) VALUES (s.k1, s.k2, s.value, s.updated_at)", + 0, + // Both AssignUniqueId nodes must be reachable from the join without crossing a + // remote exchange, otherwise the test no longer exercises a shared task + plan -> { + JoinNode join = (JoinNode) searchFrom(plan.getRoot()).where(JoinNode.class::isInstance).findOnlyElement(); + assertThat(searchFrom(join) + .recurseOnlyWhen(planNode -> !(planNode instanceof ExchangeNode exchange && exchange.getScope() == REMOTE)) + .where(AssignUniqueId.class::isInstance) + .count()) + .isEqualTo(2); + }); + assertQuery("SELECT count(*) FROM " + target.getName(), "VALUES 20"); + } + } + @Test public void testSchemaEvolutionWithNestedFieldPartitioning() { diff --git a/plugin/trino-postgresql/src/test/java/io/trino/plugin/postgresql/TestPostgreSqlConnectorTest.java b/plugin/trino-postgresql/src/test/java/io/trino/plugin/postgresql/TestPostgreSqlConnectorTest.java index c169db01b999..b751080acdc2 100644 --- a/plugin/trino-postgresql/src/test/java/io/trino/plugin/postgresql/TestPostgreSqlConnectorTest.java +++ b/plugin/trino-postgresql/src/test/java/io/trino/plugin/postgresql/TestPostgreSqlConnectorTest.java @@ -61,6 +61,7 @@ import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.collect.MoreCollectors.onlyElement; import static io.airlift.slice.Slices.utf8Slice; +import static io.trino.plugin.jdbc.JdbcMetadataSessionProperties.COMPLEX_JOIN_PUSHDOWN_ENABLED; import static io.trino.plugin.postgresql.PostgreSqlConfig.ArrayMapping.AS_ARRAY; import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.spi.type.VarcharType.VARCHAR; @@ -618,9 +619,9 @@ public void testStringJoinPushdownWithCollate() List nonEqualities = Stream.concat( Stream.of(JoinCondition.Operator.values()) - .filter(operator -> operator != JoinCondition.Operator.EQUAL && operator != JoinCondition.Operator.IDENTICAL) + .filter(operator -> operator != JoinCondition.Operator.EQUAL) .map(JoinCondition.Operator::getValue), - Stream.of("IS DISTINCT FROM", "IS NOT DISTINCT FROM")) + Stream.of("IS DISTINCT FROM")) .collect(toImmutableList()); try (TestTable nationLowercaseTable = newTrinoTable( @@ -747,6 +748,17 @@ public void testStringJoinPushdownWithCollate() } } + @Test + public void testLegacyJoinPushdownWithIdenticalCondition() + { + Session session = Session.builder(joinPushdownEnabled(getSession())) + .setCatalogSessionProperty("postgresql", COMPLEX_JOIN_PUSHDOWN_ENABLED, "false") + .build(); + + assertThat(query(session, "SELECT n1.name FROM nation n1 JOIN nation n2 ON n1.nationkey = n2.nationkey AND n1.regionkey IS NOT DISTINCT FROM n2.regionkey")) + .isFullyPushedDown(); + } + @Test public void testDecimalPredicatePushdown() { diff --git a/testing/trino-testing-containers/src/main/java/io/trino/testing/containers/KeyManagementServer.java b/testing/trino-testing-containers/src/main/java/io/trino/testing/containers/KeyManagementServer.java index be5c568a86bc..0a2549075bd2 100644 --- a/testing/trino-testing-containers/src/main/java/io/trino/testing/containers/KeyManagementServer.java +++ b/testing/trino-testing-containers/src/main/java/io/trino/testing/containers/KeyManagementServer.java @@ -27,7 +27,7 @@ public class KeyManagementServer { public static final int KES_PORT = 7373; private static final Logger log = Logger.get(KeyManagementServer.class); - private static final String DEFAULT_IMAGE = "minio/kes:2024-06-17T15-47-05Z"; + private static final String DEFAULT_IMAGE = "quay.io/minio/kes:2024-06-17T15-47-05Z"; private static final String DEFAULT_HOST_NAME = "kes"; public static KeyManagementServer.Builder builder()