java - 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 PhantomReferences 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.