io.trino.execution.buffer.LazyOutputBuffer Maven / Gradle / Ivy
/*
* 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.execution.buffer;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.SettableFuture;
import com.google.errorprone.annotations.concurrent.GuardedBy;
import io.airlift.slice.Slice;
import io.airlift.units.DataSize;
import io.trino.exchange.ExchangeManagerRegistry;
import io.trino.execution.StateMachine.StateChangeListener;
import io.trino.execution.TaskId;
import io.trino.execution.buffer.PipelinedOutputBuffers.OutputBufferId;
import io.trino.memory.context.LocalMemoryContext;
import io.trino.spi.exchange.ExchangeManager;
import io.trino.spi.exchange.ExchangeSink;
import io.trino.spi.exchange.ExchangeSinkInstanceHandle;
import jakarta.annotation.Nullable;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.Executor;
import java.util.function.Supplier;
import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkState;
import static com.google.common.util.concurrent.Futures.immediateFuture;
import static io.trino.execution.buffer.BufferResult.emptyResults;
import static io.trino.execution.buffer.BufferState.FINISHED;
import static java.util.Objects.requireNonNull;
public class LazyOutputBuffer
implements OutputBuffer
{
private final OutputBufferStateMachine stateMachine;
private final String taskInstanceId;
private final DataSize maxBufferSize;
private final DataSize maxBroadcastBufferSize;
private final Supplier memoryContextSupplier;
private final Executor executor;
private final Runnable notifyStatusChanged;
private final ExchangeManagerRegistry exchangeManagerRegistry;
// Note: this is a write once field, so an unsynchronized volatile read that returns a non-null value is safe, but if a null value is observed instead
// a subsequent synchronized read is required to ensure the writing thread can complete any in-flight initialization
@GuardedBy("this")
private volatile OutputBuffer delegate;
@GuardedBy("this")
private final Set destroyedBuffers = new HashSet<>();
@GuardedBy("this")
private final List pendingReads = new ArrayList<>();
public LazyOutputBuffer(
TaskId taskId,
String taskInstanceId,
Executor executor,
DataSize maxBufferSize,
DataSize maxBroadcastBufferSize,
Supplier memoryContextSupplier,
Runnable notifyStatusChanged,
ExchangeManagerRegistry exchangeManagerRegistry)
{
this.taskInstanceId = requireNonNull(taskInstanceId, "taskInstanceId is null");
this.executor = requireNonNull(executor, "executor is null");
stateMachine = new OutputBufferStateMachine(taskId, executor);
this.maxBufferSize = requireNonNull(maxBufferSize, "maxBufferSize is null");
this.maxBroadcastBufferSize = requireNonNull(maxBroadcastBufferSize, "maxBroadcastBufferSize is null");
checkArgument(maxBufferSize.toBytes() > 0, "maxBufferSize must be at least 1");
this.memoryContextSupplier = requireNonNull(memoryContextSupplier, "memoryContextSupplier is null");
this.notifyStatusChanged = requireNonNull(notifyStatusChanged, "notifyStatusChanged is null");
this.exchangeManagerRegistry = requireNonNull(exchangeManagerRegistry, "exchangeManagerRegistry is null");
}
@Override
public void addStateChangeListener(StateChangeListener stateChangeListener)
{
stateMachine.addStateChangeListener(stateChangeListener);
}
@Override
public double getUtilization()
{
OutputBuffer outputBuffer = getDelegateOutputBuffer();
// until output buffer is initialized, it is "full"
if (outputBuffer == null) {
return 1.0;
}
return outputBuffer.getUtilization();
}
@Override
public OutputBufferStatus getStatus()
{
OutputBuffer outputBuffer = getDelegateOutputBuffer();
if (outputBuffer == null) {
return OutputBufferStatus.initial();
}
return outputBuffer.getStatus();
}
@Override
public OutputBufferInfo getInfo()
{
OutputBuffer outputBuffer = getDelegateOutputBuffer();
if (outputBuffer == null) {
//
// NOTE: this code must be lock free to not hanging state machine updates
//
BufferState state = stateMachine.getState();
return new OutputBufferInfo(
"UNINITIALIZED",
state,
state.canAddBuffers(),
state.canAddPages(),
0,
0,
0,
0,
Optional.empty(),
Optional.empty(),
Optional.empty());
}
return outputBuffer.getInfo();
}
@Override
public BufferState getState()
{
return stateMachine.getState();
}
@Override
public void setOutputBuffers(OutputBuffers newOutputBuffers)
{
Set destroyedBuffers = ImmutableSet.of();
List pendingReads = ImmutableList.of();
OutputBuffer outputBuffer = delegate;
if (outputBuffer == null) {
synchronized (this) {
outputBuffer = delegate;
if (outputBuffer == null) {
// ignore set output if buffer was already destroyed or failed
if (stateMachine.getState().isTerminal()) {
return;
}
if (newOutputBuffers instanceof PipelinedOutputBuffers outputBuffers) {
outputBuffer = switch (outputBuffers.getType()) {
case PARTITIONED -> new PartitionedOutputBuffer(taskInstanceId, stateMachine, outputBuffers, maxBufferSize, memoryContextSupplier, executor);
case BROADCAST -> new BroadcastOutputBuffer(taskInstanceId, stateMachine, maxBroadcastBufferSize, memoryContextSupplier, executor, notifyStatusChanged);
case ARBITRARY -> new ArbitraryOutputBuffer(taskInstanceId, stateMachine, maxBufferSize, memoryContextSupplier, executor);
};
}
else if (newOutputBuffers instanceof SpoolingOutputBuffers outputBuffers) {
ExchangeSinkInstanceHandle exchangeSinkInstanceHandle = outputBuffers.getExchangeSinkInstanceHandle();
ExchangeManager exchangeManager = exchangeManagerRegistry.getExchangeManager();
ExchangeSink exchangeSink = exchangeManager.createSink(exchangeSinkInstanceHandle);
outputBuffer = new SpoolingExchangeOutputBuffer(stateMachine, outputBuffers, exchangeSink, memoryContextSupplier);
}
else {
throw new IllegalArgumentException("Unexpected output buffers type: " + newOutputBuffers.getClass());
}
// process pending aborts and reads outside of synchronized lock
destroyedBuffers = ImmutableSet.copyOf(this.destroyedBuffers);
this.destroyedBuffers.clear();
pendingReads = ImmutableList.copyOf(this.pendingReads);
this.pendingReads.clear();
// Must be assigned last to avoid a race condition with unsynchronized readers
delegate = outputBuffer;
}
}
}
outputBuffer.setOutputBuffers(newOutputBuffers);
// process pending aborts and reads outside of synchronized lock
destroyedBuffers.forEach(outputBuffer::destroy);
for (PendingRead pendingRead : pendingReads) {
pendingRead.process(outputBuffer);
}
}
@Override
public ListenableFuture get(OutputBufferId bufferId, long token, DataSize maxSize)
{
OutputBuffer outputBuffer = delegate;
if (outputBuffer == null) {
synchronized (this) {
if (delegate == null) {
if (stateMachine.getState().isTerminal()) {
// only set complete when finished, otherwise
boolean complete = stateMachine.getState() == FINISHED;
return immediateFuture(emptyResults(taskInstanceId, 0, complete));
}
PendingRead pendingRead = new PendingRead(bufferId, token, maxSize);
pendingReads.add(pendingRead);
return pendingRead.getFutureResult();
}
outputBuffer = delegate;
}
}
return outputBuffer.get(bufferId, token, maxSize);
}
@Override
public void acknowledge(OutputBufferId bufferId, long token)
{
OutputBuffer outputBuffer = getDelegateOutputBufferOrFail();
outputBuffer.acknowledge(bufferId, token);
}
@Override
public void destroy(OutputBufferId bufferId)
{
OutputBuffer outputBuffer = delegate;
if (outputBuffer == null) {
synchronized (this) {
if (delegate == null) {
destroyedBuffers.add(bufferId);
// Normally, we should free any pending readers for this buffer,
// but we assume that the real buffer will be created quickly.
return;
}
outputBuffer = delegate;
}
}
outputBuffer.destroy(bufferId);
}
@Override
public ListenableFuture isFull()
{
OutputBuffer outputBuffer = getDelegateOutputBufferOrFail();
return outputBuffer.isFull();
}
@Override
public void enqueue(List pages)
{
OutputBuffer outputBuffer = getDelegateOutputBufferOrFail();
outputBuffer.enqueue(pages);
}
@Override
public void enqueue(int partition, List pages)
{
OutputBuffer outputBuffer = getDelegateOutputBufferOrFail();
outputBuffer.enqueue(partition, pages);
}
@Override
public void setNoMorePages()
{
OutputBuffer outputBuffer = getDelegateOutputBufferOrFail();
outputBuffer.setNoMorePages();
}
@Override
public void destroy()
{
List pendingReads = ImmutableList.of();
OutputBuffer outputBuffer = delegate;
if (outputBuffer == null) {
synchronized (this) {
if (delegate == null) {
// ignore destroy if the buffer already in a terminal state.
if (!stateMachine.finish()) {
return;
}
pendingReads = ImmutableList.copyOf(this.pendingReads);
this.pendingReads.clear();
}
outputBuffer = delegate;
}
}
// if there is no output buffer, free the pending reads
if (outputBuffer == null) {
for (PendingRead pendingRead : pendingReads) {
pendingRead.getFutureResult().set(emptyResults(taskInstanceId, 0, true));
}
return;
}
outputBuffer.destroy();
}
@Override
public void abort()
{
List pendingReads = ImmutableList.of();
OutputBuffer outputBuffer = delegate;
if (outputBuffer == null) {
synchronized (this) {
if (delegate == null) {
// ignore abort if the buffer already in a terminal state.
if (!stateMachine.abort()) {
return;
}
pendingReads = ImmutableList.copyOf(this.pendingReads);
this.pendingReads.clear();
}
outputBuffer = delegate;
}
}
// if there is no output buffer, send an empty result without buffer completed signaled
if (outputBuffer == null) {
for (PendingRead pendingRead : pendingReads) {
pendingRead.getFutureResult().set(emptyResults(taskInstanceId, 0, false));
}
return;
}
outputBuffer.abort();
}
@Override
public long getPeakMemoryUsage()
{
OutputBuffer outputBuffer = getDelegateOutputBuffer();
if (outputBuffer != null) {
return outputBuffer.getPeakMemoryUsage();
}
return 0;
}
@Override
public Optional getFailureCause()
{
return stateMachine.getFailureCause();
}
@Nullable
private OutputBuffer getDelegateOutputBuffer()
{
OutputBuffer outputBuffer = delegate;
if (outputBuffer == null) {
synchronized (this) {
outputBuffer = delegate;
}
}
return outputBuffer;
}
private OutputBuffer getDelegateOutputBufferOrFail()
{
OutputBuffer outputBuffer = getDelegateOutputBuffer();
checkState(outputBuffer != null, "Buffer has not been initialized");
return outputBuffer;
}
private static class PendingRead
{
private final OutputBufferId bufferId;
private final long startingSequenceId;
private final DataSize maxSize;
private final SettableFuture futureResult = SettableFuture.create();
public PendingRead(OutputBufferId bufferId, long startingSequenceId, DataSize maxSize)
{
this.bufferId = requireNonNull(bufferId, "bufferId is null");
this.startingSequenceId = startingSequenceId;
this.maxSize = requireNonNull(maxSize, "maxSize is null");
}
public SettableFuture getFutureResult()
{
return futureResult;
}
public void process(OutputBuffer delegate)
{
if (futureResult.isDone()) {
return;
}
try {
ListenableFuture result = delegate.get(bufferId, startingSequenceId, maxSize);
futureResult.setFuture(result);
}
catch (Exception e) {
futureResult.setException(e);
}
}
}
}
© 2015 - 2025 Weber Informatics LLC | Privacy Policy