Fix race condition in Arbitrary/BroadcastOutputBuffer

A buffer may be transitioned to ABORTED while in process of setting
final OutputBuffers (with noMoreBuffers flag set to true).

core/trino-main/src/main/java/io/trino/execution/buffer/ArbitraryOutputBuffer.java

@@ -45,6 +45,7 @@
 import static com.google.common.base.MoreObjects.toStringHelper;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkState;
+import static io.trino.execution.buffer.BufferState.ABORTED;
 import static io.trino.execution.buffer.BufferState.FINISHED;
 import static io.trino.execution.buffer.BufferState.FLUSHING;
 import static io.trino.execution.buffer.BufferState.NO_MORE_PAGES;
@@ -370,7 +371,9 @@ private synchronized ClientBuffer getBuffer(OutputBufferId id)
         // NOTE: buffers are allowed to be created in the FINISHED state because destroy() can move to the finished state
         // without a clean "no-more-buffers" message from the scheduler.  This happens with limit queries and is ok because
         // the buffer will be immediately destroyed.
-        checkState(stateMachine.getState().canAddBuffers() || !outputBuffers.isNoMoreBufferIds(), "No more buffers already set");
+        BufferState state = stateMachine.getState();
+        // buffer may become aborted while the final output buffers are being set
+        checkState(state == ABORTED || state.canAddBuffers() || !outputBuffers.isNoMoreBufferIds(), "No more buffers already set");
 
         // NOTE: buffers are allowed to be created before they are explicitly declared by setOutputBuffers
         // When no-more-buffers is set, we verify that all created buffers have been declared

core/trino-main/src/main/java/io/trino/execution/buffer/BroadcastOutputBuffer.java

@@ -370,7 +370,8 @@ private synchronized ClientBuffer getBuffer(OutputBufferId id)
         // without a clean "no-more-buffers" message from the scheduler.  This happens with limit queries and is ok because
         // the buffer will be immediately destroyed.
         BufferState state = stateMachine.getState();
-        checkState(state.canAddBuffers() || !outputBuffers.isNoMoreBufferIds(), "No more buffers already set");
+        // buffer may become aborted while the final output buffers are being set
+        checkState(state == ABORTED || state.canAddBuffers() || !outputBuffers.isNoMoreBufferIds(), "No more buffers already set");
 
         // NOTE: buffers are allowed to be created before they are explicitly declared by setOutputBuffers
         // When no-more-buffers is set, we verify that all created buffers have been declared