ProcessorStateManager { public void close(Map<TopicPartition, Long> ackedOffsets) throws IOException { try { if (!stores.isEmpty()) { log.debug("Closing stores."); for (Map.Entry<String, StateStore> entry : stores.entrySet()) { log.debug("Closing storage engine {}", entry.getKey()); entry.getValue().flush(); entry.getValue().close(); } Map<TopicPartition, Long> checkpointOffsets = new HashMap<>(); for (String storeName : stores.keySet()) { TopicPartition part; if (loggingEnabled.contains(storeName)) part = new TopicPartition(storeChangelogTopic(applicationId, storeName), getPartition(storeName)); else part = new TopicPartition(storeName, getPartition(storeName)); if (stores.get(storeName).persistent()) { Long offset = ackedOffsets.get(part); if (offset != null) { checkpointOffsets.put(part, offset + 1); } else { offset = restoredOffsets.get(part); if (offset != null) checkpointOffsets.put(part, offset); } } } OffsetCheckpoint checkpoint = new OffsetCheckpoint(new File(this.baseDir, CHECKPOINT_FILE_NAME)); checkpoint.write(checkpointOffsets); } } finally { directoryLock.release(); directoryLock.channel().close(); } } ProcessorStateManager(String applicationId, int defaultPartition, Collection<TopicPartition> sources, File baseDir, Consumer<byte[], byte[]> restoreConsumer, boolean isStandby); static String storeChangelogTopic(String applicationId, String storeName); static FileLock lockStateDirectory(File stateDir); File baseDir(); void register(StateStore store, boolean loggingEnabled, StateRestoreCallback stateRestoreCallback); Map<TopicPartition, Long> checkpointedOffsets(); List<ConsumerRecord<byte[], byte[]>> updateStandbyStates(TopicPartition storePartition, List<ConsumerRecord<byte[], byte[]>> records); void putOffsetLimit(TopicPartition partition, long limit); StateStore getStore(String name); void flush(); void close(Map<TopicPartition, Long> ackedOffsets); static final String STATE_CHANGELOG_TOPIC_SUFFIX; static final String CHECKPOINT_FILE_NAME; static final String LOCK_FILE_NAME; }
@Test public void testClose() throws IOException { File baseDir = Files.createTempDirectory(stateDir).toFile(); File checkpointFile = new File(baseDir, ProcessorStateManager.CHECKPOINT_FILE_NAME); try { OffsetCheckpoint oldCheckpoint = new OffsetCheckpoint(checkpointFile); oldCheckpoint.write(Collections.<TopicPartition, Long>emptyMap()); MockRestoreConsumer restoreConsumer = new MockRestoreConsumer(); restoreConsumer.updatePartitions(persistentStoreTopicName, Utils.mkList( new PartitionInfo(persistentStoreTopicName, 1, Node.noNode(), new Node[0], new Node[0]) )); restoreConsumer.updatePartitions(nonPersistentStoreTopicName, Utils.mkList( new PartitionInfo(nonPersistentStoreTopicName, 1, Node.noNode(), new Node[0], new Node[0]) )); HashMap<TopicPartition, Long> ackedOffsets = new HashMap<>(); ackedOffsets.put(new TopicPartition(persistentStoreTopicName, 1), 123L); ackedOffsets.put(new TopicPartition(nonPersistentStoreTopicName, 1), 456L); ackedOffsets.put(new TopicPartition(ProcessorStateManager.storeChangelogTopic(applicationId, "otherTopic"), 1), 789L); MockStateStoreSupplier.MockStateStore persistentStore = new MockStateStoreSupplier.MockStateStore(persistentStoreName, true); MockStateStoreSupplier.MockStateStore nonPersistentStore = new MockStateStoreSupplier.MockStateStore(nonPersistentStoreName, false); ProcessorStateManager stateMgr = new ProcessorStateManager(applicationId, 1, noPartitions, baseDir, restoreConsumer, false); try { assertFalse(checkpointFile.exists()); restoreConsumer.reset(); stateMgr.register(persistentStore, true, persistentStore.stateRestoreCallback); restoreConsumer.reset(); stateMgr.register(nonPersistentStore, true, nonPersistentStore.stateRestoreCallback); } finally { stateMgr.close(ackedOffsets); } assertTrue(persistentStore.flushed); assertTrue(persistentStore.closed); assertTrue(nonPersistentStore.flushed); assertTrue(nonPersistentStore.closed); assertTrue(checkpointFile.exists()); OffsetCheckpoint newCheckpoint = new OffsetCheckpoint(checkpointFile); Map<TopicPartition, Long> checkpointedOffsets = newCheckpoint.read(); assertEquals(1, checkpointedOffsets.size()); assertEquals(new Long(123L + 1L), checkpointedOffsets.get(new TopicPartition(persistentStoreTopicName, 1))); } finally { Utils.delete(baseDir); } }