org.apache.baremaps.openstreetmap.stream.BufferedSpliterator Maven / Gradle / Ivy
/*
* (c) Copyright 2017 Palantir Technologies Inc. All rights reserved.
*
* 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 org.apache.baremaps.openstreetmap.stream;
import java.util.Spliterator;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.CompletableFuture;
import java.util.function.Consumer;
/**
* A spliterator that buffers the completion of a spliterator of future elements and returns them
* according to a user defined order.
*
*
* This code has been adapted from
* {@link streams} licensed under the Apache
* License 2.0.
*
*
* Copyright 2017 Palantir Technologies, Inc. All rights reserved.
*
* @param the type of elements returned by this {@code Spliterator}
*/
class BufferedSpliterator implements Spliterator> {
private final CompletionOrder completionOrder;
private final Spliterator> spliterator;
private final int bufferSize;
private final BlockingQueue> buffer;
private int pending = 0;
/**
* Constructs a {@code BufferedSpliterator} from a spliterator of futures elements.
*
* @param spliterator the spliterator to buffer
* @param bufferSize the buffer size
* @param completionOrder the completion order
*/
public BufferedSpliterator(Spliterator> spliterator, int bufferSize,
CompletionOrder completionOrder) {
this.spliterator = spliterator;
this.bufferSize = bufferSize;
this.buffer = new ArrayBlockingQueue<>(bufferSize);
this.completionOrder = completionOrder;
}
/** {@inheritDoc} */
@Override
public boolean tryAdvance(Consumer> action) {
fillBuffer();
if (pending == 0) {
return false;
}
try {
CompletableFuture future = buffer.take();
pending--;
action.accept(future);
return true;
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
throw new StreamException(e);
}
}
/** {@inheritDoc} */
@Override
public Spliterator> trySplit() {
return null;
}
/** {@inheritDoc} */
@Override
public long estimateSize() {
long estimate = pending + spliterator.estimateSize();
if (estimate < 0) {
return Long.MAX_VALUE;
}
return estimate;
}
/** {@inheritDoc} */
@Override
public int characteristics() {
return spliterator.characteristics();
}
private void fillBuffer() {
while (pending < bufferSize && spliterator
.tryAdvance(future -> completionOrder.registerCompletion(future, buffer::add))) {
pending++;
}
}
/** Represents the completion order applied to a {@code BufferedSpliterator}. */
public interface CompletionOrder {
void registerCompletion(CompletableFuture future,
Consumer> resultConsumer);
}
/** An order that registers completions when futures are completed. */
enum InCompletionOrder implements CompletionOrder {
INSTANCE;
@Override
public void registerCompletion(CompletableFuture future,
Consumer> resultConsumer) {
future.whenComplete((result, error) -> resultConsumer.accept(future));
}
}
/** An order that registers completion according to the order of the source. */
enum InSourceOrder implements CompletionOrder {
INSTANCE;
@Override
public void registerCompletion(CompletableFuture future,
Consumer> resultConsumer) {
resultConsumer.accept(future);
}
}
}