diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillOrderedList.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillOrderedList.java index 03652471a049..0b362d40a666 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillOrderedList.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillOrderedList.java @@ -233,6 +233,20 @@ public Windmill.WorkItemCommitRequest persistDirectly(WindmillStateCache.ForKeyA cleared = false; } + if (!pendingDeletes.isEmpty()) { + for (Range range : pendingDeletes.asRanges()) { + Windmill.TagSortedListDeleteRequest.Builder deletesBuilder = + updatesBuilder.addDeletesBuilder(); + deletesBuilder.setRange( + Windmill.SortedListRange.newBuilder() + .setStart(WindmillTimeUtils.harnessToWindmillTimestamp(range.lowerEndpoint())) + .setLimit(WindmillTimeUtils.harnessToWindmillTimestamp(range.upperEndpoint()))); + deletesBuilder.build(); + idTracker.remove(range); + } + pendingDeletes.clear(); + } + if (!pendingAdds.isEmpty()) { // TODO(reuvenlax): Once we start caching data, we should remove this line. We have it // here now @@ -261,20 +275,6 @@ public Windmill.WorkItemCommitRequest persistDirectly(WindmillStateCache.ForKeyA pendingAdds.clear(); insertBuilder.build(); } - - if (!pendingDeletes.isEmpty()) { - for (Range range : pendingDeletes.asRanges()) { - Windmill.TagSortedListDeleteRequest.Builder deletesBuilder = - updatesBuilder.addDeletesBuilder(); - deletesBuilder.setRange( - Windmill.SortedListRange.newBuilder() - .setStart(WindmillTimeUtils.harnessToWindmillTimestamp(range.lowerEndpoint())) - .setLimit(WindmillTimeUtils.harnessToWindmillTimestamp(range.upperEndpoint()))); - deletesBuilder.build(); - idTracker.remove(range); - } - pendingDeletes.clear(); - } } catch (ExecutionException | InterruptedException e) { throw new RuntimeException(e); } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillStateInternalsTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillStateInternalsTest.java index 87b746089f11..7a68974f4128 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillStateInternalsTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillStateInternalsTest.java @@ -2386,6 +2386,111 @@ public void testOrderedListDeleteRangePersist() { assertBuildable(commitBuilder); } + @Test + public void testOrderedListDeleteRangeThenAddPersist() throws Exception { + SettableFuture, RangeSet>> orderedListFuture = SettableFuture.create(); + orderedListFuture.set(null); + SettableFuture, RangeSet>> deletionsFuture = + SettableFuture.create(); + deletionsFuture.set(null); + when(mockReader.valueFuture( + systemKey(NAMESPACE, "orderedList" + IdTracker.IDS_AVAILABLE_STR), + STATE_FAMILY, + IdTracker.IDS_AVAILABLE_CODER)) + .thenReturn(orderedListFuture); + when(mockReader.valueFuture( + systemKey(NAMESPACE, "orderedList" + IdTracker.DELETIONS_STR), + STATE_FAMILY, + IdTracker.SUBRANGE_DELETIONS_CODER)) + .thenReturn(deletionsFuture); + + StateTag> addr = + StateTags.orderedList("orderedList", StringUtf8Coder.of()); + OrderedListState orderedListState = underTest.state(NAMESPACE, addr); + + // Clear a sub-range [2ms, 4ms) and then add an element at 2ms within that cleared range. + orderedListState.clearRange(Instant.ofEpochMilli(2), Instant.ofEpochMilli(4)); + orderedListState.add(TimestampedValue.of("hello", Instant.ofEpochMilli(2))); + + Windmill.WorkItemCommitRequest.Builder commitBuilder = + Windmill.WorkItemCommitRequest.newBuilder(); + underTest.persist(commitBuilder); + + assertEquals(1, commitBuilder.getSortedListUpdatesCount()); + TagSortedListUpdateRequest updates = commitBuilder.getSortedListUpdates(0); + assertEquals(1, updates.getDeletesCount()); + assertEquals(2000, updates.getDeletes(0).getRange().getStart()); + assertEquals(4000, updates.getDeletes(0).getRange().getLimit()); + assertEquals(1, updates.getInsertsCount()); + assertEquals(1, updates.getInserts(0).getEntriesCount()); + assertEquals("hello", updates.getInserts(0).getEntries(0).getValue().toStringUtf8()); + assertEquals(2000, updates.getInserts(0).getEntries(0).getSortKey()); + assertEquals(IdTracker.NEW_RANGE_MIN_ID, updates.getInserts(0).getEntries(0).getId()); + assertBuildable(commitBuilder); + + // Verify that the sub-range deletion for [2ms, 3ms) was removed when "hello" was added at 2ms, + // so clearing the rest of the 5-minute bucket [0, 2ms) and [4ms, 5m) in the next work item + // does not wipe out id tracking for the bucket. + resetUnderTest(); + orderedListState = underTest.state(NAMESPACE, addr); + orderedListState.clearRange(Instant.EPOCH, Instant.ofEpochMilli(2)); + orderedListState.clearRange(Instant.ofEpochMilli(4), Instant.ofEpochSecond(300)); + + commitBuilder = Windmill.WorkItemCommitRequest.newBuilder(); + underTest.persist(commitBuilder); + assertBuildable(commitBuilder); + + // Adding another element at 2ms in a subsequent work item should get the next id (1), not 0. + resetUnderTest(); + orderedListState = underTest.state(NAMESPACE, addr); + orderedListState.add(TimestampedValue.of("world", Instant.ofEpochMilli(2))); + + commitBuilder = Windmill.WorkItemCommitRequest.newBuilder(); + underTest.persist(commitBuilder); + + updates = commitBuilder.getSortedListUpdates(0); + assertEquals(1, updates.getInsertsCount()); + assertEquals(1, updates.getInserts(0).getEntriesCount()); + assertEquals("world", updates.getInserts(0).getEntries(0).getValue().toStringUtf8()); + assertEquals(2000, updates.getInserts(0).getEntries(0).getSortKey()); + assertEquals(IdTracker.NEW_RANGE_MIN_ID + 1, updates.getInserts(0).getEntries(0).getId()); + assertBuildable(commitBuilder); + + // Now clear the entire 5-minute bucket [0, 5m) and add a new element in the same work item. + // The delete should be processed before the add, resetting the bucket's ids to start at + // NEW_RANGE_MIN_ID (0) for the new element while preserving its id tracking for future adds. + resetUnderTest(); + orderedListState = underTest.state(NAMESPACE, addr); + orderedListState.clearRange(Instant.EPOCH, Instant.ofEpochSecond(300)); + orderedListState.add(TimestampedValue.of("afterFullClear1", Instant.ofEpochMilli(2))); + + commitBuilder = Windmill.WorkItemCommitRequest.newBuilder(); + underTest.persist(commitBuilder); + + updates = commitBuilder.getSortedListUpdates(0); + assertEquals(1, updates.getDeletesCount()); + assertEquals(1, updates.getInsertsCount()); + assertEquals(1, updates.getInserts(0).getEntriesCount()); + assertEquals("afterFullClear1", updates.getInserts(0).getEntries(0).getValue().toStringUtf8()); + assertEquals(IdTracker.NEW_RANGE_MIN_ID, updates.getInserts(0).getEntries(0).getId()); + assertBuildable(commitBuilder); + + // A subsequent add in the same bucket should receive NEW_RANGE_MIN_ID + 1. + resetUnderTest(); + orderedListState = underTest.state(NAMESPACE, addr); + orderedListState.add(TimestampedValue.of("afterFullClear2", Instant.ofEpochMilli(2))); + + commitBuilder = Windmill.WorkItemCommitRequest.newBuilder(); + underTest.persist(commitBuilder); + + updates = commitBuilder.getSortedListUpdates(0); + assertEquals(1, updates.getInsertsCount()); + assertEquals(1, updates.getInserts(0).getEntriesCount()); + assertEquals("afterFullClear2", updates.getInserts(0).getEntries(0).getValue().toStringUtf8()); + assertEquals(IdTracker.NEW_RANGE_MIN_ID + 1, updates.getInserts(0).getEntries(0).getId()); + assertBuildable(commitBuilder); + } + @Test public void testOrderedListMergePendingAdds() { SettableFuture, RangeSet>> orderedListFuture = SettableFuture.create();