Skip to content

Fixes some bugs in StateFetchingIterators. #40716

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
acrites wants to merge 1 commit into
base: master
Choose a base branch
from
+147 −69
Open

Fixes some bugs in StateFetchingIterators. #40716

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
84 changes: 17 additions & 67 deletions sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,10 +18,7 @@
package org.apache.beam.fn.harness.state;

import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull;
import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;

import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
Expand Down Expand Up @@ -49,7 +46,6 @@
import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Throwables;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.AbstractIterator;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.math.LongMath;

Expand Down Expand Up @@ -128,58 +124,9 @@ public UncachedStateIterable(

@Override
public PrefetchableIterator<T> createIterator() {
return new DecodingIterator<T>(
new LazyBlockingStateFetchingIterator(beamFnStateClient, stateRequestForFirstChunk),
valueCoder);
}

private static class DecodingIterator<T> extends AbstractIterator<T>
implements PrefetchableIterator<T> {
private final PrefetchableIterator<ByteString> chunkIterator;
private final Coder<T> valueCoder;
private InputStream currentChunk;

public DecodingIterator(PrefetchableIterator<ByteString> chunkIterator, Coder<T> valueCoder) {
this.chunkIterator = chunkIterator;
this.valueCoder = valueCoder;
this.currentChunk = ByteString.EMPTY.newInput();
}

@Override
protected @Nullable T computeNext() {
try {
while (currentChunk.available() == 0) {
if (chunkIterator.hasNext()) {
currentChunk = chunkIterator.next().newInput();
} else {
return endOfData();
}
}
return valueCoder.decode(currentChunk);
} catch (IOException exn) {
// Should never get here as ByteString.newInput() returns InputStreams
// that don't do actual IO operations.
throw new IllegalStateException(exn);
}
}

@Override
public boolean isReady() {
try {
return currentChunk.available() > 0 || chunkIterator.isReady();
} catch (IOException exn) {
// Should never get here as ByteString.newInput() returns InputStreams
// that don't do actual IO operations.
throw new IllegalStateException(exn);
}
}

@Override
public void prefetch() {
if (!isReady()) {
chunkIterator.prefetch();
}
}
return new DataStreamDecoder<>(
valueCoder,
new LazyBlockingStateFetchingIterator(beamFnStateClient, stateRequestForFirstChunk));
}
}

Expand Down Expand Up @@ -436,6 +383,7 @@ public void remove(Set<Object> toRemoveStructuralValues) {
// returning invalid results as part of a future request.
if (existing.getBlocks().get(existing.getBlocks().size() - 1).getNextToken() != null) {
cache.remove(IterableCacheKey.INSTANCE);
return;
}

// Combine all the individual blocks into one block containing all the values since
Expand Down Expand Up @@ -550,6 +498,7 @@ private void appendHelper(List<T> newValues, long newWeight) {
// returning invalid results as part of a future request.
if (existing.getBlocks().get(existing.getBlocks().size() - 1).getNextToken() != null) {
cache.remove(IterableCacheKey.INSTANCE);
return;
}

// Combine all the individual blocks into one block containing all the values since
Expand Down Expand Up @@ -603,8 +552,7 @@ public boolean isReady() {
@Nullable Blocks<T> existing = cache.peek(IterableCacheKey.INSTANCE);
boolean isFirstBlock = ByteString.EMPTY.equals(nextToken);
if (existing == null) {
// If there is nothing cached and we are on the first block then we are not ready.
return false;
return underlyingStateFetchingIterator.isReady();
} else {
if (isFirstBlock) {
currentBlock = existing.getBlocks().get(0);
Expand All @@ -623,7 +571,7 @@ public boolean isReady() {
currentCachedBlockValueIndex = 0;
} else {
// If not found, then we need to load it.
return false;
return underlyingStateFetchingIterator.isReady();
}
}
}
Expand Down Expand Up @@ -662,12 +610,7 @@ public boolean hasNext() {
}
} else if (isFirstBlock) {
currentBlock = existing.getBlocks().get(0);
} else {
checkState(
existing instanceof BlocksPrefix,
"Unexpected blocks type %s, expected a %s.",
existing.getClass(),
BlocksPrefix.class);
} else if (existing instanceof BlocksPrefix) {
List<Block<T>> blocks = existing.getBlocks();
int currentBlockIndex = 0;
for (; currentBlockIndex < blocks.size(); ++currentBlockIndex) {
Expand All @@ -693,12 +636,16 @@ public boolean hasNext() {
&& Objects.equals(
nextToken,
existing.getBlocks().get(existing.getBlocks().size() - 1).getNextToken())) {
List<Block<T>> newBlocks = new ArrayList<>(currentBlockIndex + 1);
List<Block<T>> newBlocks = new ArrayList<>(existing.getBlocks().size() + 1);
newBlocks.addAll(existing.getBlocks());
newBlocks.add(currentBlock);
cache.put(IterableCacheKey.INSTANCE, new BlocksPrefix<>(newBlocks));
}
}
} else {
// The cache entry was mutated or replaced and is no longer a BlocksPrefix.
// Fall back to loading from state API directly.
currentBlock = loadNextBlock(nextToken);
}
currentCachedBlockValueIndex = 0;
}
Expand Down Expand Up @@ -769,7 +716,10 @@ public void seekToContinuationToken(@Nullable ByteString continuationToken) {
return;
}
this.continuationToken = continuationToken;
this.prefetchedResponse = null;
if (this.prefetchedResponse != null) {
this.prefetchedResponse.cancel(true);
this.prefetchedResponse = null;
}
}

@Override
Expand Down
132 changes: 130 additions & 2 deletions ...va/harness/src/test/java/org/apache/beam/fn/harness/state/StateFetchingIteratorsTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@
import org.apache.beam.sdk.coders.BigEndianIntegerCoder;
import org.apache.beam.sdk.coders.NullableCoder;
import org.apache.beam.sdk.fn.data.WeightedList;
import org.apache.beam.sdk.fn.stream.PrefetchableIterable;
import org.apache.beam.sdk.fn.stream.PrefetchableIterator;
import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
Expand Down Expand Up @@ -145,8 +146,12 @@ public void testAppend() throws Exception {

// Test on a partially accessed iterable
iterable = create(5, expected);
iterable.append(Ints.asList(42, 43));
// Access the iterator to load the first block into the cache so that the cache is partially
// populated with nextToken != null before calling append. Without this call, the cache is
// empty and append would return early at the existing == null check instead of exercising
// the partial cache eviction logic.
Boolean ignored = iterable.iterator().hasNext();
iterable.append(Ints.asList(42, 43));
// The append is dropped since the entire iterable isn't in memory so next time we iterate we
// expect to be given what the runner tells us is there.
verifyFetch(iterable.iterator(), expected);
Expand Down Expand Up @@ -176,8 +181,12 @@ public void testRemove() throws Exception {

// Test on a partially accessed iterable
iterable = create(5, expected);
iterable.remove(toRemove);
// Access the iterator to load the first block into the cache so that the cache is partially
// populated with nextToken != null before calling remove. Without this call, the cache is
// empty and remove would return early at the existing == null check instead of exercising
// the partial cache eviction logic.
Boolean ignored = iterable.iterator().hasNext();
iterable.remove(toRemove);
// The remove is dropped since the entire iterable isn't in memory so next time we iterate we
// expect to be given what the runner tells us is there.
verifyFetch(iterable.iterator(), expected);
Expand Down Expand Up @@ -260,6 +269,32 @@ public void testCacheEvictionOrphansIteratorAndAllowsForIteratorToRejoin() throw
assertEquals(stateRequestCount, fakeStateClient.getCallCount());
}

@Test
public void testIsReadyBecomesTrueAfterPrefetch() throws Exception {
int[] expected = new int[] {0, 1, 2, 3};
// Chunk size 4 means 1 integer per chunk
CachingStateIterable<Integer> iterable = create(4, expected);
PrefetchableIterator<Integer> iterator = iterable.iterator();
assertFalse(iterator.isReady());
iterator.prefetch();
// After prefetch completes, isReady() must return true according to PrefetchableIterator
// contract
assertTrue(iterator.isReady());
assertEquals(0, (int) iterator.next());
}

@Test
public void testMutationDuringIterationDoesNotCrash() throws Exception {
int[] expected = new int[] {0, 1, 2, 3};
CachingStateIterable<Integer> iterable = create(4, expected);
PrefetchableIterator<Integer> iterator = iterable.iterator();
assertEquals(0, (int) iterator.next());
// Mutate while iteration is partially through
iterable.clearAndAppend(Ints.asList(99));
// hasNext() must not throw IllegalStateException on MutatedBlocks
assertTrue(iterator.hasNext());
}

@Test
public void testBlocksPrefixShrinkage() throws Exception {
List<Block<String>> originalBlocks =
Expand Down Expand Up @@ -503,6 +538,18 @@ public CompletableFuture<StateResponse> handle(StateRequest.Builder requestBuild
assertEquals(ByteString.copyFromUtf8("value" + 44), byteStrings.next());
}

@Test
public void testSeekCancelsInFlightPrefetch() throws Exception {
CompletableFuture<StateResponse> future = new CompletableFuture<>();
BeamFnStateClient fakeStateClient = requestBuilder -> future;
LazyBlockingStateFetchingIterator byteStrings =
new LazyBlockingStateFetchingIterator(fakeStateClient, StateRequest.getDefaultInstance());
byteStrings.prefetch();
assertFalse(future.isCancelled());
byteStrings.seekToContinuationToken(ByteString.copyFromUtf8("other"));
assertTrue(future.isCancelled());
}

private void testFetch(ByteString... expected) {
AtomicInteger callCount = new AtomicInteger();
BeamFnStateClient fakeStateClient = fakeStateClient(callCount, expected);
Expand All @@ -529,4 +576,85 @@ private void testFetch(ByteString... expected) {
assertEquals(Arrays.asList(expected), results);
}
}

/** Tests for {@link StateFetchingIterators#readAllAndDecodeStartingFrom} without caching. */
@RunWith(JUnit4.class)
public static class UncachedStateIterableTest {

@Test
public void testEmpty() throws Exception {
testFetch(4);
}

@Test
public void testNonEmpty() throws Exception {
testFetch(4, 0);
}

@Test
public void testMultipleElementsPerChunk() throws Exception {
testFetch(8, 0, 1, 2, 3, 4, 5);
}

@Test
public void testSingleElementPerChunk() throws Exception {
testFetch(4, 0, 1, 2, 3, 4, 5);
}

@Test
public void testChunkSmallerThanElementSize() throws Exception {
testFetch(3, 0, 1, 2, 3, 4, 5);
}

@Test
public void testChunkLargerThanElementSize() throws Exception {
testFetch(5, 0, 1, 2, 3, 4, 5);
}

@Test
public void testIsReadyAndPrefetch() throws Exception {
AtomicInteger callCount = new AtomicInteger();
BeamFnStateClient fakeStateClient =
fakeStateClient(callCount, ByteString.copyFromUtf8("BC"), ByteString.copyFromUtf8("DEF"));
PrefetchableIterable<byte[]> iterable =
StateFetchingIterators.readAllAndDecodeStartingFrom(
fakeStateClient,
StateRequest.getDefaultInstance(),
org.apache.beam.sdk.coders.ByteArrayCoder.of());
// Not yet iterated or prefetched
PrefetchableIterator<byte[]> iterator = iterable.iterator();
assertFalse(iterator.isReady());
}

private void testFetch(int chunkSize, int... expected) {
StateRequest requestForFirstChunk =
StateRequest.newBuilder()
.setStateKey(
StateKey.newBuilder()
.setBagUserState(
StateKey.BagUserState.newBuilder()
.setTransformId("transformId")
.setUserStateId("stateId")
.setKey(ByteString.copyFromUtf8("key"))
.setWindow(ByteString.copyFromUtf8("window"))))
.setGet(StateGetRequest.getDefaultInstance())
.build();
FakeBeamFnStateClient fakeStateClient =
new FakeBeamFnStateClient(
BigEndianIntegerCoder.of(),
ImmutableMap.of(requestForFirstChunk.getStateKey(), Ints.asList(expected)),
chunkSize);

PrefetchableIterable<Integer> iterable =
StateFetchingIterators.readAllAndDecodeStartingFrom(
fakeStateClient, requestForFirstChunk, BigEndianIntegerCoder.of());

PrefetchableIterator<Integer> iterator = iterable.iterator();
List<Integer> results = new ArrayList<>();
while (iterator.hasNext()) {
results.add(iterator.next());
}
assertEquals(Ints.asList(expected), results);
}
}
}
Loading