/*
 * 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
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * 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.flink.contrib.streaming.state;

import org.apache.flink.api.common.typeutils.base.IntSerializer;
import org.apache.flink.configuration.CheckpointingOptions;
import org.apache.flink.configuration.ConfigOption;
import org.apache.flink.configuration.ConfigOptions;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.configuration.CoreOptions;
import org.apache.flink.configuration.ReadableConfig;
import org.apache.flink.core.fs.CloseableRegistry;
import org.apache.flink.core.fs.FileSystem;
import org.apache.flink.core.fs.Path;
import org.apache.flink.metrics.groups.UnregisteredMetricsGroup;
import org.apache.flink.runtime.execution.Environment;
import org.apache.flink.runtime.io.disk.iomanager.IOManager;
import org.apache.flink.runtime.jobgraph.JobVertexID;
import org.apache.flink.runtime.operators.testutils.MockEnvironment;
import org.apache.flink.runtime.operators.testutils.MockEnvironmentBuilder;
import org.apache.flink.runtime.query.KvStateRegistry;
import org.apache.flink.runtime.state.AbstractKeyedStateBackend;
import org.apache.flink.runtime.state.KeyGroupRange;
import org.apache.flink.runtime.state.StateBackend;
import org.apache.flink.runtime.state.filesystem.FsStateBackend;
import org.apache.flink.runtime.state.heap.HeapPriorityQueueSetFactory;
import org.apache.flink.runtime.state.memory.MemoryStateBackend;
import org.apache.flink.runtime.state.ttl.TtlTimeProvider;
import org.apache.flink.runtime.util.TestingTaskManagerRuntimeInfo;
import org.apache.flink.util.IOUtils;

import org.junit.Assert;
import org.junit.Assume;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;
import org.rocksdb.BlockBasedTableConfig;
import org.rocksdb.ColumnFamilyOptions;
import org.rocksdb.CompactionStyle;
import org.rocksdb.DBOptions;
import org.rocksdb.util.SizeUnit;

import java.io.File;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;

import static org.apache.flink.contrib.streaming.state.RocksDBTestUtils.createKeyedStateBackend;
import static org.hamcrest.CoreMatchers.anyOf;
import static org.hamcrest.CoreMatchers.instanceOf;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.startsWith;
import static org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNotEquals;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

/**
 * Tests for configuring the RocksDB State Backend.
 */
@SuppressWarnings("serial")
public class RocksDBStateBackendConfigTest {

	@Rule
	public final TemporaryFolder tempFolder = new TemporaryFolder();

	// ------------------------------------------------------------------------
	//  default values
	// ------------------------------------------------------------------------

	@Test
	public void testDefaultsInSync() throws Exception {
		final boolean defaultIncremental = CheckpointingOptions.INCREMENTAL_CHECKPOINTS.defaultValue();

		RocksDBStateBackend backend = new RocksDBStateBackend(tempFolder.newFolder().toURI());
		assertEquals(defaultIncremental, backend.isIncrementalCheckpointsEnabled());
	}

	// ------------------------------------------------------------------------
	//  RocksDB local file directory
	// ------------------------------------------------------------------------

	/**
	 * This test checks the behavior for basic setting of local DB directories.
	 */
	@Test
	public void testSetDbPath() throws Exception {
		final RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(tempFolder.newFolder().toURI().toString());

		final String testDir1 = tempFolder.newFolder().getAbsolutePath();
		final String testDir2 = tempFolder.newFolder().getAbsolutePath();

		assertNull(rocksDbBackend.getDbStoragePaths());

		rocksDbBackend.setDbStoragePath(testDir1);
		assertArrayEquals(new String[] { testDir1 }, rocksDbBackend.getDbStoragePaths());

		rocksDbBackend.setDbStoragePath(null);
		assertNull(rocksDbBackend.getDbStoragePaths());

		rocksDbBackend.setDbStoragePaths(testDir1, testDir2);
		assertArrayEquals(new String[] { testDir1, testDir2 }, rocksDbBackend.getDbStoragePaths());

		final MockEnvironment env = getMockEnvironment(tempFolder.newFolder());
		final RocksDBKeyedStateBackend<Integer> keyedBackend = createKeyedStateBackend(rocksDbBackend, env, IntSerializer.INSTANCE);

		try {
			File instanceBasePath = keyedBackend.getInstanceBasePath();
			assertThat(instanceBasePath.getAbsolutePath(), anyOf(startsWith(testDir1), startsWith(testDir2)));

			//noinspection NullArgumentToVariableArgMethod
			rocksDbBackend.setDbStoragePaths(null);
			assertNull(rocksDbBackend.getDbStoragePaths());
		}
		finally {
			IOUtils.closeQuietly(keyedBackend);
			keyedBackend.dispose();
			env.close();
		}
	}

	@Test
	public void testConfigureTimerService() throws Exception {

		final MockEnvironment env = getMockEnvironment(tempFolder.newFolder());

		// Fix the option key string
		Assert.assertEquals("state.backend.rocksdb.timer-service.factory", RocksDBOptions.TIMER_SERVICE_FACTORY.key());

		// Fix the option value string and ensure all are covered
		Assert.assertEquals(2, RocksDBStateBackend.PriorityQueueStateType.values().length);
		Assert.assertEquals("ROCKSDB", RocksDBStateBackend.PriorityQueueStateType.ROCKSDB.toString());
		Assert.assertEquals("HEAP", RocksDBStateBackend.PriorityQueueStateType.HEAP.toString());

		// Fix the default
		Assert.assertEquals(
			RocksDBStateBackend.PriorityQueueStateType.ROCKSDB,
			RocksDBOptions.TIMER_SERVICE_FACTORY.defaultValue());

		RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(tempFolder.newFolder().toURI().toString());

		RocksDBKeyedStateBackend<Integer> keyedBackend = createKeyedStateBackend(rocksDbBackend, env, IntSerializer.INSTANCE);
		Assert.assertEquals(RocksDBPriorityQueueSetFactory.class, keyedBackend.getPriorityQueueFactory().getClass());
		keyedBackend.dispose();

		Configuration conf = new Configuration();
		conf.set(RocksDBOptions.TIMER_SERVICE_FACTORY, RocksDBStateBackend.PriorityQueueStateType.HEAP);

		rocksDbBackend = rocksDbBackend.configure(conf, Thread.currentThread().getContextClassLoader());
		keyedBackend = createKeyedStateBackend(rocksDbBackend, env, IntSerializer.INSTANCE);
		Assert.assertEquals(
			HeapPriorityQueueSetFactory.class,
			keyedBackend.getPriorityQueueFactory().getClass());
		keyedBackend.dispose();
		env.close();
	}

	/**
	 * Validates that user custom configuration from code should override the flink-conf.yaml.
	 */
	@Test
	public void testConfigureTimerServiceLoadingFromApplication() throws Exception {
		final MockEnvironment env = new MockEnvironmentBuilder().build();

		// priorityQueueStateType of the job backend
		final RocksDBStateBackend backend = new RocksDBStateBackend(tempFolder.newFolder().toURI().toString());
		backend.setPriorityQueueStateType(RocksDBStateBackend.PriorityQueueStateType.HEAP);

		// priorityQueueStateType in the cluster config
		final Configuration configFromConfFile = new Configuration();
		configFromConfFile.setString(
			RocksDBOptions.TIMER_SERVICE_FACTORY.key(),
			RocksDBStateBackend.PriorityQueueStateType.ROCKSDB.toString());

		// configure final backend from job and cluster config
		final RocksDBStateBackend configuredRocksDBStateBackend = backend.configure(
			configFromConfFile,
			Thread.currentThread().getContextClassLoader());
		final RocksDBKeyedStateBackend<Integer> keyedBackend = createKeyedStateBackend(configuredRocksDBStateBackend, env, IntSerializer.INSTANCE);

		// priorityQueueStateType of the job backend should be preserved
		assertThat(keyedBackend.getPriorityQueueFactory(), instanceOf(HeapPriorityQueueSetFactory.class));

		keyedBackend.close();
		keyedBackend.dispose();
		env.close();
	}

	@Test
	public void testStoragePathWithFilePrefix() throws Exception {
		final File folder = tempFolder.newFolder();
		final String dbStoragePath = new Path(folder.toURI().toString()).toString();

		assertTrue(dbStoragePath.startsWith("file:"));

		testLocalDbPaths(dbStoragePath, folder);
	}

	@Test
	public void testWithDefaultFsSchemeNoStoragePath() throws Exception {
		try {
			// set the default file system scheme
			Configuration config = new Configuration();
			config.setString(CoreOptions.DEFAULT_FILESYSTEM_SCHEME, "s3://mydomain.com:8020/flink");
			FileSystem.initialize(config);
			testLocalDbPaths(null, tempFolder.getRoot());
		}
		finally {
			FileSystem.initialize(new Configuration());
		}
	}

	@Test
	public void testWithDefaultFsSchemeAbsoluteStoragePath() throws Exception {
		final File folder = tempFolder.newFolder();
		final String dbStoragePath = folder.getAbsolutePath();

		try {
			// set the default file system scheme
			Configuration config = new Configuration();
			config.setString(CoreOptions.DEFAULT_FILESYSTEM_SCHEME, "s3://mydomain.com:8020/flink");
			FileSystem.initialize(config);

			testLocalDbPaths(dbStoragePath, folder);
		}
		finally {
			FileSystem.initialize(new Configuration());
		}
	}

	private void testLocalDbPaths(String configuredPath, File expectedPath) throws Exception {
		final RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(tempFolder.newFolder().toURI().toString());
		rocksDbBackend.setDbStoragePath(configuredPath);

		final MockEnvironment env = getMockEnvironment(tempFolder.newFolder());
		RocksDBKeyedStateBackend<Integer> keyedBackend = createKeyedStateBackend(rocksDbBackend, env, IntSerializer.INSTANCE);

		try {
			File instanceBasePath = keyedBackend.getInstanceBasePath();
			assertThat(instanceBasePath.getAbsolutePath(), startsWith(expectedPath.getAbsolutePath()));

			//noinspection NullArgumentToVariableArgMethod
			rocksDbBackend.setDbStoragePaths(null);
			assertNull(rocksDbBackend.getDbStoragePaths());
		} finally {
			IOUtils.closeQuietly(keyedBackend);
			keyedBackend.dispose();
			env.close();
		}
	}

	/**
	 * Validates that empty arguments for the local DB path are invalid.
	 */
	@Test(expected = IllegalArgumentException.class)
	public void testSetEmptyPaths() throws Exception {
		String checkpointPath = tempFolder.newFolder().toURI().toString();
		RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(checkpointPath);
		rocksDbBackend.setDbStoragePaths();
	}

	/**
	 * Validates that schemes other than 'file:/' are not allowed.
	 */
	@Test(expected = IllegalArgumentException.class)
	public void testNonFileSchemePath() throws Exception {
		String checkpointPath = tempFolder.newFolder().toURI().toString();
		RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(checkpointPath);
		rocksDbBackend.setDbStoragePath("hdfs:///some/path/to/perdition");
	}

	@Test(expected = IllegalArgumentException.class)
	public void testDbPathRelativePaths() throws Exception {
		RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(tempFolder.newFolder().toURI().toString());
		rocksDbBackend.setDbStoragePath("relative/path");
	}

	// ------------------------------------------------------------------------
	//  RocksDB local file automatic from temp directories
	// ------------------------------------------------------------------------

	/**
	 * This tests whether the RocksDB backends uses the temp directories that are provided
	 * from the {@link Environment} when no db storage path is set.
	 */
	@Test
	public void testUseTempDirectories() throws Exception {
		String checkpointPath = tempFolder.newFolder().toURI().toString();
		RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(checkpointPath);

		File dir1 = tempFolder.newFolder();
		File dir2 = tempFolder.newFolder();

		assertNull(rocksDbBackend.getDbStoragePaths());

		final MockEnvironment env = getMockEnvironment(dir1, dir2);
		RocksDBKeyedStateBackend<Integer> keyedBackend = (RocksDBKeyedStateBackend<Integer>) rocksDbBackend.
			createKeyedStateBackend(
				env,
				env.getJobID(),
				"test_op",
				IntSerializer.INSTANCE,
				1,
				new KeyGroupRange(0, 0),
				env.getTaskKvStateRegistry(),
				TtlTimeProvider.DEFAULT,
				new UnregisteredMetricsGroup(),
				Collections.emptyList(),
				new CloseableRegistry());

		try {
			File instanceBasePath = keyedBackend.getInstanceBasePath();
			assertThat(instanceBasePath.getAbsolutePath(), anyOf(startsWith(dir1.getAbsolutePath()), startsWith(dir2.getAbsolutePath())));
		} finally {
			IOUtils.closeQuietly(keyedBackend);
			keyedBackend.dispose();
			env.close();
		}
	}

	// ------------------------------------------------------------------------
	//  RocksDB local file directory initialization
	// ------------------------------------------------------------------------

	@Test
	public void testFailWhenNoLocalStorageDir() throws Exception {
		final File targetDir = tempFolder.newFolder();
		Assume.assumeTrue("Cannot mark directory non-writable", targetDir.setWritable(false, false));

		String checkpointPath = tempFolder.newFolder().toURI().toString();
		RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(checkpointPath);

		try (MockEnvironment env = getMockEnvironment(tempFolder.newFolder())) {
			rocksDbBackend.setDbStoragePath(targetDir.getAbsolutePath());

			boolean hasFailure = false;
			try {
				rocksDbBackend.createKeyedStateBackend(
					env,
					env.getJobID(),
					"foobar",
					IntSerializer.INSTANCE,
					1,
					new KeyGroupRange(0, 0),
					new KvStateRegistry().createTaskRegistry(env.getJobID(), new JobVertexID()),
					TtlTimeProvider.DEFAULT,
					new UnregisteredMetricsGroup(),
					Collections.emptyList(),
					new CloseableRegistry());
			} catch (Exception e) {
				assertTrue(e.getMessage().contains("No local storage directories available"));
				assertTrue(e.getMessage().contains(targetDir.getAbsolutePath()));
				hasFailure = true;
			}
			assertTrue("We must see a failure because no storaged directory is feasible.", hasFailure);
		} finally {
			//noinspection ResultOfMethodCallIgnored
			targetDir.setWritable(true, false);
		}
	}

	@Test
	public void testContinueOnSomeDbDirectoriesMissing() throws Exception {
		final File targetDir1 = tempFolder.newFolder();
		final File targetDir2 = tempFolder.newFolder();
		Assume.assumeTrue("Cannot mark directory non-writable", targetDir1.setWritable(false, false));

		String checkpointPath = tempFolder.newFolder().toURI().toString();
		RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(checkpointPath);

		try (MockEnvironment env = getMockEnvironment(tempFolder.newFolder())) {
			rocksDbBackend.setDbStoragePaths(targetDir1.getAbsolutePath(), targetDir2.getAbsolutePath());

			try {
				AbstractKeyedStateBackend<Integer> keyedStateBackend = rocksDbBackend.createKeyedStateBackend(
					env,
					env.getJobID(),
					"foobar",
					IntSerializer.INSTANCE,
					1,
					new KeyGroupRange(0, 0),
					new KvStateRegistry().createTaskRegistry(env.getJobID(), new JobVertexID()),
					TtlTimeProvider.DEFAULT,
					new UnregisteredMetricsGroup(),
					Collections.emptyList(),
					new CloseableRegistry());

				IOUtils.closeQuietly(keyedStateBackend);
				keyedStateBackend.dispose();
			}
			catch (Exception e) {
				e.printStackTrace();
				fail("Backend initialization failed even though some paths were available");
			}
		} finally {
			//noinspection ResultOfMethodCallIgnored
			targetDir1.setWritable(true, false);
		}
	}

	// ------------------------------------------------------------------------
	//  RocksDB Options
	// ------------------------------------------------------------------------

	@Test
	public void testPredefinedOptions() throws Exception {
		String checkpointPath = tempFolder.newFolder().toURI().toString();
		RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(checkpointPath);

		// verify that we would use PredefinedOptions.DEFAULT by default.
		assertEquals(PredefinedOptions.DEFAULT, rocksDbBackend.getPredefinedOptions());

		// verify that user could configure predefined options via flink-conf.yaml
		Configuration configuration = new Configuration();
		configuration.setString(RocksDBOptions.PREDEFINED_OPTIONS, PredefinedOptions.FLASH_SSD_OPTIMIZED.name());
		rocksDbBackend = new RocksDBStateBackend(checkpointPath);
		rocksDbBackend = rocksDbBackend.configure(configuration, getClass().getClassLoader());
		assertEquals(PredefinedOptions.FLASH_SSD_OPTIMIZED, rocksDbBackend.getPredefinedOptions());

		// verify that predefined options could be set programmatically and override pre-configured one.
		rocksDbBackend.setPredefinedOptions(PredefinedOptions.SPINNING_DISK_OPTIMIZED);
		assertEquals(PredefinedOptions.SPINNING_DISK_OPTIMIZED, rocksDbBackend.getPredefinedOptions());
	}

	@Test
	public void testSetConfigurableOptions() throws Exception  {
		DefaultConfigurableOptionsFactory customizedOptions = new DefaultConfigurableOptionsFactory()
			.setMaxBackgroundThreads(4)
			.setMaxOpenFiles(-1)
			.setCompactionStyle(CompactionStyle.LEVEL)
			.setUseDynamicLevelSize(true)
			.setTargetFileSizeBase("4MB")
			.setMaxSizeLevelBase("128 mb")
			.setWriteBufferSize("128 MB")
			.setMaxWriteBufferNumber(4)
			.setMinWriteBufferNumberToMerge(3)
			.setBlockSize("64KB")
			.setBlockCacheSize("512mb");

		try (RocksDBResourceContainer optionsContainer =
				new RocksDBResourceContainer(PredefinedOptions.DEFAULT, customizedOptions)) {

			DBOptions dbOptions = optionsContainer.getDbOptions();
			assertEquals(-1, dbOptions.maxOpenFiles());

			ColumnFamilyOptions columnOptions = optionsContainer.getColumnOptions();
			assertEquals(CompactionStyle.LEVEL, columnOptions.compactionStyle());
			assertTrue(columnOptions.levelCompactionDynamicLevelBytes());
			assertEquals(4 * SizeUnit.MB, columnOptions.targetFileSizeBase());
			assertEquals(128 * SizeUnit.MB, columnOptions.maxBytesForLevelBase());
			assertEquals(4, columnOptions.maxWriteBufferNumber());
			assertEquals(3, columnOptions.minWriteBufferNumberToMerge());

			BlockBasedTableConfig tableConfig = (BlockBasedTableConfig) columnOptions.tableFormatConfig();
			assertEquals(64 * SizeUnit.KB, tableConfig.blockSize());
			assertEquals(512 * SizeUnit.MB, tableConfig.blockCacheSize());
		}
	}

	@Test
	public void testConfigurableOptionsFromConfig() throws Exception {
		Configuration configuration = new Configuration();
		DefaultConfigurableOptionsFactory defaultOptionsFactory = new DefaultConfigurableOptionsFactory();
		assertTrue(defaultOptionsFactory.configure(configuration).getConfiguredOptions().isEmpty());

		// verify illegal configuration
		{
			verifyIllegalArgument(RocksDBConfigurableOptions.MAX_BACKGROUND_THREADS, "-1");
			verifyIllegalArgument(RocksDBConfigurableOptions.MAX_WRITE_BUFFER_NUMBER, "-1");
			verifyIllegalArgument(RocksDBConfigurableOptions.MIN_WRITE_BUFFER_NUMBER_TO_MERGE, "-1");

			verifyIllegalArgument(RocksDBConfigurableOptions.TARGET_FILE_SIZE_BASE, "0KB");
			verifyIllegalArgument(RocksDBConfigurableOptions.MAX_SIZE_LEVEL_BASE, "1BB");
			verifyIllegalArgument(RocksDBConfigurableOptions.WRITE_BUFFER_SIZE, "-1KB");
			verifyIllegalArgument(RocksDBConfigurableOptions.BLOCK_SIZE, "0MB");
			verifyIllegalArgument(RocksDBConfigurableOptions.BLOCK_CACHE_SIZE, "0");

			verifyIllegalArgument(RocksDBConfigurableOptions.USE_DYNAMIC_LEVEL_SIZE, "1");

			verifyIllegalArgument(RocksDBConfigurableOptions.COMPACTION_STYLE, "LEV");
		}

		// verify legal configuration
		{
			configuration.setString(RocksDBConfigurableOptions.COMPACTION_STYLE.key(), "level");
			configuration.setString(RocksDBConfigurableOptions.USE_DYNAMIC_LEVEL_SIZE.key(), "TRUE");
			configuration.setString(RocksDBConfigurableOptions.TARGET_FILE_SIZE_BASE.key(), "8 mb");
			configuration.setString(RocksDBConfigurableOptions.MAX_SIZE_LEVEL_BASE.key(), "128MB");
			configuration.setString(RocksDBConfigurableOptions.MAX_BACKGROUND_THREADS.key(), "4");
			configuration.setString(RocksDBConfigurableOptions.MAX_WRITE_BUFFER_NUMBER.key(), "4");
			configuration.setString(RocksDBConfigurableOptions.MIN_WRITE_BUFFER_NUMBER_TO_MERGE.key(), "2");
			configuration.setString(RocksDBConfigurableOptions.WRITE_BUFFER_SIZE.key(), "64 MB");
			configuration.setString(RocksDBConfigurableOptions.BLOCK_SIZE.key(), "4 kb");
			configuration.setString(RocksDBConfigurableOptions.BLOCK_CACHE_SIZE.key(), "512 mb");

			DefaultConfigurableOptionsFactory optionsFactory = new DefaultConfigurableOptionsFactory();
			optionsFactory.configure(configuration);

			try (RocksDBResourceContainer optionsContainer =
					new RocksDBResourceContainer(PredefinedOptions.DEFAULT, optionsFactory)) {

				DBOptions dbOptions = optionsContainer.getDbOptions();
				assertEquals(-1, dbOptions.maxOpenFiles());

				ColumnFamilyOptions columnOptions = optionsContainer.getColumnOptions();
				assertEquals(CompactionStyle.LEVEL, columnOptions.compactionStyle());
				assertTrue(columnOptions.levelCompactionDynamicLevelBytes());
				assertEquals(8 * SizeUnit.MB, columnOptions.targetFileSizeBase());
				assertEquals(128 * SizeUnit.MB, columnOptions.maxBytesForLevelBase());
				assertEquals(4, columnOptions.maxWriteBufferNumber());
				assertEquals(2, columnOptions.minWriteBufferNumberToMerge());
				assertEquals(64 * SizeUnit.MB, columnOptions.writeBufferSize());

				BlockBasedTableConfig tableConfig = (BlockBasedTableConfig) columnOptions.tableFormatConfig();
				assertEquals(4 * SizeUnit.KB, tableConfig.blockSize());
				assertEquals(512 * SizeUnit.MB, tableConfig.blockCacheSize());
			}
		}
	}

	@Test
	public void testOptionsFactory() throws Exception {
		String checkpointPath = tempFolder.newFolder().toURI().toString();
		RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(checkpointPath);

		// verify that user-defined options factory could be configured via flink-conf.yaml
		Configuration config = new Configuration();
		config.setString(RocksDBOptions.OPTIONS_FACTORY.key(), TestOptionsFactory.class.getName());
		config.setString(TestOptionsFactory.BACKGROUND_JOBS_OPTION.key(), "4");

		rocksDbBackend = rocksDbBackend.configure(config, getClass().getClassLoader());

		assertTrue(rocksDbBackend.getRocksDBOptions() instanceof TestOptionsFactory);

		try (RocksDBResourceContainer optionsContainer = rocksDbBackend.createOptionsAndResourceContainer()) {
			DBOptions dbOptions = optionsContainer.getDbOptions();
			assertEquals(4, dbOptions.maxBackgroundJobs());
		}

		// verify that user-defined options factory could be set programmatically and override pre-configured one.
		rocksDbBackend.setRocksDBOptions(new RocksDBOptionsFactory() {
			@Override
			public DBOptions createDBOptions(DBOptions currentOptions, Collection<AutoCloseable> handlesToClose) {
				return currentOptions;
			}

			@Override
			public ColumnFamilyOptions createColumnOptions(ColumnFamilyOptions currentOptions, Collection<AutoCloseable> handlesToClose) {
				return currentOptions.setCompactionStyle(CompactionStyle.FIFO);
			}
		});

		try (RocksDBResourceContainer optionsContainer = rocksDbBackend.createOptionsAndResourceContainer()) {
			ColumnFamilyOptions colCreated = optionsContainer.getColumnOptions();
			assertEquals(CompactionStyle.FIFO, colCreated.compactionStyle());
		}
	}

	@Test
	public void testPredefinedAndOptionsFactory() throws Exception {
		final RocksDBOptionsFactory optionsFactory = new RocksDBOptionsFactory() {
			@Override
			public DBOptions createDBOptions(DBOptions currentOptions, Collection<AutoCloseable> handlesToClose) {
				return currentOptions;
			}

			@Override
			public ColumnFamilyOptions createColumnOptions(ColumnFamilyOptions currentOptions, Collection<AutoCloseable> handlesToClose) {
				return currentOptions.setCompactionStyle(CompactionStyle.UNIVERSAL);
			}
		};

		try (final RocksDBResourceContainer optionsContainer = new RocksDBResourceContainer(
				PredefinedOptions.SPINNING_DISK_OPTIMIZED, optionsFactory)) {

			final ColumnFamilyOptions columnFamilyOptions = optionsContainer.getColumnOptions();
			assertNotNull(columnFamilyOptions);
			assertEquals(CompactionStyle.UNIVERSAL, columnFamilyOptions.compactionStyle());
		}
	}

	@Test
	public void testPredefinedOptionsEnum() {
		ArrayList<AutoCloseable> handlesToClose = new ArrayList<>();
		for (PredefinedOptions o : PredefinedOptions.values()) {
			try (DBOptions opt = o.createDBOptions(handlesToClose)) {
				assertNotNull(opt);
			}
		}
		handlesToClose.forEach(IOUtils::closeQuietly);
		handlesToClose.clear();
	}

	// ------------------------------------------------------------------------
	//  Reconfiguration
	// ------------------------------------------------------------------------

	@Test
	public void testRocksDbReconfigurationCopiesExistingValues() throws Exception {
		final FsStateBackend checkpointBackend = new FsStateBackend(tempFolder.newFolder().toURI().toString());
		final boolean incremental = !CheckpointingOptions.INCREMENTAL_CHECKPOINTS.defaultValue();

		final RocksDBStateBackend original = new RocksDBStateBackend(checkpointBackend, incremental);

		// these must not be the default options
		final PredefinedOptions predOptions = PredefinedOptions.SPINNING_DISK_OPTIMIZED_HIGH_MEM;
		assertNotEquals(predOptions, original.getPredefinedOptions());
		original.setPredefinedOptions(predOptions);

		final RocksDBOptionsFactory optionsFactory = mock(RocksDBOptionsFactory.class);
		original.setRocksDBOptions(optionsFactory);

		final String[] localDirs = new String[] {
				tempFolder.newFolder().getAbsolutePath(), tempFolder.newFolder().getAbsolutePath() };
		original.setDbStoragePaths(localDirs);

		RocksDBStateBackend copy = original.configure(new Configuration(), Thread.currentThread().getContextClassLoader());

		assertEquals(original.isIncrementalCheckpointsEnabled(), copy.isIncrementalCheckpointsEnabled());
		assertArrayEquals(original.getDbStoragePaths(), copy.getDbStoragePaths());
		assertEquals(original.getRocksDBOptions(), copy.getRocksDBOptions());
		assertEquals(original.getPredefinedOptions(), copy.getPredefinedOptions());

		FsStateBackend copyCheckpointBackend = (FsStateBackend) copy.getCheckpointBackend();
		assertEquals(checkpointBackend.getCheckpointPath(), copyCheckpointBackend.getCheckpointPath());
		assertEquals(checkpointBackend.getSavepointPath(), copyCheckpointBackend.getSavepointPath());
	}

	// ------------------------------------------------------------------------
	//  RocksDB Memory Control
	// ------------------------------------------------------------------------

	@Test
	public void testDefaultMemoryControlParameters() {
		RocksDBMemoryConfiguration memSettings = new RocksDBMemoryConfiguration();
		assertTrue(memSettings.isUsingManagedMemory());
		assertFalse(memSettings.isUsingFixedMemoryPerSlot());
		assertEquals(RocksDBOptions.HIGH_PRIORITY_POOL_RATIO.defaultValue(), memSettings.getHighPriorityPoolRatio(), 0.0);
		assertEquals(RocksDBOptions.WRITE_BUFFER_RATIO.defaultValue(), memSettings.getWriteBufferRatio(), 0.0);

		RocksDBMemoryConfiguration configured = RocksDBMemoryConfiguration.fromOtherAndConfiguration(memSettings, new Configuration());
		assertTrue(configured.isUsingManagedMemory());
		assertFalse(configured.isUsingFixedMemoryPerSlot());
		assertEquals(RocksDBOptions.HIGH_PRIORITY_POOL_RATIO.defaultValue(), configured.getHighPriorityPoolRatio(), 0.0);
		assertEquals(RocksDBOptions.WRITE_BUFFER_RATIO.defaultValue(), configured.getWriteBufferRatio(), 0.0);
	}

	@Test
	public void testConfigureManagedMemory() {
		final Configuration config = new Configuration();
		config.setBoolean(RocksDBOptions.USE_MANAGED_MEMORY, true);

		final RocksDBMemoryConfiguration memSettings = RocksDBMemoryConfiguration.fromOtherAndConfiguration(
			new RocksDBMemoryConfiguration(), config);

		assertTrue(memSettings.isUsingManagedMemory());
	}

	@Test
	public void testConfigureIllegalMemoryControlParameters() {
		RocksDBMemoryConfiguration memSettings = new RocksDBMemoryConfiguration();

		verifySetParameter(() -> memSettings.setFixedMemoryPerSlot("-1B"));
		verifySetParameter(() -> memSettings.setHighPriorityPoolRatio(-0.1));
		verifySetParameter(() -> memSettings.setHighPriorityPoolRatio(1.1));
		verifySetParameter(() -> memSettings.setWriteBufferRatio(-0.1));
		verifySetParameter(() -> memSettings.setWriteBufferRatio(1.1));

		memSettings.setFixedMemoryPerSlot("128MB");
		memSettings.setWriteBufferRatio(0.6);
		memSettings.setHighPriorityPoolRatio(0.6);

		try {
			// sum of writeBufferRatio and highPriPoolRatio larger than 1.0
			memSettings.validate();
			fail("Expected an IllegalArgumentException.");
		} catch (IllegalArgumentException expected) {
			// expected exception
		}
	}

	private void verifySetParameter(Runnable setter) {
		try {
			setter.run();
			fail("No expected IllegalArgumentException.");
		} catch (IllegalArgumentException expected) {
			// expected exception
		}
	}

	// ------------------------------------------------------------------------
	//  Contained Non-partitioned State Backend
	// ------------------------------------------------------------------------

	@Test
	public void testCallsForwardedToNonPartitionedBackend() throws Exception {
		StateBackend storageBackend = new MemoryStateBackend();
		RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(storageBackend);
		assertEquals(storageBackend, rocksDbBackend.getCheckpointBackend());
	}

	// ------------------------------------------------------------------------
	//  Utilities
	// ------------------------------------------------------------------------

	static MockEnvironment getMockEnvironment(File... tempDirs) {
		final String[] tempDirStrings = new String[tempDirs.length];
		for (int i = 0; i < tempDirs.length; i++) {
			tempDirStrings[i] = tempDirs[i].getAbsolutePath();
		}

		IOManager ioMan = mock(IOManager.class);
		when(ioMan.getSpillingDirectories()).thenReturn(tempDirs);

		return MockEnvironment.builder()
			.setUserCodeClassLoader(RocksDBStateBackendConfigTest.class.getClassLoader())
			.setTaskManagerRuntimeInfo(new TestingTaskManagerRuntimeInfo(new Configuration(), tempDirStrings))
			.setIOManager(ioMan).build();
	}

	private void verifyIllegalArgument(
			ConfigOption<?> configOption,
			String configValue) {
		Configuration configuration = new Configuration();
		configuration.setString(configOption.key(), configValue);

		DefaultConfigurableOptionsFactory optionsFactory = new DefaultConfigurableOptionsFactory();
		try {
			optionsFactory.configure(configuration);
			fail("Not throwing expected IllegalArgumentException.");
		} catch (IllegalArgumentException e) {
			// ignored
		}
	}

	/**
	 * An implementation of options factory for testing.
	 */
	public static class TestOptionsFactory implements ConfigurableRocksDBOptionsFactory {
		public static final ConfigOption<Integer> BACKGROUND_JOBS_OPTION =
			ConfigOptions.key("my.custom.rocksdb.backgroundJobs")
				.intType()
				.defaultValue(2);

		private int backgroundJobs = BACKGROUND_JOBS_OPTION.defaultValue();

		@Override
		public DBOptions createDBOptions(DBOptions currentOptions, Collection<AutoCloseable> handlesToClose) {
			return currentOptions.setMaxBackgroundJobs(backgroundJobs);
		}

		@Override
		public ColumnFamilyOptions createColumnOptions(ColumnFamilyOptions currentOptions, Collection<AutoCloseable> handlesToClose) {
			return currentOptions.setCompactionStyle(CompactionStyle.UNIVERSAL);
		}

		@Override
		public RocksDBOptionsFactory configure(ReadableConfig configuration) {
			this.backgroundJobs = configuration.get(BACKGROUND_JOBS_OPTION);
			return this;
		}
	}
}