@Override public void unacknowledgedWatermark(CheckpointableWatermark unacked) { if (_unackedWatermarkMap != null) { _unackedWatermarkMap.put(unacked.getSource(), unacked); } }
@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 int compareTo(AcknowledgableWatermark o) { return _checkpointableWatermark.compareTo(o._checkpointableWatermark); }
@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; }
@Override public void committedWatermark(CheckpointableWatermark committed) { _committedWatermarkMap.put(committed.getSource(), committed); }
@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 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 Optional<CheckpointableWatermark> getCommittableWatermark(String source) { Set<CheckpointableWatermark> unacked = unacknowledgedWatermarks.get(source); CheckpointableWatermark minUnacknowledgedWatermark = (unacked == null || unacked.isEmpty())? null: unacked.iterator().next(); CheckpointableWatermark highestCommitableWatermark = null; for (CheckpointableWatermark commitableWatermark : candidateCommittables.get(source)) { if ((minUnacknowledgedWatermark == null) || (commitableWatermark.compareTo(minUnacknowledgedWatermark) < 0)) { // commitableWatermark < minUnacknowledgedWatermark highestCommitableWatermark = commitableWatermark; } } if (highestCommitableWatermark == null) { return Optional.absent(); } else { return Optional.of(highestCommitableWatermark); } }
@Override public Map<String, CheckpointableWatermark> getAllCommitableWatermarks() { Map<String, CheckpointableWatermark> commitables = new HashMap<>(candidateCommittables.size()); for (String source: candidateCommittables.keySet()) { Optional<CheckpointableWatermark> commitable = getCommittableWatermark(source); if (commitable.isPresent()) { commitables.put(commitable.get().getSource(), commitable.get()); } } return commitables; }
@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); }
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); } } }
@Override public int compareTo(AcknowledgableWatermark o) { return _checkpointableWatermark.compareTo(o._checkpointableWatermark); }
/** * Create a CheckpointableWatermarkState object from a CheckpointableWatermark * @param watermark: the checkpointable watermark * @param gson: the instance of {@link Gson} to use for serializing the {@param watermark}. */ public CheckpointableWatermarkState(CheckpointableWatermark watermark, Gson gson) { super.setProp(watermark.getSource(), gson.toJsonTree(watermark)); super.setId(watermark.getSource()); }
@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()); }
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()); } }
public Optional<CheckpointableWatermark> getCommittableWatermark(String source) { Set<CheckpointableWatermark> unacked = unacknowledgedWatermarks.get(source); CheckpointableWatermark minUnacknowledgedWatermark = (unacked == null || unacked.isEmpty())? null: unacked.iterator().next(); CheckpointableWatermark highestCommitableWatermark = null; for (CheckpointableWatermark commitableWatermark : candidateCommittables.get(source)) { if ((minUnacknowledgedWatermark == null) || (commitableWatermark.compareTo(minUnacknowledgedWatermark) < 0)) { // commitableWatermark < minUnacknowledgedWatermark highestCommitableWatermark = commitableWatermark; } } if (highestCommitableWatermark == null) { return Optional.absent(); } else { return Optional.of(highestCommitableWatermark); } }
@Override public Map<String, CheckpointableWatermark> getAllUnacknowledgedWatermarks() { Map<String, CheckpointableWatermark> unackedMap = new HashMap<>(unacknowledgedWatermarks.size()); for (String source: unacknowledgedWatermarks.keySet()) { Optional<CheckpointableWatermark> unacked = getUnacknowledgedWatermark(source); if (unacked.isPresent()) { unackedMap.put(unacked.get().getSource(), unacked.get()); } } return unackedMap; }
@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 void committedWatermark(CheckpointableWatermark committed) { getOrCreate(candidateCommittables, committed.getSource()).add(committed); }