mirror of
https://github.com/spring-projects/spring-data-elasticsearch.git
synced 2025-06-01 09:42:11 +00:00
Automatically close scroll context when returning streamed results.
Original Pull Request #1746 Closes #1745
This commit is contained in:
parent
3500dad2bc
commit
13ab2b9e95
@ -66,10 +66,14 @@ abstract class StreamQueries {
|
|||||||
private volatile Iterator<SearchHit<T>> currentScrollHits = searchHits.iterator();
|
private volatile Iterator<SearchHit<T>> currentScrollHits = searchHits.iterator();
|
||||||
private volatile boolean continueScroll = currentScrollHits.hasNext();
|
private volatile boolean continueScroll = currentScrollHits.hasNext();
|
||||||
private volatile ScrollState scrollState = new ScrollState(searchHits.getScrollId());
|
private volatile ScrollState scrollState = new ScrollState(searchHits.getScrollId());
|
||||||
|
private volatile boolean isClosed = false;
|
||||||
|
|
||||||
@Override
|
@Override
|
||||||
public void close() {
|
public void close() {
|
||||||
clearScrollConsumer.accept(scrollState.getScrollIds());
|
if (!isClosed) {
|
||||||
|
clearScrollConsumer.accept(scrollState.getScrollIds());
|
||||||
|
isClosed = true;
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@Override
|
@Override
|
||||||
@ -96,18 +100,24 @@ abstract class StreamQueries {
|
|||||||
@Override
|
@Override
|
||||||
public boolean hasNext() {
|
public boolean hasNext() {
|
||||||
|
|
||||||
if (!continueScroll || (maxCount > 0 && currentCount.get() >= maxCount)) {
|
boolean hasNext = false;
|
||||||
return false;
|
|
||||||
|
if (!isClosed && continueScroll && (maxCount <= 0 || currentCount.get() < maxCount)) {
|
||||||
|
|
||||||
|
if (!currentScrollHits.hasNext()) {
|
||||||
|
SearchScrollHits<T> nextPage = continueScrollFunction.apply(scrollState.getScrollId());
|
||||||
|
currentScrollHits = nextPage.iterator();
|
||||||
|
scrollState.updateScrollId(nextPage.getScrollId());
|
||||||
|
continueScroll = currentScrollHits.hasNext();
|
||||||
|
}
|
||||||
|
hasNext = currentScrollHits.hasNext();
|
||||||
}
|
}
|
||||||
|
|
||||||
if (!currentScrollHits.hasNext()) {
|
if (!hasNext) {
|
||||||
SearchScrollHits<T> nextPage = continueScrollFunction.apply(scrollState.getScrollId());
|
close();
|
||||||
currentScrollHits = nextPage.iterator();
|
|
||||||
scrollState.updateScrollId(nextPage.getScrollId());
|
|
||||||
continueScroll = currentScrollHits.hasNext();
|
|
||||||
}
|
}
|
||||||
|
|
||||||
return currentScrollHits.hasNext();
|
return hasNext;
|
||||||
}
|
}
|
||||||
|
|
||||||
@Override
|
@Override
|
||||||
|
@ -24,6 +24,7 @@ import java.util.Iterator;
|
|||||||
import java.util.List;
|
import java.util.List;
|
||||||
import java.util.concurrent.atomic.AtomicBoolean;
|
import java.util.concurrent.atomic.AtomicBoolean;
|
||||||
|
|
||||||
|
import org.junit.jupiter.api.DisplayName;
|
||||||
import org.junit.jupiter.api.Test;
|
import org.junit.jupiter.api.Test;
|
||||||
import org.springframework.data.util.StreamUtils;
|
import org.springframework.data.util.StreamUtils;
|
||||||
|
|
||||||
@ -39,6 +40,8 @@ public class StreamQueriesTest {
|
|||||||
// given
|
// given
|
||||||
List<SearchHit<String>> hits = new ArrayList<>();
|
List<SearchHit<String>> hits = new ArrayList<>();
|
||||||
hits.add(getOneSearchHit());
|
hits.add(getOneSearchHit());
|
||||||
|
hits.add(getOneSearchHit());
|
||||||
|
hits.add(getOneSearchHit());
|
||||||
|
|
||||||
SearchScrollHits<String> searchHits = newSearchScrollHits(hits, "1234");
|
SearchScrollHits<String> searchHits = newSearchScrollHits(hits, "1234");
|
||||||
|
|
||||||
@ -51,9 +54,7 @@ public class StreamQueriesTest {
|
|||||||
scrollId -> newSearchScrollHits(Collections.emptyList(), scrollId), //
|
scrollId -> newSearchScrollHits(Collections.emptyList(), scrollId), //
|
||||||
scrollIds -> clearScrollCalled.set(true));
|
scrollIds -> clearScrollCalled.set(true));
|
||||||
|
|
||||||
while (iterator.hasNext()) {
|
iterator.next();
|
||||||
iterator.next();
|
|
||||||
}
|
|
||||||
iterator.close();
|
iterator.close();
|
||||||
|
|
||||||
// then
|
// then
|
||||||
@ -61,6 +62,27 @@ public class StreamQueriesTest {
|
|||||||
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@Test // #1745
|
||||||
|
@DisplayName("should call clearScroll when no more data is available")
|
||||||
|
void shouldCallClearScrollWhenNoMoreDataIsAvailable() {
|
||||||
|
|
||||||
|
List<SearchHit<String>> hits = new ArrayList<>();
|
||||||
|
hits.add(getOneSearchHit());
|
||||||
|
SearchScrollHits<String> searchHits = newSearchScrollHits(hits, "1234");
|
||||||
|
AtomicBoolean clearScrollCalled = new AtomicBoolean(false);
|
||||||
|
|
||||||
|
SearchHitsIterator<String> iterator = StreamQueries.streamResults( //
|
||||||
|
0, //
|
||||||
|
searchHits, //
|
||||||
|
scrollId -> newSearchScrollHits(Collections.emptyList(), scrollId), //
|
||||||
|
scrollIds -> clearScrollCalled.set(true));
|
||||||
|
|
||||||
|
while (iterator.hasNext()) {
|
||||||
|
iterator.next();
|
||||||
|
}
|
||||||
|
|
||||||
|
assertThat(clearScrollCalled).isTrue();
|
||||||
|
}
|
||||||
private SearchHit<String> getOneSearchHit() {
|
private SearchHit<String> getOneSearchHit() {
|
||||||
return new SearchHit<String>(null, null, null, 0, null, null, null, null, null, null, "one");
|
return new SearchHit<String>(null, null, null, 0, null, null, null, null, null, null, "one");
|
||||||
}
|
}
|
||||||
|
Loading…
x
Reference in New Issue
Block a user