@Override public void setOutputBuffers(OutputBuffers newOutputBuffers) { checkState(!Thread.holdsLock(this), "Can not set output buffers while holding a lock on this"); requireNonNull(newOutputBuffers, "newOutputBuffers is null"); synchronized (this) { // ignore buffers added after query finishes, which can happen when a query is canceled // also ignore old versions, which is normal BufferState state = this.state.get(); if (state.isTerminal() || outputBuffers.getVersion() >= newOutputBuffers.getVersion()) { return; } // verify this is valid state change outputBuffers.checkValidTransition(newOutputBuffers); outputBuffers = newOutputBuffers; // add the new buffers for (OutputBufferId outputBufferId : outputBuffers.getBuffers().keySet()) { getBuffer(outputBufferId); } // update state if no more buffers is set if (outputBuffers.isNoMoreBufferIds()) { this.state.compareAndSet(OPEN, NO_MORE_BUFFERS); this.state.compareAndSet(NO_MORE_PAGES, FLUSHING); } } if (!state.get().canAddBuffers()) { noMoreBuffers(); } checkFlushComplete(); }
@Test public void testAddQueueAfterCreation() { PartitionedOutputBuffer buffer = createPartitionedBuffer( createInitialEmptyOutputBuffers(PARTITIONED) .withBuffer(FIRST, 0) .withNoMoreBufferIds(), sizeOfPages(10)); assertFalse(buffer.isFinished()); try { buffer.setOutputBuffers(createInitialEmptyOutputBuffers(PARTITIONED) .withBuffer(FIRST, 0) .withBuffer(SECOND, 0) .withNoMoreBufferIds()); fail("Expected IllegalStateException from addQueue after noMoreQueues has been called"); } catch (IllegalArgumentException ignored) { } }
public OutputBuffers withBuffer(OutputBufferId bufferId, int partition) { requireNonNull(bufferId, "bufferId is null"); if (buffers.containsKey(bufferId)) { checkHasBuffer(bufferId, partition); return this; } // verify no new buffers is not set checkState(!noMoreBufferIds, "No more buffer ids already set"); return new OutputBuffers( type, version + 1, false, ImmutableMap.<OutputBufferId, Integer>builder() .putAll(buffers) .put(bufferId, partition) .build()); }
private synchronized void noMoreBuffers() { if (outputBuffers.isNoMoreBufferIds()) { // verify all created buffers have been declared SetView<OutputBufferId> undeclaredCreatedBuffers = Sets.difference(buffers.keySet(), outputBuffers.getBuffers().keySet()); checkState(undeclaredCreatedBuffers.isEmpty(), "Final output buffers does not contain all created buffer ids: %s", undeclaredCreatedBuffers); } }
if (outputBuffers.isNoMoreBufferIds()) { outputBuffers = outputBuffers.withBuffer(newBuffer, BROADCAST_PARTITION_ID); outputBuffers = outputBuffers.withNoMoreBufferIds();
public PartitionedOutputBufferManager(PartitioningHandle partitioningHandle, int partitionCount, Consumer<OutputBuffers> outputBufferTarget) { checkArgument(partitionCount >= 1, "partitionCount must be at least 1"); ImmutableMap.Builder<OutputBufferId, Integer> partitions = ImmutableMap.builder(); for (int partition = 0; partition < partitionCount; partition++) { partitions.put(new OutputBufferId(partition), partition); } OutputBuffers outputBuffers = createInitialEmptyOutputBuffers(requireNonNull(partitioningHandle, "partitioningHandle is null")) .withBuffers(partitions.build()) .withNoMoreBufferIds(); outputBufferTarget.accept(outputBuffers); this.outputBuffers = outputBuffers.getBuffers(); }
private static void assertOutputBuffers(OutputBuffers outputBuffers) { assertNotNull(outputBuffers); assertTrue(outputBuffers.getVersion() > 0); assertTrue(outputBuffers.isNoMoreBufferIds()); Map<OutputBufferId, Integer> buffers = outputBuffers.getBuffers(); assertEquals(buffers.size(), 4); for (int partition = 0; partition < 4; partition++) { assertEquals(buffers.get(new OutputBufferId(partition)), Integer.valueOf(partition)); } } }
@Test public void testNoMorePagesFreesReader() { ArbitraryOutputBuffer buffer = createArbitraryBuffer(createInitialEmptyOutputBuffers(ARBITRARY), sizeOfPages(10)); buffer.setOutputBuffers(createInitialEmptyOutputBuffers(ARBITRARY).withBuffer(FIRST, 0)); assertFalse(buffer.isFinished()); ListenableFuture<BufferResult> future = buffer.get(FIRST, 0, sizeOfPages(10)); assertFalse(future.isDone()); buffer.setNoMorePages(); assertTrue(future.isDone()); assertTrue(buffer.get(FIRST, 0, sizeOfPages(10)).isDone()); }
@Test public void testAddQueueAfterNoMoreQueues() { BroadcastOutputBuffer buffer = createBroadcastBuffer(createInitialEmptyOutputBuffers(BROADCAST), sizeOfPages(10)); assertFalse(buffer.isFinished()); // tell buffer no more queues will be added buffer.setOutputBuffers(createInitialEmptyOutputBuffers(BROADCAST).withNoMoreBufferIds()); assertTrue(buffer.isFinished()); // set no more queues a second time to assure that we don't get an exception or such buffer.setOutputBuffers(createInitialEmptyOutputBuffers(BROADCAST).withNoMoreBufferIds()); assertTrue(buffer.isFinished()); // set no more queues a third time to assure that we don't get an exception or such buffer.setOutputBuffers(createInitialEmptyOutputBuffers(BROADCAST).withNoMoreBufferIds()); assertTrue(buffer.isFinished()); }
public PartitionedOutputBuffer( String taskInstanceId, StateMachine<BufferState> state, OutputBuffers outputBuffers, DataSize maxBufferSize, Supplier<LocalMemoryContext> systemMemoryContextSupplier, Executor notificationExecutor) { this.state = requireNonNull(state, "state is null"); requireNonNull(outputBuffers, "outputBuffers is null"); checkArgument(outputBuffers.getType() == PARTITIONED, "Expected a PARTITIONED output buffer descriptor"); checkArgument(outputBuffers.isNoMoreBufferIds(), "Expected a final output buffer descriptor"); this.outputBuffers = outputBuffers; this.memoryManager = new OutputBufferMemoryManager( requireNonNull(maxBufferSize, "maxBufferSize is null").toBytes(), requireNonNull(systemMemoryContextSupplier, "systemMemoryContextSupplier is null"), requireNonNull(notificationExecutor, "notificationExecutor is null")); ImmutableList.Builder<ClientBuffer> partitions = ImmutableList.builder(); for (OutputBufferId bufferId : outputBuffers.getBuffers().keySet()) { ClientBuffer partition = new ClientBuffer(taskInstanceId, bufferId); partitions.add(partition); } this.partitions = partitions.build(); state.compareAndSet(OPEN, NO_MORE_BUFFERS); state.compareAndSet(NO_MORE_PAGES, FLUSHING); checkFlushComplete(); }
@Override public void setOutputBuffers(OutputBuffers newOutputBuffers) { requireNonNull(newOutputBuffers, "newOutputBuffers is null"); // ignore buffers added after query finishes, which can happen when a query is canceled // also ignore old versions, which is normal if (state.get().isTerminal() || outputBuffers.getVersion() >= newOutputBuffers.getVersion()) { return; } // no more buffers can be added but verify this is valid state change outputBuffers.checkValidTransition(newOutputBuffers); }
OutputBufferId rootBufferId = Iterables.getOnlyElement(rootOutputBuffers.getBuffers().keySet()); List<SqlStageExecution> stages = createStages( (fragmentId, tasks, noMoreExchangeLocations) -> updateQueryOutputLocations(queryStateMachine, rootBufferId, tasks, noMoreExchangeLocations),
private RemoteTask createRemoteTask(HttpRemoteTaskFactory httpRemoteTaskFactory) { return httpRemoteTaskFactory.createRemoteTask( TEST_SESSION, new TaskId("test", 1, 2), new PrestoNode("node-id", URI.create("http://fake.invalid/"), new NodeVersion("version"), false), TaskTestUtils.PLAN_FRAGMENT, ImmutableMultimap.of(), OptionalInt.empty(), createInitialEmptyOutputBuffers(OutputBuffers.BufferType.BROADCAST), new NodeTaskMap.PartitionedSplitCountTracker(i -> {}), true); }
@Override public synchronized void setOutputBuffers(OutputBuffers newOutputBuffers) { if (getTaskStatus().getState().isDone()) { return; } if (newOutputBuffers.getVersion() > outputBuffers.get().getVersion()) { outputBuffers.set(newOutputBuffers); needsUpdate.set(true); scheduleUpdate(); } }
private synchronized ClientBuffer getBuffer(OutputBufferId id) { ClientBuffer buffer = buffers.get(id); if (buffer != null) { return buffer; } // 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(state.get().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 buffer = new ClientBuffer(taskInstanceId, id); // buffer may have finished immediately before calling this method if (state.get() == FINISHED) { buffer.destroy(); } buffers.put(id, buffer); return buffer; }
public OutputBuffers withNoMoreBufferIds() { if (noMoreBufferIds) { return this; } return new OutputBuffers(type, version + 1, true, buffers); }
public void checkValidTransition(OutputBuffers newOutputBuffers) { requireNonNull(newOutputBuffers, "newOutputBuffers is null"); checkState(type == newOutputBuffers.getType(), "newOutputBuffers has a different type"); if (noMoreBufferIds) { checkArgument(this.equals(newOutputBuffers), "Expected buffer to not change after no more buffers is set"); return; } if (version > newOutputBuffers.version) { throw new IllegalArgumentException("newOutputBuffers version is older"); } if (version == newOutputBuffers.version) { checkArgument(this.equals(newOutputBuffers), "newOutputBuffers is the same version but contains different information"); } // assure we have not changed the buffer assignments for (Entry<OutputBufferId, Integer> entry : buffers.entrySet()) { if (!entry.getValue().equals(newOutputBuffers.buffers.get(entry.getKey()))) { throw new IllegalArgumentException("newOutputBuffers has changed the assignment for task " + entry.getKey()); } } }
public void checkValidTransition(OutputBuffers newOutputBuffers) { requireNonNull(newOutputBuffers, "newOutputBuffers is null"); if (version > newOutputBuffers.version) { throw new IllegalArgumentException("newOutputBuffers version is older"); } if (version == newOutputBuffers.version) { checkArgument(this.equals(newOutputBuffers), "newOutputBuffers is the same version but contains different information"); } // assure we are not removing the no more buffers flag if (noMoreBufferIds && !newOutputBuffers.noMoreBufferIds) { throw new IllegalArgumentException("Expected newOutputBuffers to have noMoreBufferIds set"); } // assure we have not changed the buffer assignments for (Entry<TaskId, Integer> entry : buffers.entrySet()) { if (!entry.getValue().equals(newOutputBuffers.buffers.get(entry.getKey()))) { throw new IllegalArgumentException("newOutputBuffers has changed the assignment for task " + entry.getKey()); } } }
if (outputBuffers.isNoMoreBufferIds()) { outputBuffers = outputBuffers.withBuffer(newBuffer, newBuffer.getId()); outputBuffers = outputBuffers.withNoMoreBufferIds();
assertTrue(outputBuffers.getVersion() > 0); assertTrue(outputBuffers.isNoMoreBufferIds()); Map<TaskId, Integer> buffers = outputBuffers.getBuffers(); assertEquals(buffers.size(), 4); assertEquals(buffers.get(new TaskId(STAGE_ID, "0")), Integer.valueOf(2));