Register a Stream "completion" hook

Question

Using the Java 8 Stream API, I would like to register a "completion hook", along the lines of:

Stream<String> stream = Stream.of("a", "b", "c");

// additional filters / mappings that I don't control
stream.onComplete((Completion c) -> {
    // This is what I'd like to do:
    closeResources();

    // This might also be useful:
    Optional<Throwable> exception = c.exception();
    exception.ifPresent(e -> throw new ExceptionWrapper(e));
});

The reason why I want to do that is because I want to wrap a resource in a Stream for API clients to consume, and I want that Stream to clean up the resource automatically once it is consumed. If that were possible, then the client could call:

Collected collectedInOneGo =
Utility.something()
       .niceLookingSQLDSL()
       .moreDSLFeatures()
       .stream()
       .filter(a -> true)
       .map(c -> c)
       .collect(collector);

Rather than what's needed currently:

try (Stream<X> meh = Utility.something()
                            .niceLookingSQLDSL()
                            .moreDSLFeatures()
                            .stream()) {

    Collected collectedWithUglySyntacticDissonance =
    meh.filter(a -> true)
       .map(c -> c)
       .collect(collector);
}

Ideally, I'd like to get into the java.util.stream.ReferencePipeline 's various methods, such as:

@Override
final void forEachWithCancel(Spliterator<P_OUT> spliterator, Sink<P_OUT> sink) {
    try {

        // Existing loop
        do { } while (!sink.cancellationRequested() && spliterator.tryAdvance(sink));
    }

    // These would be nice:
    catch (Throwable t) {
        completion.onFailure(t);
    }
    finally {
        completion.onSuccess();
    }
}

Is there an easy way to do this with existing JDK 8 API?

Answer 1

Any solution intercepting the terminal operations except flatMap -based solution (as proposed by @Holger) would be fragile to the following code:

Stream<String> stream = getAutoCloseableStream();
if(stream.iterator().hasNext()) {
    // do something if stream is non-empty
}

Such usage is absolutely legal by the specification. Do not forget that iterator() and spliterator() are terminal stream operations, but after their execution you still need an access to the stream source. Also it's perfectly valid to abandon the Iterator or Spliterator in any state, so you just cannot know whether it will be used further or not.

You may consider advicing users not to use iterator() and spliterator() , but what about this code?

Stream<String> stream = getAutoCloseableStream();
Stream.concat(stream, Stream.of("xyz")).findFirst();

This internally uses spliterator().tryAdvance() for the first stream, then abandons it (though closes if the resulting stream close() is called explicitly). You will need to ask your users not to use Stream.concat as well. And as far as I know internally in your library you are using iterator() / spliterator() pretty often, so you will need to revisit all these places for possible problems. And, of course there are plenty of other libraries which also use iterator() / spliterator() and may short-circuit after that: all of them would become incompatible with your feature.

Why flatMap -based solution works here? Because upon the first call of the hasNext() or tryAdvance() it dumps the entire stream content into the intermediate buffer and closes the original stream source. So depending on the stream size you may waste much intermediate memory or even have OutOfMemoryError .

You may also consider keeping the PhantomReference s to the Stream objects and monitoring the ReferenceQueue . In this case the completion will be triggered by garbage collector (which also has some drawbacks).

In conclusion my advice is to stay with try-with-resources.

Answer 2

The simplest solution is to wrap a stream in another stream and flatmap it to itself:

// example stream
Stream<String> original=Stream.of("bla").onClose(()->System.out.println("close action"));

// this is the trick
Stream<String> autoClosed=Stream.of(original).flatMap(Function.identity());

//example op
int sum=autoClosed.mapToInt(String::length).sum();
System.out.println(sum);

The reason why it works lies in the flatMap operation :

Each mapped stream is closed after its contents have been placed into this stream.

But the current implementation isn't as lazy as it should be when using flatMap . This has been fixed in Java 10.

---

My recommendation is to stay with the try(…) standard solution and document when a returned stream needs to be closed. After all, a stream that closes the resource after the terminal operation isn't safe as there is no guaranty that the client will actually invoke a terminal operation. Changing its mind and abandon a stream instead, is a valid use, whereas not calling the close() method, when the documentation specifies that it is required, is not.

Answer 3

Java 8 already has a precedent for how streams that need to be closed operate. In their Javadoc , it mentions:

Streams have a BaseStream.close() method and implement AutoCloseable, but nearly all stream instances do not actually need to be closed after use. Generally, only streams whose source is an IO channel (such as those returned by Files.lines(Path, Charset)) will require closing. Most streams are backed by collections, arrays, or generating functions, which require no special resource management. (If a stream does require closing, it can be declared as a resource in a try-with-resources statement.)

So Java 8's recommendation is to open those streams in a try-with-resources. And once you do that, Stream also provides a way for you to add a close hook, almost exactly as you've described: onClose(Runnable) , which accepts a lambda telling it what to do and returns a Stream that will also do that operation when it is closed.

That's the way the API design and documentation suggests to do what you're trying to do.

Answer 4

The solution I've come up with looks like this:

class AutoClosingStream<T> implements Stream<T> {

    AutoClosingStream(Stream<T> delegate, Consumer<Optional<Throwable>> onComplete) {}

    // Pipeline ops delegate the op to the real stream and wrap that again
    @Override
    public Stream<T> limit(long maxSize) {
        return new AutoClosingStream(delegate.limit(maxSize), onComplete);
    }

    // Terminal ops intercept the result and call the onComplete logic
    @Override
    public void forEach(Consumer<? super T> action) {
        terminalOp(() -> delegate.forEach(action));
    }

    private void terminalOp(Runnable runnable) {
        terminalOp(() -> { runnable.run(); return null; });
    }

    private <R> R terminalOp(Supplier<R> supplier) {
        R result = null;

        try {
            result = supplier.get();
            onComplete.accept(Optional.empty());
        }
        catch (Throwable e) {
            onComplete.accept(Optional.of(e));
            Utils.sneakyThrow(e);
        }

        return result;
    }
}

This is only a simplified sketch to illustrate the idea. The real solution would also support the primitive IntStream , LongStream , and DoubleStream

Answer 5

Check out these complete implementations of AutoClosingReferenceStream, AutoClosingIntStream, AutoClosingLongStream and AutoClosingDoubleStream at the open-source project Speedment https://github.com/speedment/speedment/tree/master/src/main/java/com/speedment/internal/core/stream/autoclose

The solution is similar to the one mentioned by @LukasEder