@VisibleForTesting
public synchronized void acknowledge(String outputId, long sequenceId) {
Preconditions.checkNotNull(outputId, "outputId is null");
NamedQueue namedQueue = namedQueues.get(outputId);
if (namedQueue == null) {
throw new NoSuchBufferException(outputId, namedQueues.keySet());
}
if (state == QueueState.FINISHED) {
return;
}
// remove queue from set before calling getPages because getPages changes
// the sequence number of the queue which is used for identity comparison in the
// sorted set
openQueuesBySequenceId.remove(namedQueue);
// acknowledge the pages
namedQueue.acknowledge(sequenceId);
// only add back the queue if it is still open
if (!closed.get()) {
openQueuesBySequenceId.add(namedQueue);
} else {
namedQueue.setFinished();
}
updateState();
}
public synchronized void noMoreQueues() {
namedQueues = ImmutableMap.copyOf(namedQueues);
if (state != QueueState.OPEN) {
return;
}
state = QueueState.NO_MORE_QUEUES;
updateState();
}
public synchronized void abort(String outputId) {
Preconditions.checkNotNull(outputId, "outputId is null");
NamedQueue namedQueue = namedQueues.get(outputId);
if (namedQueue == null || namedQueue.isFinished()) {
return;
}
namedQueue.setFinished();
openQueuesBySequenceId.remove(namedQueue);
updateState();
}
public synchronized BufferResult get(
String outputId, long startingSequenceId, DataSize maxSize, Duration maxWait)
throws InterruptedException {
Preconditions.checkNotNull(outputId, "outputId is null");
Preconditions.checkArgument(maxSize.toBytes() > 0, "maxSize must be at least 1 byte");
Preconditions.checkNotNull(maxWait, "maxWait is null");
NamedQueue namedQueue = namedQueues.get(outputId);
if (namedQueue == null) {
throw new NoSuchBufferException(outputId, namedQueues.keySet());
}
if (state == QueueState.FINISHED) {
return emptyResults(namedQueue.getSequenceId(), true);
}
// wait for pages to arrive
if (namedQueue.isEmpty()) {
long remainingNanos = maxWait.roundTo(NANOSECONDS);
long end = System.nanoTime() + remainingNanos;
while (remainingNanos > 0 && namedQueue.isEmpty() && !namedQueue.isFinished()) {
// wait for timeout or notification
NANOSECONDS.timedWait(this, remainingNanos);
remainingNanos = end - System.nanoTime();
}
}
// remove queue from set before calling getPages because getPages changes
// the sequence number of the queue which is used for identity comparison in the
// sorted set
openQueuesBySequenceId.remove(namedQueue);
// get the pages
BufferResult results = namedQueue.getPages(startingSequenceId, maxSize);
// only add back the queue if it is still open
if (!closed.get() || !results.isBufferClosed()) {
openQueuesBySequenceId.add(namedQueue);
} else {
namedQueue.setFinished();
}
updateState();
return results;
}
/**
* Marks the output as complete. After this method is called no more data can be added but there
* may still be buffered output pages.
*/
public synchronized void finish() {
closed.set(true);
// the output will only transition to finished if it isn't already marked as failed or cancel
updateState();
}