PartitionedWriter implements Closeable { public Map<Integer, Long> getStartingOffsets(Collection<Integer> partitionsId) throws IOException { synchronized (perPartitionStartOffset) { if (!perPartitionStartOffset.keySet().containsAll(partitionsId)) { final Map<Integer, Long> startingOffsets; try { startingOffsets = offsetComputer.computeOffsets(partitionsId); } catch (IOException e) { partitionsId.forEach(id -> perPartitionStartOffset.put(id, OffsetComputer.NO_OFFSET)); throw e; } perPartitionStartOffset.putAll(startingOffsets); return startingOffsets; } return partitionsId.stream() .filter(perPartitionStartOffset::containsKey) .collect(Collectors.toMap(Function.identity(), perPartitionStartOffset::get)); } } PartitionedWriter(BiFunction<Integer, LocalDateTime, ExpiringConsumer<MESSAGE_KIND>> writerBuilder,
                             OffsetComputer offsetComputer, String eventName, Message.Builder emptyMessageBuilder,
                             Checkpointer checkpointer); void dropPartition(int partitionId); void write(Instant when, Offset offset, MESSAGE_KIND msg); @Override void close(); Map<Integer, Long> getStartingOffsets(Collection<Integer> partitionsId); void heartbeat(int partition, Offset offset);  }
@Test public void getStartingOffsets() throws IOException { final Map<Integer, Long> firstOffsets = new HashMap<>(); final Map<Integer, Long> secondOffsets = new HashMap<>(); final int firstPartition = 42; final int secondPartition = 51; final OffsetComputer offsetComputer = mock(OffsetComputer.class); final PartitionedWriter<String> writer = new PartitionedWriter<>(mock(BiFunction.class), offsetComputer, "ignored", DataAccessEventProtos.FsEvent.newBuilder(), mock(Checkpointer.class)); firstOffsets.put(firstPartition, 2L); firstOffsets.put(secondPartition, 3L); secondOffsets.put(firstPartition, 4L); secondOffsets.put(secondPartition, 5L); when(offsetComputer.computeOffsets(any())).thenReturn(firstOffsets).thenReturn(secondOffsets); final List<Integer> bothPartitions = Arrays.asList(firstPartition, secondPartition); assertEquals(firstOffsets, writer.getStartingOffsets(bothPartitions)); assertEquals(firstOffsets, writer.getStartingOffsets(bothPartitions)); verify(offsetComputer, times(1)).computeOffsets(bothPartitions); writer.dropPartition(firstPartition); assertEquals(secondOffsets, writer.getStartingOffsets(bothPartitions)); }