@Override public RecordEnvelope<Object> readRecordEnvelope() throws DataRecordException, IOException { TestRecord record = (TestRecord) extractor.readRecord(null); return new RecordEnvelope<>((Object) record, new DefaultCheckpointableWatermark(""+record.getPartition(), new LongWatermark(record.getSequence()))); }
@Override public String toString() { return String.format("%s : %s ", getSource(), GSON.toJson(this.comparable.toJson())); } }
@Override public Map<String, CheckpointableWatermark> getCommittableWatermark() { watermark++; if (watermark % failEvery == 0) { throw new RuntimeException("Failed because you asked me to"); } return Collections.singletonMap("default", (CheckpointableWatermark) new DefaultCheckpointableWatermark("default", new LongWatermark(watermark))); }
@Override public String toString() { return String.format("%s : %s ", getSource(), GSON.toJson(this.comparable.toJson())); } }
@Override public Map<String, CheckpointableWatermark> getCommittableWatermark() { watermark++; return Collections.singletonMap("default", (CheckpointableWatermark) new DefaultCheckpointableWatermark("default", new LongWatermark(watermark))); }
@Override public RecordEnvelope<String> readRecordEnvelope() throws DataRecordException, IOException { if (!this.closed) { if (index == 0) { RecordEnvelope<String> recordEnvelope = new RecordEnvelope<>(record, new DefaultCheckpointableWatermark("default", new LongWatermark(index))); log.debug("Returning record with index {}", index); index++; return recordEnvelope; } else { return null; } } else { log.info("Extractor has been closed, returning null"); return null; } }
@Override public RecordEnvelope<String> readRecordEnvelope() throws DataRecordException, IOException { if (!this.closed) { try { Thread.sleep(this.sleepTimeInMillis); } catch (InterruptedException e) { Throwables.propagate(e); } String record = index + ""; RecordEnvelope<String> recordEnvelope = new RecordEnvelope<>(record, new DefaultCheckpointableWatermark("default", new LongWatermark(index))); index++; return recordEnvelope; } else { log.info("Extractor has been closed, returning null"); return null; } }
(CheckpointableWatermark) new DefaultCheckpointableWatermark("default", new LongWatermark(watermark)));
private void commits(WatermarkTracker watermarkTracker, String source, int... commit) { for (int oneCommit: commit) { watermarkTracker.committedWatermark(new DefaultCheckpointableWatermark(source, new LongWatermark(oneCommit))); } }
public void testWatermarkComputation(Long committed, Long unacknowledged, Long expected) throws IOException { State state = new State(); state.setProp(ConfigurationKeys.WRITER_PARTITIONER_CLASS, TestPartitioner.class.getCanonicalName()); String defaultSource = "default"; WatermarkAwareWriter mockDataWriter = mock(WatermarkAwareWriter.class); when(mockDataWriter.isWatermarkCapable()).thenReturn(true); when(mockDataWriter.getCommittableWatermark()).thenReturn(Collections.singletonMap(defaultSource, new DefaultCheckpointableWatermark(defaultSource, new LongWatermark(committed)))); when(mockDataWriter.getUnacknowledgedWatermark()).thenReturn(Collections.singletonMap(defaultSource, new DefaultCheckpointableWatermark(defaultSource, new LongWatermark(unacknowledged)))); PartitionAwareDataWriterBuilder builder = mock(PartitionAwareDataWriterBuilder.class); when(builder.validatePartitionSchema(any(Schema.class))).thenReturn(true); when(builder.forPartition(any(GenericRecord.class))).thenReturn(builder); when(builder.withWriterId(any(String.class))).thenReturn(builder); when(builder.build()).thenReturn(mockDataWriter); PartitionedDataWriter writer = new PartitionedDataWriter<String, String>(builder, state); RecordEnvelope<String> recordEnvelope = new RecordEnvelope<String>("0"); recordEnvelope.addCallBack( new AcknowledgableWatermark(new DefaultCheckpointableWatermark(defaultSource, new LongWatermark(0)))); writer.writeEnvelope(recordEnvelope); Map<String, CheckpointableWatermark> watermark = writer.getCommittableWatermark(); System.out.println(watermark.toString()); if (expected == null) { Assert.assertTrue(watermark.isEmpty(), "Expected watermark to be absent"); } else { Assert.assertTrue(watermark.size() == 1); Assert.assertEquals((long) expected, ((LongWatermark) watermark.values().iterator().next().getWatermark()).getValue()); } }
/** * Test that when we have commits failing to watermark storage, the manager continues to try * at every interval and keeps track of the exception it is seeing. */ @Test public void testFailingWatermarkStorage() throws IOException, InterruptedException { WatermarkStorage reallyBadWatermarkStorage = mock(WatermarkStorage.class); IOException exceptionToThrow = new IOException("Failed to write coz the programmer told me to"); doThrow(exceptionToThrow).when(reallyBadWatermarkStorage).commitWatermarks(any(Iterable.class)); long commitInterval = 1000; MultiWriterWatermarkManager watermarkManager = new MultiWriterWatermarkManager(reallyBadWatermarkStorage, commitInterval, Optional.<Logger>absent()); WatermarkAwareWriter mockWriter = mock(WatermarkAwareWriter.class); CheckpointableWatermark watermark = new DefaultCheckpointableWatermark("default", new LongWatermark(0)); when(mockWriter.getCommittableWatermark()).thenReturn(Collections.singletonMap("default", watermark)); watermarkManager.registerWriter(mockWriter); try { watermarkManager.start(); } catch (Exception e) { Assert.fail("Should not throw exception", e); } Thread.sleep(commitInterval * 2 + (commitInterval/2)); // sleep for 2.5 iterations watermarkManager.close(); int expectedCalls = 3; // 2 calls from iterations, 1 additional attempt due to close verify(reallyBadWatermarkStorage, atLeast(expectedCalls)).commitWatermarks(any(Iterable.class)); Assert.assertEquals(watermarkManager.getCommitStatus().getLastCommitException(), exceptionToThrow, "Testing tracking of failed exceptions"); }
CheckpointableWatermark checkpointableWatermark = new DefaultCheckpointableWatermark("default", new LongWatermark(i)); AcknowledgableWatermark ackable = new AcknowledgableWatermark(checkpointableWatermark); acknowledgableWatermarks[i] = ackable;
CheckpointableWatermark checkpointableWatermark = new DefaultCheckpointableWatermark("default", new LongWatermark(i)); final AcknowledgableWatermark ackable = new AcknowledgableWatermark(checkpointableWatermark); tracker.track(ackable);
@Test public void testPersistWatermarkStateToZk() throws IOException { CheckpointableWatermark watermark = new DefaultCheckpointableWatermark("source", new LongWatermark(startTime)); TaskState taskState = new TaskState(); taskState.setJobId(TEST_JOB_ID); taskState.setProp(ConfigurationKeys.JOB_NAME_KEY, "JobName-" + startTime); // watermark storage configuration taskState.setProp(StateStoreBasedWatermarkStorage.WATERMARK_STORAGE_TYPE_KEY, "zk"); taskState.setProp(StateStoreBasedWatermarkStorage.WATERMARK_STORAGE_CONFIG_PREFIX + ZkStateStoreConfigurationKeys.STATE_STORE_ZK_CONNECT_STRING_KEY, testingServer.getConnectString()); StateStoreBasedWatermarkStorage watermarkStorage = new StateStoreBasedWatermarkStorage(taskState); watermarkStorage.commitWatermarks(ImmutableList.of(watermark)); Map<String, CheckpointableWatermark> watermarkMap = watermarkStorage.getCommittedWatermarks(DefaultCheckpointableWatermark.class, ImmutableList.of("source")); Assert.assertEquals(watermarkMap.size(), 1); Assert.assertEquals(((LongWatermark) watermarkMap.get("source").getWatermark()).getValue(), startTime); }
CheckpointableWatermark checkpointableWatermark = new DefaultCheckpointableWatermark("default", new LongWatermark(i)); AcknowledgableWatermark ackable = new AcknowledgableWatermark(checkpointableWatermark); acknowledgableWatermarks[i] = ackable;
private void writeEnvelope(ConsoleWriter consoleWriter, String content, String source, long value) throws IOException { CheckpointableWatermark watermark = new DefaultCheckpointableWatermark(source, new LongWatermark(value)); AcknowledgableWatermark ackable = new AcknowledgableWatermark(watermark); RecordEnvelope<String> mockEnvelope = (RecordEnvelope<String>) new RecordEnvelope<>(content).addCallBack(ackable); consoleWriter.writeEnvelope(mockEnvelope); Assert.assertTrue(ackable.isAcked()); }
watermark++; return Collections.singletonMap("default", (CheckpointableWatermark) new DefaultCheckpointableWatermark("default", new LongWatermark(watermark)));
@Override public RecordEnvelope<Object> readRecordEnvelope() throws DataRecordException, IOException { TestRecord record = (TestRecord) extractor.readRecord(null); return new RecordEnvelope<>((Object) record, new DefaultCheckpointableWatermark(""+record.getPartition(), new LongWatermark(record.getSequence()))); }