diff --git a/CHANGES.md b/CHANGES.md index 73966a48313c..01516a7c11b3 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -69,6 +69,7 @@ ## New Features / Improvements * X feature added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). +* (Java) `Watch.growthOf` can bound its deduplication state by event time with `withTimestampCursor`, which retires an output key once the greatest emitted timestamp has moved more than the allowed lateness past it ([#18459](https://github.com/apache/beam/issues/18459)). ## Breaking Changes diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Watch.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Watch.java index 793fac048dff..736839a2ed3e 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Watch.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Watch.java @@ -120,6 +120,11 @@ * Growth.PollResult#withWatermark} if the {@link Growth.PollFn} can provide a more optimistic * estimate. * + *

By default the transform remembers the key of every output it has emitted, so the state of an + * input that is watched indefinitely grows without bound. {@link Growth#withTimestampCursor} bounds + * that state by event time, for a {@link Growth.PollFn} whose outputs arrive in roughly + * non-decreasing timestamp order. + * *

Note: This transform works only in runners supporting Splittable DoFn: see capability matrix. */ @@ -677,6 +682,8 @@ public String toString(KV state) { abstract @Nullable Coder getOutputCoder(); + abstract @Nullable Duration getTimestampCursorAllowedLateness(); + abstract Builder toBuilder(); @AutoValue.Builder @@ -695,6 +702,9 @@ abstract Builder setTerminationPerInput( abstract Builder setOutputCoder(Coder outputCoder); + abstract Builder setTimestampCursorAllowedLateness( + Duration allowedLateness); + abstract Growth build(); } @@ -728,6 +738,38 @@ public Growth withOutputCoder(Coder outputCoder) return toBuilder().setOutputCoder(outputCoder).build(); } + /** Like {@link #withTimestampCursor(Duration)} with no lateness allowed. */ + public Growth withTimestampCursor() { + return withTimestampCursor(Duration.ZERO); + } + + /** + * Bounds the deduplication state by event time. + * + *

Deduplication still goes by output key, but a key is retired once the greatest timestamp + * emitted for this input has moved more than {@code allowedLateness} past it, so the state + * holds a trailing window rather than every key ever seen. + * + *

An output whose timestamp is below that mark is taken as already seen and is dropped, so + * this suits a {@link PollFn} whose outputs arrive in roughly non-decreasing timestamp order. + * Widen {@code allowedLateness} for a source that reports outputs further out of order, at the + * cost of a larger state. + * + *

Retiring a key costs the guarantee that an output is emitted once for all time. An output + * that a {@link PollFn} reports again at or above the current floor after its key was retired + * looks new and is emitted a second time. Updating a running pipeline to widen {@code + * allowedLateness}, or to drop the cursor altogether, lowers the floor over keys that are + * already gone and can emit them again for the same reason. + */ + public Growth withTimestampCursor(Duration allowedLateness) { + checkArgument(allowedLateness != null, "allowedLateness can not be null"); + checkArgument( + !allowedLateness.isShorterThan(Duration.ZERO), + "allowedLateness must not be negative, but was %s", + allowedLateness); + return toBuilder().setTimestampCursorAllowedLateness(allowedLateness).build(); + } + @Override public PCollection> expand(PCollection input) { checkNotNull(getPollInterval(), "pollInterval"); @@ -899,20 +941,34 @@ public ProcessContinuation process( return stop(); } + PollingGrowthState pollingRestriction = + (PollingGrowthState) currentRestriction; + + @Nullable Duration allowedLateness = spec.getTimestampCursorAllowedLateness(); + @Nullable Instant cursor = pollingRestriction.getCursor(); + if (retentionFloorAtMaxTimestamp(cursor, allowedLateness)) { + // Nothing can be claimed above the floor, so claim an empty round and stop. + LOG.info("{} - will not poll, retention floor is already at max timestamp.", c.element()); + tracker.tryClaim( + KV.of( + PollResult.incomplete(Collections.emptyList()), + pollingRestriction.getTerminationState())); + return stop(); + } + // Poll for additional elements. Instant now = Instant.now(); Growth.PollResult res = spec.getPollFn().getClosure().apply(c.element(), wrapProcessContext(c)); - PollingGrowthState pollingRestriction = - (PollingGrowthState) currentRestriction; // Produce a poll result that only contains never seen before results in timestamp // sorted order. Growth.PollResult newResults = computeNeverSeenBeforeResults(pollingRestriction, res); // If we had zero new results, attempt to update the watermark if the poll result - // provided a watermark. Otherwise attempt to claim all pending outputs. + // provided a watermark or the retention floor bounds future outputs. Otherwise attempt + // to claim all pending outputs. LOG.info( "{} - current round of polling took {} ms and returned {} results, " + "of which {} were new.", @@ -944,6 +1000,26 @@ public ProcessContinuation process( // computeNeverSeenBeforeResults returns the elements in timestamp sorted order so // we can get the timestamp from the first element. computedWatermark = newResults.getOutputs().get(0).getTimestamp(); + } else if (allowedLateness != null && cursor != null) { + // Nothing below the retention floor is ever emitted, so a round with no new results and + // no explicit watermark can still advance the watermark to the floor. + computedWatermark = retentionFloor(cursor, allowedLateness); + } + + if (allowedLateness != null && !newResults.getOutputs().isEmpty()) { + // The cursor only ever advances, and lands on the greatest timestamp emitted so far. Once + // it carries the retention floor to the maximum timestamp, every later output falls below + // the floor and would be dropped, so polling stops. + Instant newCursor = + Ordering.natural() + .nullsFirst() + .max( + cursor, + newResults.getOutputs().get(newResults.getOutputs().size() - 1).getTimestamp()); + if (retentionFloorAtMaxTimestamp(newCursor, allowedLateness)) { + LOG.info("{} - will stop polling, retention floor reached max timestamp.", c.element()); + return stop(); + } } Instant currentTime = Instant.now(); @@ -979,8 +1055,14 @@ private Growth.PollResult computeNeverSeenBeforeResults( // Collect results to include as newly pending. Note that the poll result may in theory // contain multiple outputs mapping to the same output key - we need to ignore duplicates // here already. + Instant retentionFloor = retentionFloor(state, spec.getTimestampCursorAllowedLateness()); Map> newPending = Maps.newHashMap(); for (TimestampedValue output : pollResult.getOutputs()) { + if (retentionFloor != null && output.getTimestamp().isBefore(retentionFloor)) { + // The key that would prove this output already seen has been retired, so treat the + // output as seen. + continue; + } OutputT value = output.getValue(); HashCode hash = hash128(value); if (state.getCompleted().containsKey(hash) || newPending.containsKey(hash)) { @@ -989,8 +1071,8 @@ private Growth.PollResult computeNeverSeenBeforeResults( // TODO (https://github.com/apache/beam/issues/18459): // Consider adding only at most N pending elements and ignoring others, // instead relying on future poll rounds to provide them, in order to avoid - // blowing up the state. Combined with garbage collection of PollingGrowthState.completed, - // this would make the transform scalable to very large poll results. + // blowing up the state. Combined with a timestamp cursor, this would make the transform + // scalable to very large poll results. newPending.put(hash, output); } @@ -1012,7 +1094,8 @@ public GrowthState getInitialRestriction(@Element InputT element) { @NewTracker public GrowthTracker newTracker( @Restriction GrowthState restriction) { - return new GrowthTracker<>(restriction, coderFunnel); + return new GrowthTracker<>( + restriction, coderFunnel, spec.getTimestampCursorAllowedLateness()); } @GetRestrictionCoder @@ -1026,6 +1109,43 @@ public Coder getRestrictionCoder() { /** A base class for all restrictions related to the {@link Growth} SplittableDoFn. */ abstract static class GrowthState {} + /** + * The timestamp below which a key is retired from {@link PollingGrowthState#getCompleted}, or + * null when every key is retained. + * + *

An output at or above the floor is still deduplicated by key; one below it is taken as + * already seen. + */ + private static @Nullable Instant retentionFloor( + PollingGrowthState state, @Nullable Duration allowedLateness) { + if (allowedLateness == null || state.getCursor() == null) { + return null; + } + return retentionFloor(state.getCursor(), allowedLateness); + } + + /** The retention floor for a cursor, saturated at the minimum timestamp. */ + private static Instant retentionFloor(Instant cursor, Duration allowedLateness) { + long floorMillis; + try { + floorMillis = Math.subtractExact(cursor.getMillis(), allowedLateness.getMillis()); + } catch (ArithmeticException e) { + floorMillis = BoundedWindow.TIMESTAMP_MIN_VALUE.getMillis(); + } + return new Instant(Math.max(floorMillis, BoundedWindow.TIMESTAMP_MIN_VALUE.getMillis())); + } + + /** + * Whether the retention floor has reached the maximum timestamp, after which no further output + * can be claimed and polling on would only drop outputs. + */ + private static boolean retentionFloorAtMaxTimestamp( + @Nullable Instant cursor, @Nullable Duration allowedLateness) { + return cursor != null + && allowedLateness != null + && !retentionFloor(cursor, allowedLateness).isBefore(BoundedWindow.TIMESTAMP_MAX_VALUE); + } + /** * Stores the prior pending poll results related to the {@link Growth} SplittableDoFn. Used to * represent the primary restriction during checkpoint which should be replayed if the primary @@ -1055,27 +1175,40 @@ public static NonPollingGrowthState of(Growth.PollResult extends GrowthState { public static PollingGrowthState of( TerminationStateT terminationState) { - return new AutoValue_Watch_PollingGrowthState<>(ImmutableMap.of(), null, terminationState); + return new AutoValue_Watch_PollingGrowthState<>( + ImmutableMap.of(), null, terminationState, null); } public static PollingGrowthState of( ImmutableMap completed, Instant pollWatermark, TerminationStateT terminationState) { - return new AutoValue_Watch_PollingGrowthState<>(completed, pollWatermark, terminationState); + return of(completed, pollWatermark, terminationState, null); + } + + public static PollingGrowthState of( + ImmutableMap completed, + @Nullable Instant pollWatermark, + TerminationStateT terminationState, + @Nullable Instant cursor) { + return new AutoValue_Watch_PollingGrowthState<>( + completed, pollWatermark, terminationState, cursor); } // Hashes and timestamps of outputs that have already been output and should be omitted - // from future polls. Timestamps are preserved to allow garbage-collecting this state - // in the future, e.g. dropping elements from "completed" and from - // computeNeverSeenBeforeResults() if their timestamp is more than X behind the watermark. - // As of writing, we don't do this, but preserve the information for forward compatibility - // in case of pipeline update. TODO: do this. + // from future polls. Under a timestamp cursor the entries the cursor has moved past are + // dropped, which bounds this map; otherwise every key ever seen is kept. public abstract ImmutableMap getCompleted(); public abstract @Nullable Instant getPollWatermark(); public abstract TerminationStateT getTerminationState(); + + /** + * The greatest timestamp emitted for this input so far, or null when the transform is not + * bounding its state by event time or has emitted nothing yet. + */ + public abstract @Nullable Instant getCursor(); } @VisibleForTesting @@ -1088,6 +1221,10 @@ static class GrowthTracker // Used to hash values. private final Funnel coderFunnel; + // How far below the cursor a completed key is still retained, or null when the state is not + // bounded by event time. + private final @Nullable Duration allowedLateness; + // non-null after first successful tryClaim() private Growth.@Nullable PollResult claimedPollResult; private @Nullable TerminationStateT claimedTerminationState; @@ -1098,9 +1235,11 @@ static class GrowthTracker // Whether we should stop claiming poll results. private boolean shouldStop; - GrowthTracker(GrowthState state, Funnel coderFunnel) { + GrowthTracker( + GrowthState state, Funnel coderFunnel, @Nullable Duration allowedLateness) { this.state = state; this.coderFunnel = coderFunnel; + this.allowedLateness = allowedLateness; this.shouldStop = false; } @@ -1135,13 +1274,30 @@ public SplitResult trySplit(double fractionOfRemainder) { ImmutableMap.Builder newCompleted = ImmutableMap.builder(); newCompleted.putAll(currentState.getCompleted()); newCompleted.putAll(claimedHashes); + ImmutableMap completed = newCompleted.build(); + + // A round that is not bounding the state retains every key, so it drops the cursor that + // would retire them and returns the restriction to the pre-cursor format. + Instant cursor = null; + if (allowedLateness != null) { + // The cursor only ever advances, and retires the keys it has moved past. + cursor = currentState.getCursor(); + for (Instant timestamp : claimedHashes.values()) { + cursor = Ordering.natural().nullsFirst().max(cursor, timestamp); + } + if (cursor != null) { + completed = retainAtOrAfter(completed, retentionFloor(cursor, allowedLateness)); + } + } + residual = PollingGrowthState.of( - newCompleted.build(), + completed, Ordering.natural() .nullsFirst() .max(currentState.getPollWatermark(), claimedPollResult.watermark), - claimedTerminationState); + claimedTerminationState, + cursor); state = NonPollingGrowthState.of(claimedPollResult); } @@ -1153,6 +1309,18 @@ private HashCode hash128(OutputT value) { return Hashing.murmur3_128().hashObject(value, coderFunnel); } + /** Drops the completed keys the retention floor has retired, which bounds the state. */ + private static ImmutableMap retainAtOrAfter( + ImmutableMap completed, Instant floor) { + ImmutableMap.Builder retained = ImmutableMap.builder(); + for (Map.Entry entry : completed.entrySet()) { + if (!entry.getValue().isBefore(floor)) { + retained.put(entry); + } + } + return retained.build(); + } + @Override public void checkDone() throws IllegalStateException { checkState( @@ -1181,11 +1349,22 @@ public boolean tryClaim(KV, TerminationStateT> pollRe ImmutableMap newClaimedHashes = newClaimedHashesBuilder.build(); if (state instanceof PollingGrowthState) { + PollingGrowthState pollingState = (PollingGrowthState) state; // If we have previously claimed one of these hashes then return false. if (!Collections.disjoint( - newClaimedHashes.keySet(), ((PollingGrowthState) state).getCompleted().keySet())) { + newClaimedHashes.keySet(), pollingState.getCompleted().keySet())) { return false; } + // An output the cursor has already moved past cannot be claimed, since the key that would + // prove it never seen before has been retired. + Instant retentionFloor = retentionFloor(pollingState, allowedLateness); + if (retentionFloor != null) { + for (Instant timestamp : newClaimedHashes.values()) { + if (timestamp.isBefore(retentionFloor)) { + return false; + } + } + } } else { Set expectedHashesToClaim = new HashSet<>(); for (TimestampedValue value : @@ -1249,6 +1428,7 @@ static class GrowthStateCoder extends StructuredCode private static final int POLLING_GROWTH_STATE = 0; private static final int NON_POLLING_GROWTH_STATE = 1; + private static final int CURSOR_POLLING_GROWTH_STATE = 2; public static GrowthStateCoder of( Coder outputCoder, Coder terminationStateCoder) { @@ -1259,6 +1439,7 @@ public static GrowthStateCoder NULLABLE_INSTANT_CODER = NullableCoder.of(InstantCoder.of()); + private static final Coder INSTANT_CODER = InstantCoder.of(); private final Coder outputCoder; private final Coder>> timestampedOutputCoder; @@ -1275,8 +1456,17 @@ private GrowthStateCoder( @Override public void encode(GrowthState value, OutputStream os) throws IOException { if (value instanceof PollingGrowthState) { - VarInt.encode(POLLING_GROWTH_STATE, os); - encodePollingGrowthState((PollingGrowthState) value, os); + PollingGrowthState polling = + (PollingGrowthState) value; + // A state without a cursor keeps the pre-cursor byte format. + if (polling.getCursor() == null) { + VarInt.encode(POLLING_GROWTH_STATE, os); + encodePollingGrowthState(polling, os); + } else { + VarInt.encode(CURSOR_POLLING_GROWTH_STATE, os); + encodePollingGrowthState(polling, os); + INSTANT_CODER.encode(polling.getCursor(), os); + } } else if (value instanceof NonPollingGrowthState) { VarInt.encode(NON_POLLING_GROWTH_STATE, os); encodeNonPollingGrowthState((NonPollingGrowthState) value, os); @@ -1305,7 +1495,9 @@ public GrowthState decode(InputStream is) throws IOException { case NON_POLLING_GROWTH_STATE: return decodeNonPollingGrowthState(is); case POLLING_GROWTH_STATE: - return decodePollingGrowthState(is); + return decodePollingGrowthState(is, false); + case CURSOR_POLLING_GROWTH_STATE: + return decodePollingGrowthState(is, true); default: throw new IOException("Unknown growth state type " + type); } @@ -1317,11 +1509,14 @@ private GrowthState decodeNonPollingGrowthState(InputStream is) throws IOExcepti return NonPollingGrowthState.of(new Growth.PollResult<>(values, watermark)); } - private GrowthState decodePollingGrowthState(InputStream is) throws IOException { + private GrowthState decodePollingGrowthState(InputStream is, boolean hasCursor) + throws IOException { TerminationStateT terminationState = terminationStateCoder.decode(is); Instant watermark = NULLABLE_INSTANT_CODER.decode(is); Map completed = COMPLETED_CODER.decode(is); - return PollingGrowthState.of(ImmutableMap.copyOf(completed), watermark, terminationState); + Instant cursor = hasCursor ? INSTANT_CODER.decode(is) : null; + return PollingGrowthState.of( + ImmutableMap.copyOf(completed), watermark, terminationState, cursor); } @Override diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/WatchTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/WatchTest.java index 277d49a7240d..ae32c6314e9e 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/WatchTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/WatchTest.java @@ -28,9 +28,12 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; import java.io.IOException; import java.io.Serializable; @@ -58,6 +61,7 @@ import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator; import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.CoderUtils; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PCollection; import org.apache.beam.sdk.values.PCollectionView; @@ -69,6 +73,7 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.hash.Funnels; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.hash.HashCode; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.hash.Hashing; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.io.BaseEncoding; import org.joda.time.Duration; import org.joda.time.Instant; import org.joda.time.ReadableDuration; @@ -183,6 +188,36 @@ private void testMultiplePolls(boolean terminationConditionElapsesBeforeOutputIs p.run(); } + @Test + @Category({NeedsRunner.class, UsesUnboundedSplittableParDo.class}) + public void testMultiplePollsWithTimestampCursor() { + List all = Arrays.asList(0, 1, 2, 3, 4, 5, 6, 7, 8, 9); + + PCollection res = + p.apply(Create.of("a")) + .apply( + Watch.growthOf( + new StablyTimedPollFn( + all, standardSeconds(3) /* timeToOutputEverything */)) + .withPollInterval(Duration.millis(300)) + .withTimestampCursor() + .withOutputCoder(VarIntCoder.of())) + .apply("Drop input", Values.create()); + + PAssert.that(res).containsInAnyOrder(all); + + p.run(); + } + + @Test + public void testTimestampCursorRejectsNegativeAllowedLateness() { + Watch.Growth growth = + Watch.growthOf( + new StablyTimedPollFn(Arrays.asList(0), standardSeconds(1))); + assertThrows( + IllegalArgumentException.class, () -> growth.withTimestampCursor(standardSeconds(-1))); + } + @Test @Category({NeedsRunner.class, UsesUnboundedSplittableParDo.class}) public void testMultiplePollsWithKeyExtractor() { @@ -323,6 +358,45 @@ public void testCoder() throws Exception { CoderProperties.coderDecodeEncodeEqual(coder, nonPollingState); } + @Test + public void testCoderWithTimestampCursor() throws Exception { + Instant now = Instant.now(); + ImmutableMap completed = + ImmutableMap.of(HashCode.fromString("0123456789abcdef0123456789abcdef"), now); + GrowthState withoutCursor = PollingGrowthState.of(completed, now, "STATE"); + GrowthState withCursor = PollingGrowthState.of(completed, now, "STATE", now); + Coder coder = + Watch.GrowthStateCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()); + + CoderProperties.coderDecodeEncodeEqual(coder, withCursor); + // A state without a cursor keeps the pre-cursor bytes, so a pipeline written before the cursor + // existed can still be updated. + assertEquals(0, CoderUtils.encodeToByteArray(coder, withoutCursor)[0]); + } + + @Test + public void testCoderKeepsPreCursorEncodedForm() throws Exception { + // Encoded forms produced before the cursor existed; an update must keep them byte for byte. + Instant ts = new Instant(1234567890123L); + GrowthState polling = + PollingGrowthState.of( + ImmutableMap.of(HashCode.fromString("0123456789abcdef0123456789abcdef"), ts), + ts, + "STATE"); + GrowthState nonPolling = + NonPollingGrowthState.of(Growth.PollResult.incomplete(ts, Arrays.asList("A", "B"))); + Coder coder = + Watch.GrowthStateCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()); + + assertEquals( + "00055354415445018000011F71FB04CB0000000101234567" + + "89ABCDEF0123456789ABCDEF8000011F71FB04CB", + BaseEncoding.base16().encode(CoderUtils.encodeToByteArray(coder, polling))); + assertEquals( + "01000000000201418000011F71FB04CB01428000011F71FB04CB", + BaseEncoding.base16().encode(CoderUtils.encodeToByteArray(coder, nonPolling))); + } + /** * Gradually emits all items from the given list, pairing each one with a UUID that identifies the * round of polling, so a client can check how many rounds of polling there were. @@ -370,6 +444,38 @@ public PollResult apply(InputT element, Context c) throws Exception { } } + /** + * Gradually emits all items from the given list, giving item {@code i} a timestamp of its own so + * that every poll reports an item at the same timestamp. Items are paired onto shared timestamps, + * so a poll boundary can fall between two items that carry the same one. + */ + private static class StablyTimedPollFn extends PollFn { + + private final Instant baseTime; + private final List outputs; + private final Duration timeToOutputEverything; + + StablyTimedPollFn(List outputs, Duration timeToOutputEverything) { + this.baseTime = Instant.now(); + this.outputs = outputs; + this.timeToOutputEverything = timeToOutputEverything; + } + + @Override + public PollResult apply(InputT element, Context c) throws Exception { + Duration elapsed = new Duration(baseTime, Instant.now()); + double fractionElapsed = 1.0 * elapsed.getMillis() / timeToOutputEverything.getMillis(); + int numToEmit = (int) Math.min(outputs.size(), fractionElapsed * outputs.size()); + List> toEmit = Lists.newArrayList(); + for (int i = 0; i < numToEmit; ++i) { + toEmit.add(TimestampedValue.of(outputs.get(i), baseTime.plus(standardSeconds(i / 2)))); + } + return numToEmit == outputs.size() + ? PollResult.complete(toEmit) + : PollResult.incomplete(toEmit); + } + } + @Test public void testTerminationConditionsNever() { Watch.Growth.Never c = never(); @@ -441,6 +547,11 @@ public void testTerminationConditionsAllOf() { } private static GrowthTracker newTracker(GrowthState state) { + return newTracker(state, null); + } + + private static GrowthTracker newTracker( + GrowthState state, Duration allowedLateness) { Funnel coderFunnel = (from, into) -> { try { @@ -449,7 +560,7 @@ private static GrowthTracker newTracker(GrowthState state) { throw new RuntimeException(e); } }; - return new GrowthTracker<>(state, coderFunnel); + return new GrowthTracker<>(state, coderFunnel, allowedLateness); } private static HashCode hash128(String value) { @@ -468,6 +579,11 @@ private static GrowthTracker newPollingGrowthTracker() { return newTracker(PollingGrowthState.of(never().forNewInput(Instant.now(), null))); } + private static GrowthTracker newPollingGrowthTracker(Duration allowedLateness) { + return newTracker( + PollingGrowthState.of(never().forNewInput(Instant.now(), null)), allowedLateness); + } + @Test public void testPollingGrowthTrackerUsesElementTimestampIfNoWatermarkProvided() throws Exception { Instant now = Instant.now(); @@ -500,6 +616,91 @@ public PollResult apply(String element, Context c) throws Exception { assertTrue(processContinuation.shouldResume()); } + @Test + public void testPollingGrowthTrackerDropsOutputsBehindCursor() throws Exception { + Instant now = Instant.now(); + Watch.Growth growth = + Watch.growthOf( + new PollFn() { + @Override + public PollResult apply(String element, Context c) throws Exception { + return PollResult.incomplete( + Arrays.asList( + TimestampedValue.of("retired", now.plus(standardSeconds(1))), + TimestampedValue.of("atCursor", now.plus(standardSeconds(3))), + TimestampedValue.of("fresh", now.plus(standardSeconds(5))))); + } + }) + .withPollInterval(standardSeconds(10)) + .withTimestampCursor(); + WatchGrowthFn growthFn = + new WatchGrowthFn( + growth, StringUtf8Coder.of(), SerializableFunctions.identity(), StringUtf8Coder.of()); + GrowthTracker tracker = + newTracker( + PollingGrowthState.of( + ImmutableMap.of(), + null, + never().forNewInput(now, null), + now.plus(standardSeconds(3))), + Duration.ZERO); + DoFn.ProcessContext context = mock(DoFn.ProcessContext.class); + ManualWatermarkEstimator watermarkEstimator = + new WatermarkEstimators.Manual(BoundedWindow.TIMESTAMP_MIN_VALUE); + + ProcessContinuation processContinuation = + growthFn.process(context, tracker, watermarkEstimator); + + // The output below the cursor has no key left to prove it was seen, so it is taken as seen. An + // output at the cursor is still retained, so it is deduplicated by key rather than dropped. + verify(context) + .output( + KV.of( + null, + Arrays.asList( + TimestampedValue.of("atCursor", now.plus(standardSeconds(3))), + TimestampedValue.of("fresh", now.plus(standardSeconds(5)))))); + assertEquals(now.plus(standardSeconds(3)), watermarkEstimator.currentWatermark()); + assertTrue(processContinuation.shouldResume()); + } + + @Test + public void testPollingGrowthTrackerEmptyRoundAdvancesWatermarkToFloor() throws Exception { + Instant now = Instant.now(); + Watch.Growth growth = + Watch.growthOf( + new PollFn() { + @Override + public PollResult apply(String element, Context c) throws Exception { + // A re-listed output below the floor, and no explicit watermark. + return PollResult.incomplete( + Arrays.asList( + TimestampedValue.of("retired", now.minus(standardSeconds(1))))); + } + }) + .withPollInterval(standardSeconds(10)) + .withTimestampCursor(); + WatchGrowthFn growthFn = + new WatchGrowthFn( + growth, StringUtf8Coder.of(), SerializableFunctions.identity(), StringUtf8Coder.of()); + GrowthTracker tracker = + newTracker( + PollingGrowthState.of(ImmutableMap.of(), null, never().forNewInput(now, null), now), + Duration.ZERO); + DoFn.ProcessContext context = mock(DoFn.ProcessContext.class); + ManualWatermarkEstimator watermarkEstimator = + new WatermarkEstimators.Manual(BoundedWindow.TIMESTAMP_MIN_VALUE); + + ProcessContinuation processContinuation = + growthFn.process(context, tracker, watermarkEstimator); + + // Nothing below the retention floor is ever emitted, so the floor is a sound watermark for a + // round that computed none. + verify(context, org.mockito.Mockito.never()).output(any()); + assertEquals(now, watermarkEstimator.currentWatermark()); + assertTrue(processContinuation.shouldResume()); + } + @Test public void testPollingGrowthTrackerCheckpointNonEmpty() { Instant now = Instant.now(); @@ -531,6 +732,151 @@ public void testPollingGrowthTrackerCheckpointNonEmpty() { residual.getCompleted().keySet(), containsInAnyOrder(hash128("a"), hash128("b"), hash128("c"), hash128("d"))); assertEquals(1, (int) residual.getTerminationState()); + assertNull(residual.getCursor()); + } + + @Test + public void testPollingGrowthTrackerRetiresCompletedBehindCursor() { + Instant now = Instant.now(); + GrowthTracker tracker = newPollingGrowthTracker(Duration.ZERO); + + PollResult claim = + PollResult.incomplete( + Arrays.asList( + TimestampedValue.of("a", now.plus(standardSeconds(1))), + TimestampedValue.of("b", now.plus(standardSeconds(2))), + TimestampedValue.of("c", now.plus(standardSeconds(4))), + TimestampedValue.of("d", now.plus(standardSeconds(4))))); + + assertTrue(tracker.tryClaim(KV.of(claim, 1 /* termination state */))); + + PollingGrowthState residual = + (PollingGrowthState) tracker.trySplit(0).getResidual(); + + assertEquals(now.plus(standardSeconds(4)), residual.getCursor()); + // A key at the cursor is retained, so an output that arrives later at the same timestamp is + // still deduplicated by key. + assertThat(residual.getCompleted().keySet(), containsInAnyOrder(hash128("c"), hash128("d"))); + } + + @Test + public void testPollingGrowthTrackerAllowedLatenessRetainsCompleted() { + Instant now = Instant.now(); + GrowthTracker tracker = newPollingGrowthTracker(standardSeconds(2)); + + PollResult claim = + PollResult.incomplete( + Arrays.asList( + TimestampedValue.of("a", now.plus(standardSeconds(1))), + TimestampedValue.of("b", now.plus(standardSeconds(2))), + TimestampedValue.of("c", now.plus(standardSeconds(4))), + TimestampedValue.of("d", now.plus(standardSeconds(4))))); + + assertTrue(tracker.tryClaim(KV.of(claim, 1 /* termination state */))); + + PollingGrowthState residual = + (PollingGrowthState) tracker.trySplit(0).getResidual(); + + assertEquals(now.plus(standardSeconds(4)), residual.getCursor()); + assertThat( + residual.getCompleted().keySet(), + containsInAnyOrder(hash128("b"), hash128("c"), hash128("d"))); + } + + @Test + public void testPollingGrowthTrackerRoundWithoutCursorDropsStaleCursor() throws Exception { + Instant now = Instant.now(); + // A round that is not bounding the state retains every key, so the cursor that would retire + // them is dropped and the restriction returns to the pre-cursor encoding. + GrowthState state = + PollingGrowthState.of( + ImmutableMap.of(), null, never().forNewInput(now, null), now.plus(standardSeconds(10))); + GrowthTracker tracker = newTracker(state, null); + + assertTrue( + tracker.tryClaim( + KV.of( + PollResult.incomplete( + Arrays.asList(TimestampedValue.of("a", now.plus(standardSeconds(20))))), + 1))); + + PollingGrowthState residual = + (PollingGrowthState) tracker.trySplit(0).getResidual(); + + assertNull(residual.getCursor()); + assertEquals(1, residual.getCompleted().size()); + Coder coder = Watch.GrowthStateCoder.of(StringUtf8Coder.of(), VarIntCoder.of()); + assertEquals(0, CoderUtils.encodeToByteArray(coder, residual)[0]); + } + + @Test + public void testPollingGrowthTrackerAllowedLatenessKeepsMaxCursorClaimable() throws Exception { + // A cursor at the maximum timestamp still leaves the allowed lateness window claimable. + GrowthState state = + PollingGrowthState.of( + ImmutableMap.of(), + null, + never().forNewInput(Instant.now(), null), + BoundedWindow.TIMESTAMP_MAX_VALUE); + GrowthTracker tracker = newTracker(state, Duration.standardHours(1)); + + assertTrue( + tracker.tryClaim( + KV.of( + PollResult.incomplete( + Arrays.asList( + TimestampedValue.of( + "late", + BoundedWindow.TIMESTAMP_MAX_VALUE.minus( + Duration.standardMinutes(30))))), + 1))); + } + + @Test + public void testPollingGrowthTrackerHugeAllowedLatenessDoesNotOverflow() { + Instant now = Instant.now(); + GrowthState state = + PollingGrowthState.of(ImmutableMap.of(), null, never().forNewInput(now, null), now); + GrowthTracker tracker = newTracker(state, Duration.millis(Long.MAX_VALUE)); + + // The floor saturates at the minimum timestamp rather than throwing. + assertTrue( + tracker.tryClaim( + KV.of( + PollResult.incomplete( + Arrays.asList(TimestampedValue.of("a", BoundedWindow.TIMESTAMP_MIN_VALUE))), + 1))); + } + + @Test + public void testPollingGrowthTrackerRejectsClaimBehindCursor() { + Instant now = Instant.now(); + GrowthTracker tracker = newPollingGrowthTracker(Duration.ZERO); + + assertTrue( + tracker.tryClaim( + KV.of( + PollResult.incomplete( + Arrays.asList(TimestampedValue.of("a", now.plus(standardSeconds(4))))), + 1))); + + PollingGrowthState residual = + (PollingGrowthState) tracker.trySplit(0).getResidual(); + + assertFalse( + newTracker(residual, Duration.ZERO) + .tryClaim( + KV.of( + PollResult.incomplete( + Arrays.asList(TimestampedValue.of("b", now.plus(standardSeconds(3))))), + 2))); + assertTrue( + newTracker(residual, Duration.ZERO) + .tryClaim( + KV.of( + PollResult.incomplete( + Arrays.asList(TimestampedValue.of("b", now.plus(standardSeconds(4))))), + 2))); } @Test