streams/src/main/java/org/apache/kafka/streams/processor/internals/ProcessorStateManager.java - kafka - Git at Google

 /**
  * Licensed to the Apache Software Foundation (ASF) under one or more
  * contributor license agreements.  See the NOTICE file distributed with
  * this work for additional information regarding copyright ownership.
  * The ASF licenses this file to You under the Apache License, Version 2.0
  * (the "License"); you may not use this file except in compliance with
  * the License.  You may obtain a copy of the License at
  * <p>
  * http://www.apache.org/licenses/LICENSE-2.0
  * <p>
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */

 package org.apache.kafka.streams.processor.internals;

 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.common.PartitionInfo;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.streams.errors.StreamsException;
 import org.apache.kafka.streams.processor.StateRestoreCallback;
 import org.apache.kafka.streams.processor.StateStore;
 import org.apache.kafka.streams.state.internals.OffsetCheckpoint;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 import java.io.File;
 import java.io.IOException;
 import java.io.RandomAccessFile;
 import java.nio.channels.FileChannel;
 import java.nio.channels.FileLock;
 import java.nio.channels.OverlappingFileLockException;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;

 import static java.util.Collections.singleton;

 public class ProcessorStateManager {

     private static final Logger log = LoggerFactory.getLogger(ProcessorStateManager.class);

     public static final String STATE_CHANGELOG_TOPIC_SUFFIX = "-changelog";
     public static final String CHECKPOINT_FILE_NAME = ".checkpoint";
     public static final String LOCK_FILE_NAME = ".lock";

     private final String applicationId;
     private final int defaultPartition;
     private final Map<String, TopicPartition> partitionForTopic;
     private final File baseDir;
     private final FileLock directoryLock;
     private final Map<String, StateStore> stores;
     private final Set<String> loggingEnabled;
     private final Consumer<byte[], byte[]> restoreConsumer;
     private final Map<TopicPartition, Long> restoredOffsets;
     private final Map<TopicPartition, Long> checkpointedOffsets;
     private final Map<TopicPartition, Long> offsetLimits;
     private final boolean isStandby;
     private final Map<String, StateRestoreCallback> restoreCallbacks; // used for standby tasks, keyed by state topic name

     /**
      * @throws IOException if any error happens while creating or locking the state directory
      */
     public ProcessorStateManager(String applicationId, int defaultPartition, Collection<TopicPartition> sources, File baseDir, Consumer<byte[], byte[]> restoreConsumer, boolean isStandby) throws IOException {
         this.applicationId = applicationId;
         this.defaultPartition = defaultPartition;
         this.partitionForTopic = new HashMap<>();
         for (TopicPartition source : sources) {
             this.partitionForTopic.put(source.topic(), source);
         }
         this.baseDir = baseDir;
         this.stores = new HashMap<>();
         this.loggingEnabled = new HashSet<>();
         this.restoreConsumer = restoreConsumer;
         this.restoredOffsets = new HashMap<>();
         this.isStandby = isStandby;
         this.restoreCallbacks = isStandby ? new HashMap<String, StateRestoreCallback>() : null;
         this.offsetLimits = new HashMap<>();

         // create the state directory for this task if missing (we won't create the parent directory)
         createStateDirectory(baseDir);

         // try to acquire the exclusive lock on the state directory
         directoryLock = lockStateDirectory(baseDir, 5);
         if (directoryLock == null) {
             throw new IOException("Failed to lock the state directory: " + baseDir.getCanonicalPath());
         }

         // load the checkpoint information
         OffsetCheckpoint checkpoint = new OffsetCheckpoint(new File(this.baseDir, CHECKPOINT_FILE_NAME));
         this.checkpointedOffsets = new HashMap<>(checkpoint.read());

         // delete the checkpoint file after finish loading its stored offsets
         checkpoint.delete();
     }

     private static void createStateDirectory(File stateDir) throws IOException {
         if (!stateDir.exists()) {
             stateDir.mkdir();
         }
     }

     public static String storeChangelogTopic(String applicationId, String storeName) {
         return applicationId + "-" + storeName + STATE_CHANGELOG_TOPIC_SUFFIX;
     }

     /**
      * @throws IOException if any error happens when locking the state directory
      */
     public static FileLock lockStateDirectory(File stateDir) throws IOException {
         return lockStateDirectory(stateDir, 0);
     }

     private static FileLock lockStateDirectory(File stateDir, int retry) throws IOException {
         File lockFile = new File(stateDir, ProcessorStateManager.LOCK_FILE_NAME);
         FileChannel channel = new RandomAccessFile(lockFile, "rw").getChannel();

         FileLock lock = lockStateDirectory(channel);
         while (lock == null && retry > 0) {
             try {
                 Thread.sleep(200);
             } catch (Exception ex) {
                 // do nothing
             }
             retry--;
             lock = lockStateDirectory(channel);
         }
         if (lock == null) {
             channel.close();
         }
         return lock;
     }

     private static FileLock lockStateDirectory(FileChannel channel) throws IOException {
         try {
             return channel.tryLock();
         } catch (OverlappingFileLockException e) {
             return null;
         }
     }

     public File baseDir() {
         return this.baseDir;
     }

     /**
      * @throws IllegalArgumentException if the store name has already been registered or if it is not a valid name
      * (e.g., when it conflicts with the names of internal topics, like the checkpoint file name)
      * @throws StreamsException if the store's change log does not contain the partition
      */
     public void register(StateStore store, boolean loggingEnabled, StateRestoreCallback stateRestoreCallback) {
         if (store.name().equals(CHECKPOINT_FILE_NAME))
             throw new IllegalArgumentException("Illegal store name: " + CHECKPOINT_FILE_NAME);

         if (this.stores.containsKey(store.name()))
             throw new IllegalArgumentException("Store " + store.name() + " has already been registered.");

         if (loggingEnabled)
             this.loggingEnabled.add(store.name());

         // check that the underlying change log topic exist or not
         String topic;
         if (loggingEnabled)
             topic = storeChangelogTopic(this.applicationId, store.name());
         else topic = store.name();

         // block until the partition is ready for this state changelog topic or time has elapsed
         int partition = getPartition(topic);
         boolean partitionNotFound = true;
         long startTime = System.currentTimeMillis();
         long waitTime = 5000L;      // hard-code the value since we should not block after KIP-4

         do {
             try {
                 Thread.sleep(50L);
             } catch (InterruptedException e) {
                 // ignore
             }

             for (PartitionInfo partitionInfo : restoreConsumer.partitionsFor(topic)) {
                 if (partitionInfo.partition() == partition) {
                     partitionNotFound = false;
                     break;
                 }
             }
         } while (partitionNotFound && System.currentTimeMillis() < startTime + waitTime);

         if (partitionNotFound)
             throw new StreamsException("Store " + store.name() + "'s change log (" + topic + ") does not contain partition " + partition);

         this.stores.put(store.name(), store);

         if (isStandby) {
             if (store.persistent())
                 restoreCallbacks.put(topic, stateRestoreCallback);
         } else {
             restoreActiveState(topic, stateRestoreCallback);
         }
     }

     private void restoreActiveState(String topicName, StateRestoreCallback stateRestoreCallback) {
         // ---- try to restore the state from change-log ---- //

         // subscribe to the store's partition
         if (!restoreConsumer.subscription().isEmpty()) {
             throw new IllegalStateException("Restore consumer should have not subscribed to any partitions beforehand");
         }
         TopicPartition storePartition = new TopicPartition(topicName, getPartition(topicName));
         restoreConsumer.assign(Collections.singletonList(storePartition));

         try {
             // calculate the end offset of the partition
             // TODO: this is a bit hacky to first seek then position to get the end offset
             restoreConsumer.seekToEnd(singleton(storePartition));
             long endOffset = restoreConsumer.position(storePartition);

             // restore from the checkpointed offset of the change log if it is persistent and the offset exists;
             // restore the state from the beginning of the change log otherwise
             if (checkpointedOffsets.containsKey(storePartition)) {
                 restoreConsumer.seek(storePartition, checkpointedOffsets.get(storePartition));
             } else {
                 restoreConsumer.seekToBeginning(singleton(storePartition));
             }

             // restore its state from changelog records
             long limit = offsetLimit(storePartition);
             while (true) {
                 long offset = 0L;
                 for (ConsumerRecord<byte[], byte[]> record : restoreConsumer.poll(100).records(storePartition)) {
                     offset = record.offset();
                     if (offset >= limit) break;
                     stateRestoreCallback.restore(record.key(), record.value());
                 }

                 if (offset >= limit) {
                     break;
                 } else if (restoreConsumer.position(storePartition) == endOffset) {
                     break;
                 } else if (restoreConsumer.position(storePartition) > endOffset) {
                     // For a logging enabled changelog (no offset limit),
                     // the log end offset should not change while restoring since it is only written by this thread.
                     throw new IllegalStateException("Log end offset should not change while restoring");
                 }
             }

             // record the restored offset for its change log partition
             long newOffset = Math.min(limit, restoreConsumer.position(storePartition));
             restoredOffsets.put(storePartition, newOffset);
         } finally {
             // un-assign the change log partition
             restoreConsumer.assign(Collections.<TopicPartition>emptyList());
         }
     }

     public Map<TopicPartition, Long> checkpointedOffsets() {
         Map<TopicPartition, Long> partitionsAndOffsets = new HashMap<>();

         for (Map.Entry<String, StateRestoreCallback> entry : restoreCallbacks.entrySet()) {
             String topicName = entry.getKey();
             int partition = getPartition(topicName);
             TopicPartition storePartition = new TopicPartition(topicName, partition);

             if (checkpointedOffsets.containsKey(storePartition)) {
                 partitionsAndOffsets.put(storePartition, checkpointedOffsets.get(storePartition));
             } else {
                 partitionsAndOffsets.put(storePartition, -1L);
             }
         }
         return partitionsAndOffsets;
     }

     public List<ConsumerRecord<byte[], byte[]>> updateStandbyStates(TopicPartition storePartition, List<ConsumerRecord<byte[], byte[]>> records) {
         long limit = offsetLimit(storePartition);
         List<ConsumerRecord<byte[], byte[]>> remainingRecords = null;

         // restore states from changelog records

         StateRestoreCallback restoreCallback = restoreCallbacks.get(storePartition.topic());

         long lastOffset = -1L;
         int count = 0;
         for (ConsumerRecord<byte[], byte[]> record : records) {
             if (record.offset() < limit) {
                 restoreCallback.restore(record.key(), record.value());
                 lastOffset = record.offset();
             } else {
                 if (remainingRecords == null)
                     remainingRecords = new ArrayList<>(records.size() - count);

                 remainingRecords.add(record);
             }
             count++;
         }
         // record the restored offset for its change log partition
         restoredOffsets.put(storePartition, lastOffset + 1);

         return remainingRecords;
     }

     public void putOffsetLimit(TopicPartition partition, long limit) {
         offsetLimits.put(partition, limit);
     }

     private long offsetLimit(TopicPartition partition) {
         Long limit = offsetLimits.get(partition);
         return limit != null ? limit : Long.MAX_VALUE;
     }

     public StateStore getStore(String name) {
         return stores.get(name);
     }

     public void flush() {
         if (!this.stores.isEmpty()) {
             log.debug("Flushing stores.");
             for (StateStore store : this.stores.values())
                 store.flush();
         }
     }

     /**
      * @throws IOException if any error happens when flushing or closing the state stores
      */
     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));

                     // only checkpoint the offset to the offsets file if it is persistent;
                     if (stores.get(storeName).persistent()) {
                         Long offset = ackedOffsets.get(part);

                         if (offset != null) {
                             // store the last offset + 1 (the log position after restoration)
                             checkpointOffsets.put(part, offset + 1);
                         } else {
                             // if no record was produced. we need to check the restored offset.
                             offset = restoredOffsets.get(part);
                             if (offset != null)
                                 checkpointOffsets.put(part, offset);
                         }
                     }
                 }

                 // write the checkpoint file before closing, to indicate clean shutdown
                 OffsetCheckpoint checkpoint = new OffsetCheckpoint(new File(this.baseDir, CHECKPOINT_FILE_NAME));
                 checkpoint.write(checkpointOffsets);
             }
         } finally {
             // release the state directory directoryLock
             directoryLock.release();
             directoryLock.channel().close();
         }
     }

     private int getPartition(String topic) {
         TopicPartition partition = partitionForTopic.get(topic);

         return partition == null ? defaultPartition : partition.partition();
     }
 }
	/**
	* Licensed to the Apache Software Foundation (ASF) under one or more
	* contributor license agreements. See the NOTICE file distributed with
	* this work for additional information regarding copyright ownership.
	* The ASF licenses this file to You under the Apache License, Version 2.0
	* (the "License"); you may not use this file except in compliance with
	* the License. You may obtain a copy of the License at
	* <p>
	* http://www.apache.org/licenses/LICENSE-2.0
	* <p>
	* Unless required by applicable law or agreed to in writing, software
	* distributed under the License is distributed on an "AS IS" BASIS,
	* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	* See the License for the specific language governing permissions and
	* limitations under the License.
	*/

	package org.apache.kafka.streams.processor.internals;

	import org.apache.kafka.clients.consumer.Consumer;
	import org.apache.kafka.clients.consumer.ConsumerRecord;
	import org.apache.kafka.common.PartitionInfo;
	import org.apache.kafka.common.TopicPartition;
	import org.apache.kafka.streams.errors.StreamsException;
	import org.apache.kafka.streams.processor.StateRestoreCallback;
	import org.apache.kafka.streams.processor.StateStore;
	import org.apache.kafka.streams.state.internals.OffsetCheckpoint;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;

	import java.io.File;
	import java.io.IOException;
	import java.io.RandomAccessFile;
	import java.nio.channels.FileChannel;
	import java.nio.channels.FileLock;
	import java.nio.channels.OverlappingFileLockException;
	import java.util.ArrayList;
	import java.util.Collection;
	import java.util.Collections;
	import java.util.HashMap;
	import java.util.HashSet;
	import java.util.List;
	import java.util.Map;
	import java.util.Set;

	import static java.util.Collections.singleton;

	public class ProcessorStateManager {

	private static final Logger log = LoggerFactory.getLogger(ProcessorStateManager.class);

	public static final String STATE_CHANGELOG_TOPIC_SUFFIX = "-changelog";
	public static final String CHECKPOINT_FILE_NAME = ".checkpoint";
	public static final String LOCK_FILE_NAME = ".lock";

	private final String applicationId;
	private final int defaultPartition;
	private final Map<String, TopicPartition> partitionForTopic;
	private final File baseDir;
	private final FileLock directoryLock;
	private final Map<String, StateStore> stores;
	private final Set<String> loggingEnabled;
	private final Consumer<byte[], byte[]> restoreConsumer;
	private final Map<TopicPartition, Long> restoredOffsets;
	private final Map<TopicPartition, Long> checkpointedOffsets;
	private final Map<TopicPartition, Long> offsetLimits;
	private final boolean isStandby;
	private final Map<String, StateRestoreCallback> restoreCallbacks; // used for standby tasks, keyed by state topic name

	/**
	* @throws IOException if any error happens while creating or locking the state directory
	*/
	public ProcessorStateManager(String applicationId, int defaultPartition, Collection<TopicPartition> sources, File baseDir, Consumer<byte[], byte[]> restoreConsumer, boolean isStandby) throws IOException {
	this.applicationId = applicationId;
	this.defaultPartition = defaultPartition;
	this.partitionForTopic = new HashMap<>();
	for (TopicPartition source : sources) {
	this.partitionForTopic.put(source.topic(), source);
	}
	this.baseDir = baseDir;
	this.stores = new HashMap<>();
	this.loggingEnabled = new HashSet<>();
	this.restoreConsumer = restoreConsumer;
	this.restoredOffsets = new HashMap<>();
	this.isStandby = isStandby;
	this.restoreCallbacks = isStandby ? new HashMap<String, StateRestoreCallback>() : null;
	this.offsetLimits = new HashMap<>();

	// create the state directory for this task if missing (we won't create the parent directory)
	createStateDirectory(baseDir);

	// try to acquire the exclusive lock on the state directory
	directoryLock = lockStateDirectory(baseDir, 5);
	if (directoryLock == null) {
	throw new IOException("Failed to lock the state directory: " + baseDir.getCanonicalPath());
	}

	// load the checkpoint information
	OffsetCheckpoint checkpoint = new OffsetCheckpoint(new File(this.baseDir, CHECKPOINT_FILE_NAME));
	this.checkpointedOffsets = new HashMap<>(checkpoint.read());

	// delete the checkpoint file after finish loading its stored offsets
	checkpoint.delete();
	}

	private static void createStateDirectory(File stateDir) throws IOException {
	if (!stateDir.exists()) {
	stateDir.mkdir();
	}
	}

	public static String storeChangelogTopic(String applicationId, String storeName) {
	return applicationId + "-" + storeName + STATE_CHANGELOG_TOPIC_SUFFIX;
	}

	/**
	* @throws IOException if any error happens when locking the state directory
	*/
	public static FileLock lockStateDirectory(File stateDir) throws IOException {
	return lockStateDirectory(stateDir, 0);
	}

	private static FileLock lockStateDirectory(File stateDir, int retry) throws IOException {
	File lockFile = new File(stateDir, ProcessorStateManager.LOCK_FILE_NAME);
	FileChannel channel = new RandomAccessFile(lockFile, "rw").getChannel();

	FileLock lock = lockStateDirectory(channel);
	while (lock == null && retry > 0) {
	try {
	Thread.sleep(200);
	} catch (Exception ex) {
	// do nothing
	}
	retry--;
	lock = lockStateDirectory(channel);
	}
	if (lock == null) {
	channel.close();
	}
	return lock;
	}

	private static FileLock lockStateDirectory(FileChannel channel) throws IOException {
	try {
	return channel.tryLock();
	} catch (OverlappingFileLockException e) {
	return null;
	}
	}

	public File baseDir() {
	return this.baseDir;
	}

	/**
	* @throws IllegalArgumentException if the store name has already been registered or if it is not a valid name
	* (e.g., when it conflicts with the names of internal topics, like the checkpoint file name)
	* @throws StreamsException if the store's change log does not contain the partition
	*/
	public void register(StateStore store, boolean loggingEnabled, StateRestoreCallback stateRestoreCallback) {
	if (store.name().equals(CHECKPOINT_FILE_NAME))
	throw new IllegalArgumentException("Illegal store name: " + CHECKPOINT_FILE_NAME);

	if (this.stores.containsKey(store.name()))
	throw new IllegalArgumentException("Store " + store.name() + " has already been registered.");

	if (loggingEnabled)
	this.loggingEnabled.add(store.name());

	// check that the underlying change log topic exist or not
	String topic;
	if (loggingEnabled)
	topic = storeChangelogTopic(this.applicationId, store.name());
	else topic = store.name();

	// block until the partition is ready for this state changelog topic or time has elapsed
	int partition = getPartition(topic);
	boolean partitionNotFound = true;
	long startTime = System.currentTimeMillis();
	long waitTime = 5000L; // hard-code the value since we should not block after KIP-4

	do {
	try {
	Thread.sleep(50L);
	} catch (InterruptedException e) {
	// ignore
	}

	for (PartitionInfo partitionInfo : restoreConsumer.partitionsFor(topic)) {
	if (partitionInfo.partition() == partition) {
	partitionNotFound = false;
	break;
	}
	}
	} while (partitionNotFound && System.currentTimeMillis() < startTime + waitTime);

	if (partitionNotFound)
	throw new StreamsException("Store " + store.name() + "'s change log (" + topic + ") does not contain partition " + partition);

	this.stores.put(store.name(), store);

	if (isStandby) {
	if (store.persistent())
	restoreCallbacks.put(topic, stateRestoreCallback);
	} else {
	restoreActiveState(topic, stateRestoreCallback);
	}
	}

	private void restoreActiveState(String topicName, StateRestoreCallback stateRestoreCallback) {
	// ---- try to restore the state from change-log ---- //

	// subscribe to the store's partition
	if (!restoreConsumer.subscription().isEmpty()) {
	throw new IllegalStateException("Restore consumer should have not subscribed to any partitions beforehand");
	}
	TopicPartition storePartition = new TopicPartition(topicName, getPartition(topicName));
	restoreConsumer.assign(Collections.singletonList(storePartition));

	try {
	// calculate the end offset of the partition
	// TODO: this is a bit hacky to first seek then position to get the end offset
	restoreConsumer.seekToEnd(singleton(storePartition));
	long endOffset = restoreConsumer.position(storePartition);

	// restore from the checkpointed offset of the change log if it is persistent and the offset exists;
	// restore the state from the beginning of the change log otherwise
	if (checkpointedOffsets.containsKey(storePartition)) {
	restoreConsumer.seek(storePartition, checkpointedOffsets.get(storePartition));
	} else {
	restoreConsumer.seekToBeginning(singleton(storePartition));
	}

	// restore its state from changelog records
	long limit = offsetLimit(storePartition);
	while (true) {
	long offset = 0L;
	for (ConsumerRecord<byte[], byte[]> record : restoreConsumer.poll(100).records(storePartition)) {
	offset = record.offset();
	if (offset >= limit) break;
	stateRestoreCallback.restore(record.key(), record.value());
	}

	if (offset >= limit) {
	break;
	} else if (restoreConsumer.position(storePartition) == endOffset) {
	break;
	} else if (restoreConsumer.position(storePartition) > endOffset) {
	// For a logging enabled changelog (no offset limit),
	// the log end offset should not change while restoring since it is only written by this thread.
	throw new IllegalStateException("Log end offset should not change while restoring");
	}
	}

	// record the restored offset for its change log partition
	long newOffset = Math.min(limit, restoreConsumer.position(storePartition));
	restoredOffsets.put(storePartition, newOffset);
	} finally {
	// un-assign the change log partition
	restoreConsumer.assign(Collections.<TopicPartition>emptyList());
	}
	}

	public Map<TopicPartition, Long> checkpointedOffsets() {
	Map<TopicPartition, Long> partitionsAndOffsets = new HashMap<>();

	for (Map.Entry<String, StateRestoreCallback> entry : restoreCallbacks.entrySet()) {
	String topicName = entry.getKey();
	int partition = getPartition(topicName);
	TopicPartition storePartition = new TopicPartition(topicName, partition);

	if (checkpointedOffsets.containsKey(storePartition)) {
	partitionsAndOffsets.put(storePartition, checkpointedOffsets.get(storePartition));
	} else {
	partitionsAndOffsets.put(storePartition, -1L);
	}
	}
	return partitionsAndOffsets;
	}

	public List<ConsumerRecord<byte[], byte[]>> updateStandbyStates(TopicPartition storePartition, List<ConsumerRecord<byte[], byte[]>> records) {
	long limit = offsetLimit(storePartition);
	List<ConsumerRecord<byte[], byte[]>> remainingRecords = null;

	// restore states from changelog records

	StateRestoreCallback restoreCallback = restoreCallbacks.get(storePartition.topic());

	long lastOffset = -1L;
	int count = 0;
	for (ConsumerRecord<byte[], byte[]> record : records) {
	if (record.offset() < limit) {
	restoreCallback.restore(record.key(), record.value());
	lastOffset = record.offset();
	} else {
	if (remainingRecords == null)
	remainingRecords = new ArrayList<>(records.size() - count);

	remainingRecords.add(record);
	}
	count++;
	}
	// record the restored offset for its change log partition
	restoredOffsets.put(storePartition, lastOffset + 1);

	return remainingRecords;
	}

	public void putOffsetLimit(TopicPartition partition, long limit) {
	offsetLimits.put(partition, limit);
	}

	private long offsetLimit(TopicPartition partition) {
	Long limit = offsetLimits.get(partition);
	return limit != null ? limit : Long.MAX_VALUE;
	}

	public StateStore getStore(String name) {
	return stores.get(name);
	}

	public void flush() {
	if (!this.stores.isEmpty()) {
	log.debug("Flushing stores.");
	for (StateStore store : this.stores.values())
	store.flush();
	}
	}

	/**
	* @throws IOException if any error happens when flushing or closing the state stores
	*/
	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));

	// only checkpoint the offset to the offsets file if it is persistent;
	if (stores.get(storeName).persistent()) {
	Long offset = ackedOffsets.get(part);

	if (offset != null) {
	// store the last offset + 1 (the log position after restoration)
	checkpointOffsets.put(part, offset + 1);
	} else {
	// if no record was produced. we need to check the restored offset.
	offset = restoredOffsets.get(part);
	if (offset != null)
	checkpointOffsets.put(part, offset);
	}
	}
	}

	// write the checkpoint file before closing, to indicate clean shutdown
	OffsetCheckpoint checkpoint = new OffsetCheckpoint(new File(this.baseDir, CHECKPOINT_FILE_NAME));
	checkpoint.write(checkpointOffsets);
	}
	} finally {
	// release the state directory directoryLock
	directoryLock.release();
	directoryLock.channel().close();
	}
	}

	private int getPartition(String topic) {
	TopicPartition partition = partitionForTopic.get(topic);

	return partition == null ? defaultPartition : partition.partition();
	}
	}