diff --git a/server-spi-private/src/main/java/org/keycloak/utils/ClosingDoubleStream.java b/server-spi-private/src/main/java/org/keycloak/utils/ClosingDoubleStream.java index afebc8b7ed..4c8dc16b6d 100644 --- a/server-spi-private/src/main/java/org/keycloak/utils/ClosingDoubleStream.java +++ b/server-spi-private/src/main/java/org/keycloak/utils/ClosingDoubleStream.java @@ -16,6 +16,7 @@ */ package org.keycloak.utils; +import java.util.Comparator; import java.util.DoubleSummaryStatistics; import java.util.OptionalDouble; import java.util.PrimitiveIterator; @@ -250,16 +251,110 @@ class ClosingDoubleStream implements DoubleStream { @Override public PrimitiveIterator.OfDouble iterator() { - return delegate.iterator(); + return new ClosingIterator(delegate.iterator()); } @Override public Spliterator.OfDouble spliterator() { - return delegate.spliterator(); + return new ClosingSpliterator(delegate.spliterator()); } @Override public boolean isParallel() { return delegate.isParallel(); } + + private class ClosingIterator implements PrimitiveIterator.OfDouble { + + private final PrimitiveIterator.OfDouble iterator; + + public ClosingIterator(PrimitiveIterator.OfDouble iterator) { + this.iterator = iterator; + } + + @Override + public boolean hasNext() { + final boolean res = iterator.hasNext(); + if (! res) { + close(); + } + return res; + } + + @Override + public Double next() { + return iterator.next(); + } + + @Override + public void remove() { + iterator.remove(); + } + + @Override + public void forEachRemaining(DoubleConsumer action) { + iterator.forEachRemaining(action); + close(); + } + + @Override + public double nextDouble() { + return iterator.nextDouble(); + } + } + + private class ClosingSpliterator implements Spliterator.OfDouble { + + private final Spliterator.OfDouble spliterator; + + public ClosingSpliterator(Spliterator.OfDouble spliterator) { + this.spliterator = spliterator; + } + + @Override + public boolean tryAdvance(DoubleConsumer action) { + final boolean res = spliterator.tryAdvance(action); + if (! res) { + close(); + } + return res; + } + + @Override + public void forEachRemaining(DoubleConsumer action) { + spliterator.forEachRemaining(action); + close(); + } + + @Override + public Spliterator.OfDouble trySplit() { + return spliterator.trySplit(); + } + + @Override + public long estimateSize() { + return spliterator.estimateSize(); + } + + @Override + public long getExactSizeIfKnown() { + return spliterator.getExactSizeIfKnown(); + } + + @Override + public int characteristics() { + return spliterator.characteristics(); + } + + @Override + public boolean hasCharacteristics(int characteristics) { + return spliterator.hasCharacteristics(characteristics); + } + + @Override + public Comparator getComparator() { + return spliterator.getComparator(); + } + + } } diff --git a/server-spi-private/src/main/java/org/keycloak/utils/ClosingIntStream.java b/server-spi-private/src/main/java/org/keycloak/utils/ClosingIntStream.java index 5e86f02123..9c3570bbee 100644 --- a/server-spi-private/src/main/java/org/keycloak/utils/ClosingIntStream.java +++ b/server-spi-private/src/main/java/org/keycloak/utils/ClosingIntStream.java @@ -16,6 +16,7 @@ */ package org.keycloak.utils; +import java.util.Comparator; import java.util.IntSummaryStatistics; import java.util.OptionalDouble; import java.util.OptionalInt; @@ -265,16 +266,110 @@ class ClosingIntStream implements IntStream { @Override public PrimitiveIterator.OfInt iterator() { - return delegate.iterator(); + return new ClosingIterator(delegate.iterator()); } @Override public Spliterator.OfInt spliterator() { - return delegate.spliterator(); + return new ClosingSpliterator(delegate.spliterator()); } @Override public boolean isParallel() { return delegate.isParallel(); } + + private class ClosingIterator implements PrimitiveIterator.OfInt { + + private final PrimitiveIterator.OfInt iterator; + + public ClosingIterator(PrimitiveIterator.OfInt iterator) { + this.iterator = iterator; + } + + @Override + public boolean hasNext() { + final boolean res = iterator.hasNext(); + if (! res) { + close(); + } + return res; + } + + @Override + public Integer next() { + return iterator.next(); + } + + @Override + public void remove() { + iterator.remove(); + } + + @Override + public void forEachRemaining(IntConsumer action) { + iterator.forEachRemaining(action); + close(); + } + + @Override + public int nextInt() { + return iterator.nextInt(); + } + } + + private class ClosingSpliterator implements Spliterator.OfInt { + + private final Spliterator.OfInt spliterator; + + public ClosingSpliterator(Spliterator.OfInt spliterator) { + this.spliterator = spliterator; + } + + @Override + public boolean tryAdvance(IntConsumer action) { + final boolean res = spliterator.tryAdvance(action); + if (! res) { + close(); + } + return res; + } + + @Override + public void forEachRemaining(IntConsumer action) { + spliterator.forEachRemaining(action); + close(); + } + + @Override + public Spliterator.OfInt trySplit() { + return spliterator.trySplit(); + } + + @Override + public long estimateSize() { + return spliterator.estimateSize(); + } + + @Override + public long getExactSizeIfKnown() { + return spliterator.getExactSizeIfKnown(); + } + + @Override + public int characteristics() { + return spliterator.characteristics(); + } + + @Override + public boolean hasCharacteristics(int characteristics) { + return spliterator.hasCharacteristics(characteristics); + } + + @Override + public Comparator getComparator() { + return spliterator.getComparator(); + } + + } } diff --git a/server-spi-private/src/main/java/org/keycloak/utils/ClosingLongStream.java b/server-spi-private/src/main/java/org/keycloak/utils/ClosingLongStream.java index b2832ab2d0..f9fbf8f35c 100644 --- a/server-spi-private/src/main/java/org/keycloak/utils/ClosingLongStream.java +++ b/server-spi-private/src/main/java/org/keycloak/utils/ClosingLongStream.java @@ -16,6 +16,7 @@ */ package org.keycloak.utils; +import java.util.Comparator; import java.util.LongSummaryStatistics; import java.util.OptionalDouble; import java.util.OptionalLong; @@ -258,16 +259,110 @@ class ClosingLongStream implements LongStream { @Override public PrimitiveIterator.OfLong iterator() { - return delegate.iterator(); + return new ClosingIterator(delegate.iterator()); } @Override public Spliterator.OfLong spliterator() { - return delegate.spliterator(); + return new ClosingSpliterator(delegate.spliterator()); } @Override public boolean isParallel() { return delegate.isParallel(); } + + private class ClosingIterator implements PrimitiveIterator.OfLong { + + private final PrimitiveIterator.OfLong iterator; + + public ClosingIterator(PrimitiveIterator.OfLong iterator) { + this.iterator = iterator; + } + + @Override + public boolean hasNext() { + final boolean res = iterator.hasNext(); + if (! res) { + close(); + } + return res; + } + + @Override + public Long next() { + return iterator.next(); + } + + @Override + public void remove() { + iterator.remove(); + } + + @Override + public void forEachRemaining(LongConsumer action) { + iterator.forEachRemaining(action); + close(); + } + + @Override + public long nextLong() { + return iterator.nextLong(); + } + } + + private class ClosingSpliterator implements Spliterator.OfLong { + + private final Spliterator.OfLong spliterator; + + public ClosingSpliterator(Spliterator.OfLong spliterator) { + this.spliterator = spliterator; + } + + @Override + public boolean tryAdvance(LongConsumer action) { + final boolean res = spliterator.tryAdvance(action); + if (! res) { + close(); + } + return res; + } + + @Override + public void forEachRemaining(LongConsumer action) { + spliterator.forEachRemaining(action); + close(); + } + + @Override + public Spliterator.OfLong trySplit() { + return spliterator.trySplit(); + } + + @Override + public long estimateSize() { + return spliterator.estimateSize(); + } + + @Override + public long getExactSizeIfKnown() { + return spliterator.getExactSizeIfKnown(); + } + + @Override + public int characteristics() { + return spliterator.characteristics(); + } + + @Override + public boolean hasCharacteristics(int characteristics) { + return spliterator.hasCharacteristics(characteristics); + } + + @Override + public Comparator getComparator() { + return spliterator.getComparator(); + } + + } } diff --git a/server-spi-private/src/main/java/org/keycloak/utils/ClosingStream.java b/server-spi-private/src/main/java/org/keycloak/utils/ClosingStream.java index 5848752747..2c554720f4 100644 --- a/server-spi-private/src/main/java/org/keycloak/utils/ClosingStream.java +++ b/server-spi-private/src/main/java/org/keycloak/utils/ClosingStream.java @@ -242,12 +242,12 @@ class ClosingStream implements Stream { @Override public Iterator iterator() { - return delegate.iterator(); + return new ClosingIterator(delegate.iterator()); } @Override public Spliterator spliterator() { - return delegate.spliterator(); + return new ClosingSpliterator(delegate.spliterator()); } @Override @@ -280,4 +280,92 @@ class ClosingStream implements Stream { delegate.close(); } + private class ClosingIterator implements Iterator { + + private final Iterator iterator; + + public ClosingIterator(Iterator iterator) { + this.iterator = iterator; + } + + @Override + public boolean hasNext() { + final boolean res = iterator.hasNext(); + if (! res) { + close(); + } + return res; + } + + @Override + public R next() { + return iterator.next(); + } + + @Override + public void remove() { + iterator.remove(); + } + + @Override + public void forEachRemaining(Consumer action) { + iterator.forEachRemaining(action); + close(); + } + } + + private class ClosingSpliterator implements Spliterator { + + private final Spliterator spliterator; + + public ClosingSpliterator(Spliterator spliterator) { + this.spliterator = spliterator; + } + + @Override + public boolean tryAdvance(Consumer action) { + final boolean res = spliterator.tryAdvance(action); + if (! res) { + close(); + } + return res; + } + + @Override + public void forEachRemaining(Consumer action) { + spliterator.forEachRemaining(action); + close(); + } + + @Override + public Spliterator trySplit() { + return spliterator.trySplit(); + } + + @Override + public long estimateSize() { + return spliterator.estimateSize(); + } + + @Override + public long getExactSizeIfKnown() { + return spliterator.getExactSizeIfKnown(); + } + + @Override + public int characteristics() { + return spliterator.characteristics(); + } + + @Override + public boolean hasCharacteristics(int characteristics) { + return spliterator.hasCharacteristics(characteristics); + } + + @Override + public Comparator getComparator() { + return spliterator.getComparator(); + } + + } } diff --git a/server-spi-private/src/test/java/org/keycloak/utils/StreamsUtilTest.java b/server-spi-private/src/test/java/org/keycloak/utils/StreamsUtilTest.java index 5546da8370..9719a4451e 100644 --- a/server-spi-private/src/test/java/org/keycloak/utils/StreamsUtilTest.java +++ b/server-spi-private/src/test/java/org/keycloak/utils/StreamsUtilTest.java @@ -36,6 +36,43 @@ public class StreamsUtilTest { Assert.assertTrue(closed.get()); } + @Test + public void testAutoClosingOfClosingStreamOuter() { + AtomicBoolean closed = new AtomicBoolean(); + StreamsUtil.closing(Stream.of(1, 2, 3)).onClose(() -> closed.set(true)).forEach(NOOP); + + Assert.assertTrue(closed.get()); + } + + @Test + public void testAutoClosingOfClosingStreamFlatMap() { + AtomicBoolean closed = new AtomicBoolean(); + Stream.of("value") + .flatMap(v -> StreamsUtil.closing(Stream.of(1, 2, 3)).onClose(() -> closed.set(true))) + .forEach(NOOP); + + Assert.assertTrue(closed.get()); + } + + @Test + public void testAutoClosingOfClosingUsingIterator() { + AtomicBoolean closed = new AtomicBoolean(); + StreamsUtil.closing(Stream.of(1, 2, 3).onClose(() -> closed.set(true))).iterator().forEachRemaining(NOOP); + + Assert.assertTrue(closed.get()); + } + + @Test + public void testAutoClosingOfClosingUsingConcat() { + AtomicBoolean closed = new AtomicBoolean(); + Stream.concat( + Stream.of(4, 5), + StreamsUtil.closing(Stream.of(1, 2, 3).onClose(() -> closed.set(true))) + ).iterator().forEachRemaining(NOOP); + + Assert.assertTrue(closed.get()); + } + @Test public void testMultipleClosingHandlersOnClosingStream() { AtomicInteger firstHandlerFiringCount = new AtomicInteger();