Bulk Loading Part 1: Transaction-Based User Interface

fdbclient/NativeAPI.actor.cpp

@@ -1608,6 +1608,28 @@ DatabaseContext::DatabaseContext(Reference<AsyncVar<Reference<IClusterConnection
 		    std::make_unique<FaultToleranceMetricsImpl>(
 		        singleKeyRange("fault_tolerance_metrics_json"_sr)
 		            .withPrefix(SpecialKeySpace::getModuleRange(SpecialKeySpace::MODULE::METRICS).begin)));
+		registerSpecialKeysImpl(
+		    SpecialKeySpace::MODULE::BULKLOADING,
+		    SpecialKeySpace::IMPLTYPE::READONLY,
+		    std::make_unique<BulkLoadStatusImpl>(
+		        KeyRangeRef("status/"_sr, "status0"_sr)
+		            .withPrefix(SpecialKeySpace::getModuleRange(SpecialKeySpace::MODULE::BULKLOADING).begin)));
+		registerSpecialKeysImpl(
+		    SpecialKeySpace::MODULE::BULKLOADING,
+		    SpecialKeySpace::IMPLTYPE::READWRITE,
+		    std::make_unique<BulkLoadTaskImpl>(
+		        KeyRangeRef("task/"_sr, "task0"_sr)
+		            .withPrefix(SpecialKeySpace::getModuleRange(SpecialKeySpace::MODULE::BULKLOADING).begin)));
+		registerSpecialKeysImpl(
+		    SpecialKeySpace::MODULE::BULKLOADING,
+		    SpecialKeySpace::IMPLTYPE::READWRITE,
+		    std::make_unique<BulkLoadCancelImpl>(
+		        KeyRangeRef("cancel/"_sr, "cancel0"_sr)
+		            .withPrefix(SpecialKeySpace::getModuleRange(SpecialKeySpace::MODULE::BULKLOADING).begin)));
+		registerSpecialKeysImpl(SpecialKeySpace::MODULE::BULKLOADING,
+		                        SpecialKeySpace::IMPLTYPE::READWRITE,
+		                        std::make_unique<BulkLoadModeImpl>(singleKeyRange("mode"_sr).withPrefix(
+		                            SpecialKeySpace::getModuleRange(SpecialKeySpace::MODULE::BULKLOADING).begin)));
 	}
 
 	if (apiVersion.version() >= 700) {

fdbclient/SystemData.cpp

@@ -1215,6 +1215,21 @@ const KeyRef moveKeysLockWriteKey = "\xff/moveKeysLock/Write"_sr;
 const KeyRef dataDistributionModeKey = "\xff/dataDistributionMode"_sr;
 const UID dataDistributionModeLock = UID(6345, 3425);
 
+// Bulk loading keys
+const KeyRef bulkLoadModeKey = "\xff/bulkLoadMode"_sr;
+const KeyRef bulkLoadPrefix = "\xff/bulkLoad/"_sr;
+
+const Value bulkLoadStateValue(const BulkLoadState& bulkLoadState) {
+	return ObjectWriter::toValue(bulkLoadState, IncludeVersion());
+}
+
+BulkLoadState decodeBulkLoadState(const ValueRef& value) {
+	BulkLoadState bulkLoadState;
+	ObjectReader reader(value.begin(), IncludeVersion());
+	reader.deserialize(bulkLoadState);
+	return bulkLoadState;
+}
+
 // Keys to view and control tag throttling
 const KeyRangeRef tagThrottleKeys = KeyRangeRef("\xff\x02/throttledTags/tag/"_sr, "\xff\x02/throttledTags/tag0"_sr);
 const KeyRef tagThrottleKeysPrefix = tagThrottleKeys.begin;

fdbclient/include/fdbclient/BulkLoading.h

@@ -0,0 +1,69 @@
+/*
+ * BulkLoading.h
+ *
+ * This source file is part of the FoundationDB open source project
+ *
+ * Copyright 2013-2024 Apple Inc. and the FoundationDB project authors
+ *
+ * Licensed 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.
+ */
+
+#ifndef FDBCLIENT_BULKLOADING_H
+#define FDBCLIENT_BULKLOADING_H
+#pragma once
+
+#include "fdbclient/FDBTypes.h"
+#include "fdbrpc/fdbrpc.h"
+
+enum class BulkLoadPhase : uint8_t {
+	Invalid = 0,
+	Running = 1,
+	Complete = 2,
+	Error = 3,
+};
+
+enum class BulkLoadType : uint8_t {
+	Invalid = 0,
+	SQLite = 1,
+	RocksDB = 2,
+	ShardedRocksDB = 3,
+};
+
+struct BulkLoadState {
+	constexpr static FileIdentifier file_identifier = 1384499;
+
+	BulkLoadState() = default;
+
+	BulkLoadState(BulkLoadType loadType, std::set<std::string> filePaths) : loadType(loadType), filePaths(filePaths) {}
+
+	BulkLoadState(KeyRange range, BulkLoadType loadType, std::set<std::string> filePaths)
+	  : range(range), loadType(loadType), filePaths(filePaths) {}
+
+	bool isValid() const { return loadType != BulkLoadType::Invalid; }
+
+	std::string toString() const {
+		return "BulkLoadState: [Range]: " + Traceable<KeyRangeRef>::toString(range) +
+		       ", [Type]: " + std::to_string(static_cast<uint8_t>(loadType)) + ", [FilePath]: " + describe(filePaths);
+	}
+
+	template <class Ar>
+	void serialize(Ar& ar) {
+		serializer(ar, range, loadType, filePaths);
+	}
+
+	KeyRange range;
+	BulkLoadType loadType;
+	std::set<std::string> filePaths;
+};
+
+#endif

fdbclient/include/fdbclient/SpecialKeySpace.actor.h

@@ -182,6 +182,7 @@ class SpecialKeySpace {
 		STATUSJSON,
 		UNKNOWN, // default value for all unregistered range
 		WORKERINTERFACE,
+		BULKLOADING,
 	};
 
 	enum class IMPLTYPE {
@@ -562,6 +563,50 @@ class DataDistributionImpl : public SpecialKeyRangeRWImpl {
 	Future<Optional<std::string>> commit(ReadYourWritesTransaction* ryw) override;
 };
 
+class BulkLoadStatusImpl : public SpecialKeyRangeReadImpl {
+public:
+	explicit BulkLoadStatusImpl(KeyRangeRef kr);
+	Future<RangeResult> getRange(ReadYourWritesTransaction* ryw,
+	                             KeyRangeRef kr,
+	                             GetRangeLimits limitsHint) const override;
+};
+
+class BulkLoadTaskImpl : public SpecialKeyRangeRWImpl {
+public:
+	explicit BulkLoadTaskImpl(KeyRangeRef kr);
+	Future<RangeResult> getRange(ReadYourWritesTransaction* ryw,
+	                             KeyRangeRef kr,
+	                             GetRangeLimits limitsHint) const override;
+	void set(ReadYourWritesTransaction* ryw, const KeyRef& key, const ValueRef& value) override;
+	Future<Optional<std::string>> commit(ReadYourWritesTransaction* ryw) override;
+	void clear(ReadYourWritesTransaction* ryw, const KeyRangeRef& range) override;
+	void clear(ReadYourWritesTransaction* ryw, const KeyRef& key) override;
+};
+
+class BulkLoadCancelImpl : public SpecialKeyRangeRWImpl {
+public:
+	explicit BulkLoadCancelImpl(KeyRangeRef kr);
+	Future<RangeResult> getRange(ReadYourWritesTransaction* ryw,
+	                             KeyRangeRef kr,
+	                             GetRangeLimits limitsHint) const override;
+	void set(ReadYourWritesTransaction* ryw, const KeyRef& key, const ValueRef& value) override;
+	Future<Optional<std::string>> commit(ReadYourWritesTransaction* ryw) override;
+	void clear(ReadYourWritesTransaction* ryw, const KeyRangeRef& range) override;
+	void clear(ReadYourWritesTransaction* ryw, const KeyRef& key) override;
+};
+
+class BulkLoadModeImpl : public SpecialKeyRangeRWImpl {
+public:
+	explicit BulkLoadModeImpl(KeyRangeRef kr);
+	Future<RangeResult> getRange(ReadYourWritesTransaction* ryw,
+	                             KeyRangeRef kr,
+	                             GetRangeLimits limitsHint) const override;
+	void set(ReadYourWritesTransaction* ryw, const KeyRef& key, const ValueRef& value) override;
+	Future<Optional<std::string>> commit(ReadYourWritesTransaction* ryw) override;
+	void clear(ReadYourWritesTransaction* ryw, const KeyRangeRef& range) override;
+	void clear(ReadYourWritesTransaction* ryw, const KeyRef& key) override;
+};
+
 class WorkerInterfacesSpecialKeyImpl : public SpecialKeyRangeReadImpl {
 public:
 	explicit WorkerInterfacesSpecialKeyImpl(KeyRangeRef kr);

fdbclient/include/fdbclient/SystemData.h

@@ -25,6 +25,7 @@
 // Functions and constants documenting the organization of the reserved keyspace in the database beginning with "\xFF"
 
 #include "fdbclient/AccumulativeChecksum.h"
+#include "fdbclient/BulkLoading.h"
 #include "fdbclient/FDBTypes.h"
 #include "fdbclient/BlobWorkerInterface.h" // TODO move the functions that depend on this out of here and into BlobWorkerInterface.h to remove this dependency
 #include "fdbclient/StorageServerInterface.h"
@@ -518,6 +519,11 @@ extern const KeyRef moveKeysLockOwnerKey, moveKeysLockWriteKey;
 extern const KeyRef dataDistributionModeKey;
 extern const UID dataDistributionModeLock;
 
+extern const KeyRef bulkLoadModeKey;
+extern const KeyRef bulkLoadPrefix;
+const Value bulkLoadStateValue(const BulkLoadState& bulkLoadState);
+BulkLoadState decodeBulkLoadState(const ValueRef& value);
+
 // Keys to view and control tag throttling
 extern const KeyRangeRef tagThrottleKeys;
 extern const KeyRef tagThrottleKeysPrefix;

fdbserver/DataDistribution.actor.cpp

@@ -994,6 +994,26 @@ ACTOR Future<Void> serveBlobMigratorRequests(Reference<DataDistributor> self,
 	}
 }
 
+ACTOR Future<Void> bulkLoadingCore(Reference<DataDistributor> self) {
+	loop {
+		try {
+			Database cx = self->txnProcessor->context();
+			Transaction tr(cx);
+			Optional<Value> mode = wait(tr.get(bulkLoadModeKey));
+			if (mode.present()) {
+				int bulkLoadMode = BinaryReader::fromStringRef<int>(mode.get(), Unversioned());
+				if (bulkLoadMode == 1) {
+					TraceEvent("BulkLoadingCore").detail("Status", "Mode On").log();
+					return Void();
+				}
+			}
+		} catch (Error& e) {
+			TraceEvent("BulkLoadingCore").detail("Status", "Error").errorUnsuppressed(e).log();
+			wait(delay(5.0));
+		}
+	}
+}
+
 // Runs the data distribution algorithm for FDB, including the DD Queue, DD tracker, and DD team collection
 ACTOR Future<Void> dataDistribution(Reference<DataDistributor> self,
                                     PromiseStream<GetMetricsListRequest> getShardMetricsList,
@@ -1072,6 +1092,8 @@ ACTOR Future<Void> dataDistribution(Reference<DataDistributor> self,
 
 			actors.push_back(self->pollMoveKeysLock());
 
+			actors.push_back(bulkLoadingCore(self));
+
 			self->context->tracker = makeReference<DataDistributionTracker>(
 			    DataDistributionTrackerInitParams{ .db = self->txnProcessor,
 			                                       .distributorId = self->ddId,