Support multiple range

Author: yao-xiao-github

fdbclient/ServerKnobs.cpp

@@ -623,6 +623,7 @@ void ServerKnobs::initialize(Randomize randomize, ClientKnobs* clientKnobs, IsSi
 	// Block cache key-value checksum. Checksum is validated during read, so has non-trivial impact on read performance.
 	init( ROCKSDB_BLOCK_PROTECTION_BYTES_PER_KEY,                  0 ); if ( randomize && BUGGIFY ) ROCKSDB_BLOCK_PROTECTION_BYTES_PER_KEY = 8; // Default: 0 (disabled). Supported values: 0, 1, 2, 4, 8.
 	init( ROCKSDB_ENABLE_NONDETERMINISM,                      false );
+	init( SHARDED_ROCKSDB_ALLOW_MULTIPLE_RANGES,              false );
 	init( SHARDED_ROCKSDB_ALLOW_WRITE_STALL_ON_FLUSH,          false );	
 	init( SHARDED_ROCKSDB_VALIDATE_MAPPING_RATIO,               0.01 ); if (isSimulated) SHARDED_ROCKSDB_VALIDATE_MAPPING_RATIO = deterministicRandom()->random01();
 	init( SHARD_METADATA_SCAN_BYTES_LIMIT,                  10485760 ); // 10MB

fdbclient/include/fdbclient/ServerKnobs.h

@@ -607,6 +607,7 @@ class SWIFT_CXX_IMMORTAL_SINGLETON_TYPE ServerKnobs : public KnobsImpl<ServerKno
 	                                    // Note that turning this on in simulation could lead to non-deterministic runs
 	                                    // since we rely on rocksdb metadata. This knob also applies to sharded rocks
 	                                    // storage engine.
+	bool SHARDED_ROCKSDB_ALLOW_MULTIPLE_RANGES;
 	bool SHARDED_ROCKSDB_ALLOW_WRITE_STALL_ON_FLUSH;
 	int SHARDED_ROCKSDB_MEMTABLE_MAX_RANGE_DELETIONS;
 	double SHARDED_ROCKSDB_VALIDATE_MAPPING_RATIO;

fdbserver/KeyValueStoreShardedRocksDB.actor.cpp

@@ -1546,7 +1546,7 @@ class ShardManager {
 		for (auto it = ranges.begin(); it != ranges.end(); ++it) {
 			if (it.value()) {
 				if (it.value()->physicalShard->id == id) {
-					TraceEvent(SevError, "ShardedRocksDBAddRange")
+					TraceEvent(SevWarn, "ShardedRocksDBAddRange")
 					    .detail("ErrorType", "RangeAlreadyExist")
 					    .detail("IntersectingRange", it->range())
 					    .detail("DataShardRange", it->value()->range)
@@ -1564,15 +1564,21 @@ class ShardManager {
 			}
 		}
 
-		auto currentCfOptions = active ? rState->getCFOptions() : rState->getCFOptionsForInactiveShard();
-		auto [it, inserted] = physicalShards.emplace(id, std::make_shared<PhysicalShard>(db, id, currentCfOptions));
-		std::shared_ptr<PhysicalShard>& shard = it->second;
-
-		activePhysicalShardIds.emplace(id);
+		auto it = physicalShards.find(id);
+		std::shared_ptr<PhysicalShard> physicalShard = nullptr;
+		if (it != physicalShards.end()) {
+			physicalShard = it->second;
+			ASSERT(SERVER_KNOBS->SHARDED_ROCKSDB_ALLOW_MULTIPLE_RANGES);
+		} else {
+			auto currentCfOptions = active ? rState->getCFOptions() : rState->getCFOptionsForInactiveShard();
+			auto [it, inserted] = physicalShards.emplace(id, std::make_shared<PhysicalShard>(db, id, currentCfOptions));
+			physicalShard = it->second;
+			activePhysicalShardIds.emplace(id);
+		}
 
-		auto dataShard = std::make_unique<DataShard>(range, shard.get());
+		auto dataShard = std::make_unique<DataShard>(range, physicalShard.get());
 		dataShardMap.insert(range, dataShard.get());
-		shard->dataShards[range.begin.toString()] = std::move(dataShard);
+		physicalShard->dataShards[range.begin.toString()] = std::move(dataShard);
 
 		validate();
 
@@ -1581,7 +1587,7 @@ class ShardManager {
 		    .detail("ShardId", id)
 		    .detail("Active", active);
 
-		return shard.get();
+		return physicalShard.get();
 	}
 
 	std::vector<std::string> removeRange(KeyRange range) {
@@ -1636,6 +1642,7 @@ class ShardManager {
 
 			// Range modification could result in more than one segments. Remove the original segment key here.
 			existingShard->dataShards.erase(shardRange.begin.toString());
+			int count = 0;
 			if (shardRange.begin < range.begin) {
 				auto dataShard =
 				    std::make_unique<DataShard>(KeyRange(KeyRangeRef(shardRange.begin, range.begin)), existingShard);
@@ -1646,6 +1653,7 @@ class ShardManager {
 
 				existingShard->dataShards[shardRange.begin.toString()] = std::move(dataShard);
 				logShardEvent(existingShard->id, shardRange, ShardOp::MODIFY_RANGE, SevInfo, msg);
+				count++;
 			}
 
 			if (shardRange.end > range.end) {
@@ -1658,6 +1666,12 @@ class ShardManager {
 
 				existingShard->dataShards[range.end.toString()] = std::move(dataShard);
 				logShardEvent(existingShard->id, shardRange, ShardOp::MODIFY_RANGE, SevInfo, msg);
+				count++;
+			}
+
+			if (count > 1) {
+				TraceEvent("RangeSplit").detail("OriginalRange", shardRange).detail("RemovedRange", range);
+				ASSERT(SERVER_KNOBS->SHARDED_ROCKSDB_ALLOW_MULTIPLE_RANGES);
 			}
 		}
 
@@ -1986,28 +2000,37 @@ class ShardManager {
 		}
 
 		TraceEvent(SevVerbose, "ShardedRocksValidateShardManager", this->logId);
+		int totalDataShards = 0;
+		int expectedDataShards = 0;
 		for (auto s = dataShardMap.ranges().begin(); s != dataShardMap.ranges().end(); ++s) {
 			TraceEvent e(SevVerbose, "ShardedRocksValidateDataShardMap", this->logId);
 			e.detail("Range", s->range());
 			const DataShard* shard = s->value();
 			e.detail("ShardAddress", reinterpret_cast<std::uintptr_t>(shard));
-			if (shard != nullptr) {
-				e.detail("PhysicalShard", shard->physicalShard->id);
-			} else {
-				e.detail("Shard", "Empty");
+			if (shard == nullptr) {
+				e.detail("RangeUnassigned", "True");
+				continue;
 			}
-			if (shard != nullptr) {
-				if (shard->range != static_cast<KeyRangeRef>(s->range())) {
-					TraceEvent(SevWarnAlways, "ShardRangeMismatch").detail("Range", s->range());
-				}
-
-				ASSERT(shard->range == static_cast<KeyRangeRef>(s->range()));
-				ASSERT(shard->physicalShard != nullptr);
-				auto it = shard->physicalShard->dataShards.find(shard->range.begin.toString());
-				ASSERT(it != shard->physicalShard->dataShards.end());
-				ASSERT(it->second.get() == shard);
+			totalDataShards++;
+			if (shard->range != static_cast<KeyRangeRef>(s->range())) {
+				TraceEvent(SevWarnAlways, "ShardRangeMismatch")
+				    .detail("Range", s->range())
+				    .detail("DataShardRange", shard->range)
+				    .detail("PhysicalShardId", shard->physicalShard->id);
 			}
+
+			ASSERT(shard->range == static_cast<KeyRangeRef>(s->range()));
+			ASSERT(shard->physicalShard != nullptr);
+			auto it = shard->physicalShard->dataShards.find(shard->range.begin.toString());
+			ASSERT(it != shard->physicalShard->dataShards.end());
+			ASSERT(it->second.get() == shard);
+		}
+
+		for (auto [shardId, physicalShard] : physicalShards) {
+			ASSERT(physicalShard);
+			expectedDataShards += physicalShard->dataShards.size();
 		}
+		ASSERT_EQ(expectedDataShards, totalDataShards);
 	}
 
 private:
@@ -4403,6 +4426,81 @@ TEST_CASE("noSim/ShardedRocksDB/Metadata") {
 	return Void();
 }
 
+TEST_CASE("noSim/ShardedRocksDBRangeOps/RemoveSplitRange") {
+	state std::string rocksDBTestDir = "sharded-rocksdb-kvs-test-db";
+	platform::eraseDirectoryRecursive(rocksDBTestDir);
+
+	state ShardedRocksDBKeyValueStore* rocksdbStore =
+	    new ShardedRocksDBKeyValueStore(rocksDBTestDir, deterministicRandom()->randomUniqueID());
+	state IKeyValueStore* kvStore = rocksdbStore;
+	wait(kvStore->init());
+
+	// Add two ranges to the same shard.
+	{
+		std::vector<Future<Void>> addRangeFutures;
+		addRangeFutures.push_back(kvStore->addRange(KeyRangeRef("a"_sr, "d"_sr), "shard-1"));
+		addRangeFutures.push_back(kvStore->addRange(KeyRangeRef("g"_sr, "n"_sr), "shard-1"));
+
+		wait(waitForAll(addRangeFutures));
+	}
+
+	state std::set<std::string> originalKeys = { "a", "b", "c", "g", "h", "m" };
+	state std::set<std::string> currentKeys = originalKeys;
+	for (auto key : originalKeys) {
+		kvStore->set(KeyValueRef(key, key));
+	}
+	wait(kvStore->commit());
+
+	state std::string key;
+	for (key : currentKeys) {
+		Optional<Value> val = wait(kvStore->readValue(key));
+		ASSERT(val.present());
+		ASSERT(val.get().toString() == key);
+	}
+
+	// Remove single range.
+	std::vector<std::string> shardsToCleanUp;
+	auto shardIds = kvStore->removeRange(KeyRangeRef("b"_sr, "c"_sr));
+	// Remove range didn't create empty shard.
+	ASSERT_EQ(shardIds.size(), 0);
+
+	currentKeys.erase("b");
+	for (key : originalKeys) {
+		Optional<Value> val = wait(kvStore->readValue(key));
+		if (currentKeys.contains(key)) {
+			ASSERT(val.present());
+			ASSERT(val.get().toString() == key);
+		} else {
+			ASSERT(!val.present());
+		}
+	}
+
+	// Remove range spanning on multple sub range.
+	auto shardIds = kvStore->removeRange(KeyRangeRef("c"_sr, "k"_sr));
+	ASSERT(shardIds.empty());
+
+	currentKeys.erase("c");
+	currentKeys.erase("g");
+	currentKeys.erase("h");
+	for (key : originalKeys) {
+		Optional<Value> val = wait(kvStore->readValue(key));
+		if (currentKeys.contains(key)) {
+			ASSERT(val.present());
+			ASSERT(val.get().toString() == key);
+		} else {
+			ASSERT(!val.present());
+		}
+	}
+
+	{
+		Future<Void> closed = kvStore->onClosed();
+		kvStore->dispose();
+		wait(closed);
+	}
+	ASSERT(!directoryExists(rocksDBTestDir));
+	return Void();
+}
+
 TEST_CASE("noSim/ShardedRocksDBCheckpoint/CheckpointBasic") {
 	state std::string rocksDBTestDir = "sharded-rocks-checkpoint-restore";
 	state std::map<Key, Value> kvs({ { "a"_sr, "TestValueA"_sr },

tests/noSim/ShardedRocksDBTest.toml

@@ -2,6 +2,8 @@
 rocksdb_disable_auto_compactions = true
 rocksdb_suggest_compact_clear_range = false
 rocksdb_empty_range_check = false
+sharded_rocksdb_validate_mapping_ratio=1.0
+sharded_rocksdb_allow_multiple_ranges=true
 
 [[test]]
  testTitle = 'UnitTests'
@@ -11,4 +13,5 @@ rocksdb_empty_range_check = false
      [[test.workload]]
      testName = 'UnitTests'
      maxTestCases = 10
-     testsMatching = 'noSim/ShardedRocksDB/'
+     # testsMatching = 'noSim/ShardedRocksDB/'
+     testsMatching = 'noSim/ShardedRocksDBRangeOps/RemoveSplitRange'