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..174136d779b7 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 @@ -436,6 +436,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 +551,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 diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BagUserStateTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BagUserStateTest.java index 05a280245d13..952fe0a8cd1f 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BagUserStateTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BagUserStateTest.java @@ -196,6 +196,45 @@ public void testAppendCached() throws Exception { assertEquals(encode("A1", "A2", "A3", "A4", "A5"), fakeClient.getData().get(key("A"))); } + @Test + public void testPartialIterationThenAppendDoesNotTruncateCache() throws Exception { + // Initial data contains multiple elements across multiple chunks. + FakeBeamFnStateClient fakeClient = + new FakeBeamFnStateClient( + StringUtf8Coder.of(), ImmutableMap.of(key("A"), asList("A1", "A2", "A3", "A4")), 2); + + Cache cache = Caches.eternal(); + { + // First user state: simulates BagState.isEmpty().read() or partial iteration, followed by + // append. + BagUserState userState = + new BagUserState<>(cache, fakeClient, "instructionId", key("A"), StringUtf8Coder.of()); + + // Calling iterator().hasNext() only loads the first chunk into the cache with nextToken != + // null. + assertTrue(userState.get().iterator().hasNext()); + + // User appends a new element. + userState.append("A5"); + + // asyncClose persists the append to runner and updates the cache. + userState.asyncClose(); + } + + { + // The next user state accesses the state for the same key. + BagUserState userState = + new BagUserState<>(cache, fakeClient, "instructionId", key("A"), StringUtf8Coder.of()); + + // All elements (A1, A2, A3, A4, A5) must be present; unread pages (A2, A3, A4) must not be + // lost. + assertArrayEquals( + new String[] {"A1", "A2", "A3", "A4", "A5"}, + Iterables.toArray(userState.get(), String.class)); + userState.asyncClose(); + } + } + @Test public void testClear() throws Exception { FakeBeamFnStateClient fakeClient = 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..ee57ea6bbc66 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 @@ -145,8 +145,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 +180,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);