Lazily process sessions from ISPN to avoid fetching client sessions (#39639)

Closes #39638

Signed-off-by: Alexander Schwartz <aschwart@redhat.com>
This commit is contained in:
Alexander Schwartz 2025-05-13 13:16:41 +02:00 committed by GitHub
parent 9b324b9228
commit 4b47697c83
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 53 additions and 2 deletions

View File

@ -81,6 +81,7 @@ import org.keycloak.models.sessions.infinispan.stream.UserSessionPredicate;
import org.keycloak.models.sessions.infinispan.util.FuturesHelper;
import org.keycloak.models.sessions.infinispan.util.InfinispanKeyGenerator;
import org.keycloak.models.sessions.infinispan.util.SessionTimeouts;
import org.keycloak.utils.StreamsUtil;
import static org.keycloak.models.Constants.SESSION_NOTE_LIGHTWEIGHT_USER;
import static org.keycloak.utils.StreamsUtil.paginatedStream;
@ -452,8 +453,10 @@ public class InfinispanUserSessionProvider implements UserSessionProvider, Sessi
UserSessionPredicate predicate = UserSessionPredicate.create(realm.getId()).client(client.getId());
return paginatedStream(getUserSessionsStream(realm, predicate, offline)
.sorted(Comparator.comparing(UserSessionModel::getLastSessionRefresh)), firstResult, maxResults);
// If the sorted stream is used within a flatMap (like in SessionsResource), it will not terminate early unless wrapped with
// StreamsUtil.prepareSortedStreamToWorkInsideOfFlatMapWithTerminalOperations causing unnecessary operations.
return paginatedStream(StreamsUtil.prepareSortedStreamToWorkInsideOfFlatMapWithTerminalOperations(getUserSessionsStream(realm, predicate, offline)
.sorted(Comparator.comparing(UserSessionModel::getLastSessionRefresh))), firstResult, maxResults);
}
@Override

View File

@ -142,4 +142,19 @@ public class StreamsUtil {
}
}, false);
}
/**
* This works around a bug in JDK 21 (but no longer in JDK 24) where a sorted stream has all its elements processed
* when used inside of a flatmap and a terminal operation like limit() is used outside of it.
* See StreamUtilTests.testSortedInsideOfFlatMapShouldRespectTerminalOperation for an example.
* Possible <a href="https://bugs.openjdk.org/browse/JDK-8196106">JDK-8196106</a> as the reference to the bug.
*
* @param <T> The type of the stream
* @param originalStream The original stream
* @return The stream that is lazily evaluating
*/
public static <T> Stream<T> prepareSortedStreamToWorkInsideOfFlatMapWithTerminalOperations(Stream<T> originalStream) {
return StreamSupport.stream(originalStream.spliterator(), false);
}
}

View File

@ -101,4 +101,37 @@ public class StreamsUtilTest {
Assert.assertEquals(1, numberOfFetchedElements.get());
}
@Test
public void testSortedInsideOfFlatMapShouldRespectTerminalOperation() {
AtomicInteger numberOfFetchedElements = new AtomicInteger();
Stream.of(new Object())
.flatMap(
o -> Stream.of(1, 2, 3).peek(integer -> numberOfFetchedElements.incrementAndGet()))
.limit(1).forEach(NOOP);
Assert.assertEquals(1, numberOfFetchedElements.get());
numberOfFetchedElements.set(0);
Stream.of(new Object())
.flatMap(
o -> Stream.of(1, 2, 3).sorted().peek(integer -> numberOfFetchedElements.incrementAndGet()))
.limit(1).forEach(NOOP);
// Expect actually 1, but always delivery 3 on JDK 21, but will work on JDK 24
// Assert.assertEquals(1, numberOfFetchedElements.get());
numberOfFetchedElements.set(0);
Stream.of(new Object())
.flatMap(
o -> StreamsUtil.prepareSortedStreamToWorkInsideOfFlatMapWithTerminalOperations(Stream.of(1, 2, 3).sorted()).peek(integer -> numberOfFetchedElements.incrementAndGet()))
.limit(1).forEach(NOOP);
// With the workaround it is only 1 as expected
Assert.assertEquals(1, numberOfFetchedElements.get());
}
}