Skip to content

[BEAM-90] TestCountingSource can throw on checkpointing #14

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

Closed
mshields822 wants to merge 2 commits into from
Closed

[BEAM-90] TestCountingSource can throw on checkpointing #14

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f2b7c60
[BEAM-90] TestCountingSource can throw on checkpointing
Mar 3, 2016
5dc0c49
Address review comments.
Mar 3, 2016
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
31 changes: 27 additions & 4 deletions sdk/src/test/java/com/google/cloud/dataflow/sdk/runners/dataflow/TestCountingSource.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,8 @@
import com.google.cloud.dataflow.sdk.values.KV;

import org.joda.time.Instant;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.ArrayList;
Expand All @@ -46,31 +48,47 @@
*/
public class TestCountingSource
extends UnboundedSource<KV<Integer, Integer>, TestCountingSource.CounterMark> {
private static final Logger LOG = LoggerFactory.getLogger(TestCountingSource.class);

private static List<Integer> finalizeTracker;
private final int numMessagesPerShard;
private final int shardNumber;
private final boolean dedup;
private final boolean throwOnFirstSnapshot;

/**
* We only allow an exception to be thrown from getCheckpointMark
* at most once. This must be static since the entire TestCountingSource
* instance may re-serialized when the pipeline recovers and retries.
*/
private static boolean thrown = false;
Copy link
Member

@kennknowles kennknowles Mar 3, 2016

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you leave a little comment here to explain?


public static void setFinalizeTracker(List<Integer> finalizeTracker) {
TestCountingSource.finalizeTracker = finalizeTracker;
}

public TestCountingSource(int numMessagesPerShard) {
this(numMessagesPerShard, 0, false);
this(numMessagesPerShard, 0, false, false);
}

public TestCountingSource withDedup() {
return new TestCountingSource(numMessagesPerShard, shardNumber, true);
return new TestCountingSource(numMessagesPerShard, shardNumber, true, throwOnFirstSnapshot);
}

private TestCountingSource withShardNumber(int shardNumber) {
return new TestCountingSource(numMessagesPerShard, shardNumber, dedup);
return new TestCountingSource(numMessagesPerShard, shardNumber, dedup, throwOnFirstSnapshot);
}

private TestCountingSource(int numMessagesPerShard, int shardNumber, boolean dedup) {
public TestCountingSource withThrowOnFirstSnapshot(boolean throwOnFirstSnapshot) {
return new TestCountingSource(numMessagesPerShard, shardNumber, dedup, throwOnFirstSnapshot);
}

private TestCountingSource(
int numMessagesPerShard, int shardNumber, boolean dedup, boolean throwOnFirstSnapshot) {
this.numMessagesPerShard = numMessagesPerShard;
this.shardNumber = shardNumber;
this.dedup = dedup;
this.throwOnFirstSnapshot = throwOnFirstSnapshot;
}

public int getShardNumber() {
Expand Down Expand Up @@ -187,6 +205,11 @@ public Instant getWatermark() {

@Override
public CheckpointMark getCheckpointMark() {
if (throwOnFirstSnapshot && !thrown) {
thrown = true;
LOG.error("Throwing exception while checkpointing counter");
throw new RuntimeException("failed during checkpoint");
}
return new CounterMark(current);
}

Expand Down