@Override public int compareTo(CheckpointableWatermark o) { if (!(this.source.equals(o.getSource()))) { throw new RuntimeException("Could not compare two checkpointable watermarks because they have different sources " + this.source + ":" + o.getSource()); } return this.comparable.compareTo(o.getWatermark()); }
@Override public void start(WatermarkStorage watermarkStorage) throws IOException { this.watermarkStorage = Optional.of(watermarkStorage); Map<String, CheckpointableWatermark> lastCommitted; try { lastCommitted = this.watermarkStorage.get() .getCommittedWatermarks(DefaultCheckpointableWatermark.class, ImmutableList.of("" + extractor.partition)); } catch (IOException e) { // failed to get watermarks ... log a warning message log.warn("Failed to get watermarks... will start from the beginning", e); lastCommitted = Collections.EMPTY_MAP; } for (Map.Entry entry: lastCommitted.entrySet()) { log.info("{}: Found these committed watermarks: key: {}, value: {}", this, entry.getKey(), entry.getValue()); } LongWatermark currentWatermark; if (!lastCommitted.isEmpty() && lastCommitted.containsKey(""+extractor.partition)) { currentWatermark = (LongWatermark) (lastCommitted.get("" + extractor.partition)).getWatermark(); } else { // first record starts from 0 currentWatermark = new LongWatermark(-1); } extractor.setCurrentWatermark(currentWatermark); log.info("{}: Set current watermark to : {}", this, currentWatermark); } };
public boolean validateWatermarks(boolean exact, Map<String, CheckpointableWatermark> watermarkMap) { if (!watermarkMap.isEmpty()) { // watermark must be <= the index LongWatermark longWatermark = (LongWatermark) watermarkMap.values().iterator().next().getWatermark(); if (exact) { System.out.println(index-1 + ":" + longWatermark.getValue()); return ((index-1) == longWatermark.getValue()); } else { return (index > longWatermark.getValue()); } } return true; }
public boolean validateWatermarks(boolean exact, Map<String, CheckpointableWatermark> watermarkMap) { if (!watermarkMap.isEmpty()) { // watermark must be <= the index LongWatermark longWatermark = (LongWatermark) watermarkMap.values().iterator().next().getWatermark(); if (exact) { System.out.println(index-1 + ":" + longWatermark.getValue()); return ((index-1) == longWatermark.getValue()); } else { return (index > longWatermark.getValue()); } } return true; }
private static void verifyCommitables(FineGrainedWatermarkTracker tracker, SortedSet<Integer> holes, long maxWatermark) { // commitable should be the first hole -1 // uncommitable should be the first hole Map<String, CheckpointableWatermark> uncommitted = tracker.getUnacknowledgedWatermarks(); if (holes.isEmpty()) { Assert.assertEquals(uncommitted.size(), 0); } else { Assert.assertEquals(uncommitted.size(), 1); CheckpointableWatermark uncommitable = uncommitted.get("default"); Assert.assertEquals(((LongWatermark) uncommitable.getWatermark()).getValue(), (long) holes.first()); } Map<String, CheckpointableWatermark> commitables = tracker.getCommittableWatermarks(); if (holes.contains(0)) { // if the first record didn't get an ack Assert.assertEquals(commitables.size(), 0); } else { Assert.assertEquals(commitables.size(), 1); CheckpointableWatermark commitable = commitables.get("default"); if (holes.isEmpty()) { Assert.assertEquals(((LongWatermark) commitable.getWatermark()).getValue(), maxWatermark); } else { Assert.assertEquals(((LongWatermark) commitable.getWatermark()).getValue(), holes.first() - 1); } } }
@Test public void testSingleSource() { MultiWriterWatermarkTracker watermarkTracker = new MultiWriterWatermarkTracker(); commits(watermarkTracker, "default", 0, 4, 5, 6); Assert.assertEquals(watermarkTracker.getCommittableWatermark("default").get().getSource(), "default"); Assert.assertEquals(((LongWatermark) watermarkTracker.getCommittableWatermark("default") .get().getWatermark()).getValue(), 6L); }
@Test public void testMultiSource() { MultiWriterWatermarkTracker watermarkTracker = new MultiWriterWatermarkTracker(); commits(watermarkTracker, "default", 0, 4, 5, 6); commits(watermarkTracker, "other", 1, 3, 5, 7); Assert.assertEquals(watermarkTracker.getCommittableWatermark("default").get().getSource(), "default"); Assert.assertEquals(((LongWatermark) watermarkTracker.getCommittableWatermark("default") .get().getWatermark()).getValue(), 6L); Assert.assertEquals(watermarkTracker.getCommittableWatermark("other").get().getSource(), "other"); Assert.assertEquals(((LongWatermark) watermarkTracker.getCommittableWatermark("other") .get().getWatermark()).getValue(), 7L); }
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 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); }
@Override public int compareTo(CheckpointableWatermark o) { if (!(this.source.equals(o.getSource()))) { throw new RuntimeException("Could not compare two checkpointable watermarks because they have different sources " + this.source + ":" + o.getSource()); } return this.comparable.compareTo(o.getWatermark()); }
@Override public void start(WatermarkStorage watermarkStorage) throws IOException { this.watermarkStorage = Optional.of(watermarkStorage); Map<String, CheckpointableWatermark> lastCommitted; try { lastCommitted = this.watermarkStorage.get() .getCommittedWatermarks(DefaultCheckpointableWatermark.class, ImmutableList.of("" + extractor.partition)); } catch (IOException e) { // failed to get watermarks ... log a warning message log.warn("Failed to get watermarks... will start from the beginning", e); lastCommitted = Collections.EMPTY_MAP; } for (Map.Entry entry: lastCommitted.entrySet()) { log.info("{}: Found these committed watermarks: key: {}, value: {}", this, entry.getKey(), entry.getValue()); } LongWatermark currentWatermark; if (!lastCommitted.isEmpty() && lastCommitted.containsKey(""+extractor.partition)) { currentWatermark = (LongWatermark) (lastCommitted.get("" + extractor.partition)).getWatermark(); } else { // first record starts from 0 currentWatermark = new LongWatermark(-1); } extractor.setCurrentWatermark(currentWatermark); log.info("{}: Set current watermark to : {}", this, currentWatermark); } };