From 0613ff892f8701e3772a5d3d75558b4ea342d722 Mon Sep 17 00:00:00 2001
From: Daniel Smith <daniel_b_smith@apple.com>
Date: Fri, 9 Jul 2021 16:25:14 -0400
Subject: [PATCH] Enable RocksDB in simulation

---
 fdbclient/ServerKnobs.cpp                |  4 +++-
 fdbserver/KeyValueStoreRocksDB.actor.cpp | 14 +++++++++++---
 fdbserver/SimulatedCluster.actor.cpp     | 13 +++++++++++--
 3 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/fdbclient/ServerKnobs.cpp b/fdbclient/ServerKnobs.cpp
index 994f367292..d46336099d 100644
--- a/fdbclient/ServerKnobs.cpp
+++ b/fdbclient/ServerKnobs.cpp
@@ -335,7 +335,9 @@ void ServerKnobs::initialize(Randomize randomize, ClientKnobs* clientKnobs, IsSi
 	// KeyValueStoreRocksDB
 	init( ROCKSDB_BACKGROUND_PARALLELISM,                          0 );
 	init( ROCKSDB_READ_PARALLELISM,                                4 );
-	init( ROCKSDB_MEMTABLE_BYTES,                  512 * 1024 * 1024 );
+	// Use a smaller memtable in simulation to avoid OOMs.
+	int64_t memtableBytes = isSimulated ? 64 * 1014 : 512 * 1024 * 1024;
+	init( ROCKSDB_MEMTABLE_BYTES,                      memtableBytes );
 	init( ROCKSDB_UNSAFE_AUTO_FSYNC,                           false );
 	init( ROCKSDB_PERIODIC_COMPACTION_SECONDS,                     0 );
 	init( ROCKSDB_PREFIX_LEN,                                      0 );
diff --git a/fdbserver/KeyValueStoreRocksDB.actor.cpp b/fdbserver/KeyValueStoreRocksDB.actor.cpp
index 99278a7862..edae3412d0 100644
--- a/fdbserver/KeyValueStoreRocksDB.actor.cpp
+++ b/fdbserver/KeyValueStoreRocksDB.actor.cpp
@@ -7,6 +7,7 @@
 #include <rocksdb/slice_transform.h>
 #include <rocksdb/table.h>
 #include <rocksdb/utilities/table_properties_collectors.h>
+#include "fdbserver/CoroFlow.h"
 #include "flow/flow.h"
 #include "flow/IThreadPool.h"
 
@@ -283,7 +284,9 @@ struct RocksDBKeyValueStore : IKeyValueStore {
 				                              std::min(value.size(), size_t(a.maxLength)))));
 			} else {
 				if (!s.IsNotFound()) {
-					TraceEvent(SevError, "RocksDBError").detail("Error", s.ToString()).detail("Method", "ReadValuePrefix");
+					TraceEvent(SevError, "RocksDBError")
+					    .detail("Error", s.ToString())
+					    .detail("Method", "ReadValuePrefix");
 				}
 				a.result.send(Optional<Value>());
 			}
@@ -367,8 +370,13 @@ struct RocksDBKeyValueStore : IKeyValueStore {
 	std::unique_ptr<rocksdb::WriteBatch> writeBatch;
 
 	explicit RocksDBKeyValueStore(const std::string& path, UID id) : path(path), id(id) {
-		writeThread = createGenericThreadPool();
-		readThreads = createGenericThreadPool();
+		if (g_network->isSimulated()) {
+			writeThread = CoroThreadPool::createThreadPool();
+			readThreads = CoroThreadPool::createThreadPool();
+		} else {
+			writeThread = createGenericThreadPool();
+			readThreads = createGenericThreadPool();
+		}
 		writeThread->addThread(new Writer(db, id), "fdb-rocksdb-wr");
 		for (unsigned i = 0; i < SERVER_KNOBS->ROCKSDB_READ_PARALLELISM; ++i) {
 			readThreads->addThread(new Reader(db), "fdb-rocksdb-re");
diff --git a/fdbserver/SimulatedCluster.actor.cpp b/fdbserver/SimulatedCluster.actor.cpp
index a8a152f141..831fbaf630 100644
--- a/fdbserver/SimulatedCluster.actor.cpp
+++ b/fdbserver/SimulatedCluster.actor.cpp
@@ -232,6 +232,7 @@ public:
 	//	1 = "memory"
 	//	2 = "memory-radixtree-beta"
 	//	3 = "ssd-redwood-experimental"
+	//	4 = "ssd-rocksdb-experimental"
 	// Requires a comma-separated list of numbers WITHOUT whitespaces
 	std::vector<int> storageEngineExcludeTypes;
 	// Set the maximum TLog version that can be selected for a test
@@ -1252,7 +1253,7 @@ void SimulationConfig::setDatacenters(const TestConfig& testConfig) {
 
 // Sets storage engine based on testConfig details
 void SimulationConfig::setStorageEngine(const TestConfig& testConfig) {
-	int storage_engine_type = deterministicRandom()->randomInt(0, 4);
+	int storage_engine_type = deterministicRandom()->randomInt(0, 5);
 	if (testConfig.storageEngineType.present()) {
 		storage_engine_type = testConfig.storageEngineType.get();
 	} else {
@@ -1260,7 +1261,7 @@ void SimulationConfig::setStorageEngine(const TestConfig& testConfig) {
 		while (std::find(testConfig.storageEngineExcludeTypes.begin(),
 		                 testConfig.storageEngineExcludeTypes.end(),
 		                 storage_engine_type) != testConfig.storageEngineExcludeTypes.end()) {
-			storage_engine_type = deterministicRandom()->randomInt(0, 4);
+			storage_engine_type = deterministicRandom()->randomInt(0, 5);
 		}
 	}
 
@@ -1285,6 +1286,14 @@ void SimulationConfig::setStorageEngine(const TestConfig& testConfig) {
 		set_config("ssd-redwood-experimental");
 		break;
 	}
+	case 4: {
+		TEST(true); // Simulated cluster using RocksDB storage engine
+		set_config("ssd-rocksdb-experimental");
+		// Tests using the RocksDB engine are necessarily non-deterministic because of RocksDB
+		// background threads.
+		noUnseed = true;
+		break;
+	}
 	default:
 		ASSERT(false); // Programmer forgot to adjust cases.
 	}