Skip to content

Commit

Permalink
Fix snapshot metrics (#182)
Browse files Browse the repository at this point in the history 
Shift since snapshot metrics to `snapshotState()` rather than first
write to prevent leaking the metric values to scope of the next
checkpoint.

/gcbrun
  • Loading branch information
@prashastia
prashastia authored Dec 10, 2024
1 parent 5462bef commit 5d8b34a
Show file tree
Hide file tree
Showing 2 changed files with 49 additions and 16 deletions.
10 changes: 5 additions & 5 deletions ...ery/src/main/java/com/google/cloud/flink/bigquery/sink/writer/BigQueryBufferedWriter.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -174,11 +174,6 @@ private void preWriteOpsAfterCommit() {
long numberOfRecordsWrittenInLastCommit = totalRecordsWritten - totalRecordsCommitted;
totalRecordsCommitted = totalRecordsWritten;
numberOfRecordsWrittenToBigQuery.inc(numberOfRecordsWrittenInLastCommit);
// Reset the "Since Checkpoint" values to 0.
numberOfRecordsBufferedByBigQuerySinceCheckpoint.dec(
numberOfRecordsBufferedByBigQuerySinceCheckpoint.getCount());
numberOfRecordsSeenByWriterSinceCheckpoint.dec(
numberOfRecordsSeenByWriterSinceCheckpoint.getCount());
}

/**
Expand Down Expand Up @@ -272,6 +267,11 @@ public List<BigQueryWriterState> snapshotState(long checkpointId) {
isFirstWriteAfterCheckpoint = true;
streamNameInState = streamName;
streamOffsetInState = streamOffset;
// Reset the "Since Checkpoint" values to 0.
numberOfRecordsBufferedByBigQuerySinceCheckpoint.dec(
numberOfRecordsBufferedByBigQuerySinceCheckpoint.getCount());
numberOfRecordsSeenByWriterSinceCheckpoint.dec(
numberOfRecordsSeenByWriterSinceCheckpoint.getCount());
return Collections.singletonList(
// Note that it's possible to store the associated checkpointId in writer's state.
// For now, we're not leveraging this due to absence of a use case.
Expand Down
55 changes: 44 additions & 11 deletions ...src/test/java/com/google/cloud/flink/bigquery/sink/writer/BigQueryBufferedWriterTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,7 @@ public void testConstructor_withNewWriter() {
assertTrue(bufferedWriter.getProtoRows().getSerializedRowsList().isEmpty());
assertTrue(bufferedWriter.getAppendResponseFuturesQueue().isEmpty());
// Test Flink Metrics
// All should be 0 since the writer is newly created
assertEquals(0, bufferedWriter.numberOfRecordsSeenByWriter.getCount());
assertEquals(0, bufferedWriter.numberOfRecordsSeenByWriterSinceCheckpoint.getCount());
assertEquals(0, bufferedWriter.numberOfRecordsWrittenToBigQuery.getCount());
Expand All @@ -116,6 +117,7 @@ public void testConstructor_withRestoredWriter() {
assertTrue(bufferedWriter.getProtoRows().getSerializedRowsList().isEmpty());
assertTrue(bufferedWriter.getAppendResponseFuturesQueue().isEmpty());
// Test Flink Metrics
// Since Checkpoint metrics are 0, rest are restored to saved values.
assertEquals(210, bufferedWriter.numberOfRecordsSeenByWriter.getCount());
assertEquals(0, bufferedWriter.numberOfRecordsSeenByWriterSinceCheckpoint.getCount());
assertEquals(100, bufferedWriter.numberOfRecordsWrittenToBigQuery.getCount());
Expand Down Expand Up @@ -146,6 +148,7 @@ public void testWrite_withoutAppend() {
assertEquals(1, bufferedWriter.numberOfRecordsSeenByWriter.getCount());
assertEquals(1, bufferedWriter.numberOfRecordsSeenByWriterSinceCheckpoint.getCount());
assertEquals(0, bufferedWriter.numberOfRecordsWrittenToBigQuery.getCount());
// Since append is not called yet.
assertEquals(0, bufferedWriter.numberOfRecordsBufferedByBigQuerySinceCheckpoint.getCount());

bufferedWriter.write(new Object(), null);
Expand All @@ -157,6 +160,7 @@ public void testWrite_withoutAppend() {
assertEquals(2, bufferedWriter.getProtoRows().getSerializedRowsCount());
assertTrue(bufferedWriter.getAppendResponseFuturesQueue().isEmpty());
// Test Flink Metrics
// Since append is not called yet.
assertEquals(2, bufferedWriter.numberOfRecordsSeenByWriter.getCount());
assertEquals(2, bufferedWriter.numberOfRecordsSeenByWriterSinceCheckpoint.getCount());
assertEquals(0, bufferedWriter.numberOfRecordsWrittenToBigQuery.getCount());
Expand Down Expand Up @@ -203,6 +207,7 @@ public void testWrite_withAppend_withNewStream() {
// Test Flink Metrics
assertEquals(1, bufferedWriter.numberOfRecordsSeenByWriter.getCount());
assertEquals(1, bufferedWriter.numberOfRecordsSeenByWriterSinceCheckpoint.getCount());
// Append() not called.
assertEquals(0, bufferedWriter.numberOfRecordsWrittenToBigQuery.getCount());
assertEquals(0, bufferedWriter.numberOfRecordsBufferedByBigQuerySinceCheckpoint.getCount());

Expand Down Expand Up @@ -244,6 +249,7 @@ public void testWrite_withAppend_withNewStream() {
assertEquals(3, bufferedWriter.numberOfRecordsSeenByWriter.getCount());
assertEquals(3, bufferedWriter.numberOfRecordsSeenByWriterSinceCheckpoint.getCount());
assertEquals(0, bufferedWriter.numberOfRecordsWrittenToBigQuery.getCount());
// Post append
assertEquals(1, bufferedWriter.numberOfRecordsBufferedByBigQuerySinceCheckpoint.getCount());

((FakeBigQueryStorageWriteClient) bufferedWriter.writeClient)
Expand Down Expand Up @@ -299,6 +305,7 @@ public void testWrite_withAppend_withUsableRestoredStream() {
assertEquals(211, bufferedWriter.numberOfRecordsSeenByWriter.getCount());
assertEquals(1, bufferedWriter.numberOfRecordsSeenByWriterSinceCheckpoint.getCount());
assertEquals(200, bufferedWriter.numberOfRecordsWrittenToBigQuery.getCount());
// 0 since restored writer, append() not called yet.
assertEquals(0, bufferedWriter.numberOfRecordsBufferedByBigQuerySinceCheckpoint.getCount());

// Second element will exceed append request's size, so append will be called with
Expand Down Expand Up @@ -713,6 +720,11 @@ public void testSnapshotState_withNewWriter() {
bufferedWriter.flush(false);
assertEquals("", bufferedWriter.getStreamNameInState());
assertEquals(0, bufferedWriter.getStreamOffsetInState());
// Test Flink Metrics
assertEquals(3, bufferedWriter.numberOfRecordsSeenByWriter.getCount());
assertEquals(3, bufferedWriter.numberOfRecordsSeenByWriterSinceCheckpoint.getCount());
assertEquals(0, bufferedWriter.numberOfRecordsWrittenToBigQuery.getCount());
assertEquals(3, bufferedWriter.numberOfRecordsBufferedByBigQuerySinceCheckpoint.getCount());
Collection<BigQueryWriterState> writerStates = bufferedWriter.snapshotState(1);
BigQueryWriterState writerState = (BigQueryWriterState) writerStates.toArray()[0];
assertEquals(1, writerStates.size());
Expand All @@ -725,9 +737,9 @@ public void testSnapshotState_withNewWriter() {
assertEquals(3, bufferedWriter.getStreamOffsetInState());
// Test Flink Metrics
assertEquals(3, bufferedWriter.numberOfRecordsSeenByWriter.getCount());
assertEquals(3, bufferedWriter.numberOfRecordsSeenByWriterSinceCheckpoint.getCount());
assertEquals(0, bufferedWriter.numberOfRecordsSeenByWriterSinceCheckpoint.getCount());
assertEquals(0, bufferedWriter.numberOfRecordsWrittenToBigQuery.getCount());
assertEquals(3, bufferedWriter.numberOfRecordsBufferedByBigQuerySinceCheckpoint.getCount());
assertEquals(0, bufferedWriter.numberOfRecordsBufferedByBigQuerySinceCheckpoint.getCount());
}

@Test
Expand Down Expand Up @@ -763,14 +775,19 @@ public void testSnapshotState_withNewWriter_metrics() {
bufferedWriter.write(new Object(), null);
bufferedWriter.write(new Object(), null);
bufferedWriter.flush(false);
// Test Flink Metrics
assertEquals(3, bufferedWriter.numberOfRecordsSeenByWriter.getCount());
assertEquals(3, bufferedWriter.numberOfRecordsSeenByWriterSinceCheckpoint.getCount());
assertEquals(0, bufferedWriter.numberOfRecordsWrittenToBigQuery.getCount());
assertEquals(3, bufferedWriter.numberOfRecordsBufferedByBigQuerySinceCheckpoint.getCount());
Collection<BigQueryWriterState> writerStates = bufferedWriter.snapshotState(1);
BigQueryWriterState writerState = (BigQueryWriterState) writerStates.toArray()[0];
// Test Flink Metrics
assertEquals(3, bufferedWriter.numberOfRecordsSeenByWriter.getCount());
assertEquals(3, bufferedWriter.numberOfRecordsSeenByWriterSinceCheckpoint.getCount());
assertEquals(0, bufferedWriter.numberOfRecordsSeenByWriterSinceCheckpoint.getCount());
// Updated at first write after checkpoint.
assertEquals(0, bufferedWriter.numberOfRecordsWrittenToBigQuery.getCount());
assertEquals(3, bufferedWriter.numberOfRecordsBufferedByBigQuerySinceCheckpoint.getCount());
assertEquals(0, bufferedWriter.numberOfRecordsBufferedByBigQuerySinceCheckpoint.getCount());
bufferedWriter.write(new Object(), null);
// Test Flink Metrics
assertEquals(4, bufferedWriter.numberOfRecordsSeenByWriter.getCount());
Expand Down Expand Up @@ -819,6 +836,11 @@ public void testSnapshotState_withRestoredWriter_withUsableStream() {
bufferedWriter.write(new Object(), null);
bufferedWriter.write(new Object(), null);
bufferedWriter.flush(false);
// Test Flink Metrics
assertEquals(213, bufferedWriter.numberOfRecordsSeenByWriter.getCount());
assertEquals(3, bufferedWriter.numberOfRecordsSeenByWriterSinceCheckpoint.getCount());
assertEquals(200, bufferedWriter.numberOfRecordsWrittenToBigQuery.getCount());
assertEquals(3, bufferedWriter.numberOfRecordsBufferedByBigQuerySinceCheckpoint.getCount());
assertEquals("restored_stream", bufferedWriter.getStreamNameInState());
assertEquals(100, bufferedWriter.getStreamOffsetInState());
Collection<BigQueryWriterState> writerStates = bufferedWriter.snapshotState(1);
Expand All @@ -833,9 +855,9 @@ public void testSnapshotState_withRestoredWriter_withUsableStream() {
assertEquals(103, bufferedWriter.getStreamOffsetInState());
// Test Flink Metrics
assertEquals(213, bufferedWriter.numberOfRecordsSeenByWriter.getCount());
assertEquals(3, bufferedWriter.numberOfRecordsSeenByWriterSinceCheckpoint.getCount());
assertEquals(0, bufferedWriter.numberOfRecordsSeenByWriterSinceCheckpoint.getCount());
assertEquals(200, bufferedWriter.numberOfRecordsWrittenToBigQuery.getCount());
assertEquals(3, bufferedWriter.numberOfRecordsBufferedByBigQuerySinceCheckpoint.getCount());
assertEquals(0, bufferedWriter.numberOfRecordsBufferedByBigQuerySinceCheckpoint.getCount());
}

@Test
Expand Down Expand Up @@ -876,13 +898,19 @@ public void testSnapshotState_withRestoredWriter_withUsableStream_testMetrics()
bufferedWriter.write(new Object(), null);
bufferedWriter.write(new Object(), null);
bufferedWriter.write(new Object(), null);
bufferedWriter.flush(false);
Collection<BigQueryWriterState> writerStates = bufferedWriter.snapshotState(1);
// Test Flink Metrics
assertEquals(213, bufferedWriter.numberOfRecordsSeenByWriter.getCount());
assertEquals(3, bufferedWriter.numberOfRecordsSeenByWriterSinceCheckpoint.getCount());
assertEquals(200, bufferedWriter.numberOfRecordsWrittenToBigQuery.getCount());
assertEquals(1, bufferedWriter.numberOfRecordsBufferedByBigQuerySinceCheckpoint.getCount());
bufferedWriter.flush(false);
assertEquals(3, bufferedWriter.numberOfRecordsBufferedByBigQuerySinceCheckpoint.getCount());
Collection<BigQueryWriterState> writerStates = bufferedWriter.snapshotState(1);
// Test Flink Metrics
assertEquals(213, bufferedWriter.numberOfRecordsSeenByWriter.getCount());
assertEquals(0, bufferedWriter.numberOfRecordsSeenByWriterSinceCheckpoint.getCount());
assertEquals(200, bufferedWriter.numberOfRecordsWrittenToBigQuery.getCount());
assertEquals(0, bufferedWriter.numberOfRecordsBufferedByBigQuerySinceCheckpoint.getCount());
bufferedWriter.write(new Object(), null);
bufferedWriter.write(new Object(), null);
assertEquals(215, bufferedWriter.numberOfRecordsSeenByWriter.getCount());
Expand Down Expand Up @@ -927,6 +955,11 @@ public void testSnapshotState_withRestoredWriter_withUnusableStream() {
bufferedWriter.flush(false);
assertEquals("restored_stream", bufferedWriter.getStreamNameInState());
assertEquals(100, bufferedWriter.getStreamOffsetInState());
// Test Flink Metrics
assertEquals(213, bufferedWriter.numberOfRecordsSeenByWriter.getCount());
assertEquals(3, bufferedWriter.numberOfRecordsSeenByWriterSinceCheckpoint.getCount());
assertEquals(200, bufferedWriter.numberOfRecordsWrittenToBigQuery.getCount());
assertEquals(3, bufferedWriter.numberOfRecordsBufferedByBigQuerySinceCheckpoint.getCount());
Collection<BigQueryWriterState> writerStates = bufferedWriter.snapshotState(1);
BigQueryWriterState writerState = (BigQueryWriterState) writerStates.toArray()[0];
assertEquals(1, writerStates.size());
Expand All @@ -937,11 +970,11 @@ public void testSnapshotState_withRestoredWriter_withUnusableStream() {
assertEquals(1, writerState.getCheckpointId());
assertEquals("new_stream", bufferedWriter.getStreamNameInState());
assertEquals(3, bufferedWriter.getStreamOffsetInState());
// Test Flink Metrics
// Test Flink Metrics - Set to 0.
assertEquals(213, bufferedWriter.numberOfRecordsSeenByWriter.getCount());
assertEquals(3, bufferedWriter.numberOfRecordsSeenByWriterSinceCheckpoint.getCount());
assertEquals(0, bufferedWriter.numberOfRecordsSeenByWriterSinceCheckpoint.getCount());
assertEquals(200, bufferedWriter.numberOfRecordsWrittenToBigQuery.getCount());
assertEquals(3, bufferedWriter.numberOfRecordsBufferedByBigQuerySinceCheckpoint.getCount());
assertEquals(0, bufferedWriter.numberOfRecordsBufferedByBigQuerySinceCheckpoint.getCount());
}

@Test
Expand Down

0 comments on commit 5d8b34a

Please sign in to comment.