From a9ea216ef6a16fe51679e615e204b88ef7e9953c Mon Sep 17 00:00:00 2001 From: Andrew Crites Date: Wed, 21 Oct 2026 17:12:54 -0700 Subject: [PATCH] Fixes some bugs in StateFetchingIterators. --- .../harness/state/StateFetchingIterators.java | 84 +++-------- .../state/StateFetchingIteratorsTest.java | 132 +++++++++++++++++- 2 files changed, 147 insertions(+), 69 deletions(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java index 1b4d7ec371ba..c51a83795c0c 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java @@ -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; @@ -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; @@ -128,58 +124,9 @@ public UncachedStateIterable( @Override public PrefetchableIterator createIterator() { - return new DecodingIterator( - new LazyBlockingStateFetchingIterator(beamFnStateClient, stateRequestForFirstChunk), - valueCoder); - } - - private static class DecodingIterator extends AbstractIterator - implements PrefetchableIterator { - private final PrefetchableIterator chunkIterator; - private final Coder valueCoder; - private InputStream currentChunk; - - public DecodingIterator(PrefetchableIterator chunkIterator, Coder 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)); } } @@ -436,6 +383,7 @@ public void remove(Set 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 @@ -550,6 +498,7 @@ private void appendHelper(List 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 @@ -603,8 +552,7 @@ public boolean isReady() { @Nullable Blocks 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); @@ -623,7 +571,7 @@ public boolean isReady() { currentCachedBlockValueIndex = 0; } else { // If not found, then we need to load it. - return false; + return underlyingStateFetchingIterator.isReady(); } } } @@ -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> blocks = existing.getBlocks(); int currentBlockIndex = 0; for (; currentBlockIndex < blocks.size(); ++currentBlockIndex) { @@ -693,12 +636,16 @@ public boolean hasNext() { && Objects.equals( nextToken, existing.getBlocks().get(existing.getBlocks().size() - 1).getNextToken())) { - List> newBlocks = new ArrayList<>(currentBlockIndex + 1); + List> 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; } @@ -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 diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/StateFetchingIteratorsTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/StateFetchingIteratorsTest.java index 43087cabbd37..2e374a543efe 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/StateFetchingIteratorsTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/StateFetchingIteratorsTest.java @@ -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; @@ -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); @@ -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); @@ -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 iterable = create(4, expected); + PrefetchableIterator 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 iterable = create(4, expected); + PrefetchableIterator 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> originalBlocks = @@ -503,6 +538,18 @@ public CompletableFuture handle(StateRequest.Builder requestBuild assertEquals(ByteString.copyFromUtf8("value" + 44), byteStrings.next()); } + @Test + public void testSeekCancelsInFlightPrefetch() throws Exception { + CompletableFuture 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); @@ -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 iterable = + StateFetchingIterators.readAllAndDecodeStartingFrom( + fakeStateClient, + StateRequest.getDefaultInstance(), + org.apache.beam.sdk.coders.ByteArrayCoder.of()); + // Not yet iterated or prefetched + PrefetchableIterator 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 iterable = + StateFetchingIterators.readAllAndDecodeStartingFrom( + fakeStateClient, requestForFirstChunk, BigEndianIntegerCoder.of()); + + PrefetchableIterator iterator = iterable.iterator(); + List results = new ArrayList<>(); + while (iterator.hasNext()) { + results.add(iterator.next()); + } + assertEquals(Ints.asList(expected), results); + } + } }