WatermarkStates { void update(WatermarkMessage watermarkMessage, SystemStreamPartition ssp) { WatermarkState state = watermarkStates.get(ssp); if (state != null) { state.update(watermarkMessage.getTimestamp(), watermarkMessage.getTaskName()); } else { LOG.error("SSP {} doesn't have watermark states", ssp); } } WatermarkStates(
      Set<SystemStreamPartition> ssps,
      Map<SystemStream, Integer> producerTaskCounts,
      MetricsRegistry metricsRegistry);  static final long WATERMARK_NOT_EXIST; }
@Test public void testUpdate() { SystemStream input = new SystemStream("system", "input"); SystemStream intermediate = new SystemStream("system", "intermediate"); Set<SystemStreamPartition> ssps = new HashSet<>(); SystemStreamPartition inputPartition0 = new SystemStreamPartition(input, new Partition(0)); SystemStreamPartition intPartition0 = new SystemStreamPartition(intermediate, new Partition(0)); SystemStreamPartition intPartition1 = new SystemStreamPartition(intermediate, new Partition(1)); ssps.add(inputPartition0); ssps.add(intPartition0); ssps.add(intPartition1); Map<SystemStream, Integer> producerCounts = new HashMap<>(); producerCounts.put(intermediate, 2); WatermarkStates watermarkStates = new WatermarkStates(ssps, producerCounts, new MetricsRegistryMap()); IncomingMessageEnvelope envelope = IncomingMessageEnvelope.buildWatermarkEnvelope(inputPartition0, 5L); watermarkStates.update((WatermarkMessage) envelope.getMessage(), envelope.getSystemStreamPartition()); assertEquals(watermarkStates.getWatermark(input), 5L); assertEquals(watermarkStates.getWatermark(intermediate), WATERMARK_NOT_EXIST); WatermarkMessage watermarkMessage = new WatermarkMessage(6L, "task 0"); watermarkStates.update(watermarkMessage, intPartition0); assertEquals(watermarkStates.getWatermarkPerSSP(intPartition0), WATERMARK_NOT_EXIST); assertEquals(watermarkStates.getWatermark(intermediate), WATERMARK_NOT_EXIST); watermarkMessage = new WatermarkMessage(3L, "task 1"); watermarkStates.update(watermarkMessage, intPartition0); assertEquals(watermarkStates.getWatermarkPerSSP(intPartition0), 3L); assertEquals(watermarkStates.getWatermark(intermediate), WATERMARK_NOT_EXIST); watermarkMessage = new WatermarkMessage(10L, "task 0"); watermarkStates.update(watermarkMessage, intPartition1); assertEquals(watermarkStates.getWatermarkPerSSP(intPartition1), WATERMARK_NOT_EXIST); assertEquals(watermarkStates.getWatermark(intermediate), WATERMARK_NOT_EXIST); watermarkMessage = new WatermarkMessage(4L, "task 1"); watermarkStates.update(watermarkMessage, intPartition1); assertEquals(watermarkStates.getWatermarkPerSSP(intPartition1), 4L); assertEquals(watermarkStates.getWatermark(intermediate), 3L); watermarkMessage = new WatermarkMessage(8L, "task 1"); watermarkStates.update(watermarkMessage, intPartition0); assertEquals(watermarkStates.getWatermarkPerSSP(intPartition0), 6L); assertEquals(watermarkStates.getWatermark(intermediate), 4L); watermarkMessage = new WatermarkMessage(7L, "task 1"); watermarkStates.update(watermarkMessage, intPartition1); assertEquals(watermarkStates.getWatermarkPerSSP(intPartition1), 7L); assertEquals(watermarkStates.getWatermark(intermediate), 6L); }