From de4753a5db3daa09172ebfb93b27eb5c793279f5 Mon Sep 17 00:00:00 2001
From: Chaoguang Lin <chaoguang.lin@snowflake.com>
Date: Fri, 23 Apr 2021 01:32:30 -0700
Subject: [PATCH] Add a workaround to temporily use the ryw to create a
 ThreadTransaction; Make sure we are using the same underlying ryw object

---
 fdbcli/ConsistencycheckCommand.actor.cpp |  7 +++---
 fdbcli/fdbcli.actor.cpp                  | 28 ++++++++++++++++++++++--
 fdbcli/fdbcli.h                          |  2 +-
 fdbclient/ThreadSafeTransaction.cpp      |  6 +++++
 fdbclient/ThreadSafeTransaction.h        |  3 +++
 5 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/fdbcli/ConsistencycheckCommand.actor.cpp b/fdbcli/ConsistencycheckCommand.actor.cpp
index 4279727b1d..1d099212d6 100644
--- a/fdbcli/ConsistencycheckCommand.actor.cpp
+++ b/fdbcli/ConsistencycheckCommand.actor.cpp
@@ -10,8 +10,7 @@
 
 using namespace fdb_cli;
 
-ACTOR static Future<bool> consistencycheckCommandActor(Reference<IDatabase> db, std::vector<StringRef> tokens) {
-	state Reference<ITransaction> tr = db->createTransaction();
+ACTOR static Future<bool> consistencycheckCommandActor(Reference<ITransaction> tr, std::vector<StringRef> tokens) {
 	tr->setOption(FDBTransactionOptions::SPECIAL_KEY_SPACE_ENABLE_WRITES);
 	if (tokens.size() == 1) {
 		Optional<Value> suspended = wait(safeThreadFutureToFuture(tr->get(consistencyCheckSpeicalKey)));
@@ -33,8 +32,8 @@ namespace fdb_cli {
 
 const KeyRef consistencyCheckSpeicalKey = LiteralStringRef("\xff\xff/management/consistency_check_suspended");
 
-Future<bool> consistencycheckCommand(Reference<IDatabase> db, std::vector<StringRef> tokens) {
-	return consistencycheckCommandActor(db, tokens);
+Future<bool> consistencycheckCommand(Reference<ITransaction> tr, std::vector<StringRef> tokens) {
+	return consistencycheckCommandActor(tr, tokens);
 }
 
 CommandFactory consistencycheckFactory(
diff --git a/fdbcli/fdbcli.actor.cpp b/fdbcli/fdbcli.actor.cpp
index 062aca4399..8eac79fba1 100644
--- a/fdbcli/fdbcli.actor.cpp
+++ b/fdbcli/fdbcli.actor.cpp
@@ -36,6 +36,7 @@
 #include "fdbclient/FDBOptions.g.h"
 #include "fdbclient/TagThrottle.h"
 
+#include "fdbclient/ThreadSafeTransaction.h"
 #include "flow/DeterministicRandom.h"
 #include "flow/Platform.h"
 
@@ -183,6 +184,13 @@ public:
 		}
 	}
 
+	// TODO: replace the above function after we refactor all fdbcli code
+	void apply(Reference<ITransaction> tr) {
+		for (const auto& [name, value] : transactionOptions.options) {
+			tr->setOption(name, value.castTo<StringRef>());
+		}
+	}
+
 	// Returns true if any options have been set
 	bool hasAnyOptionsEnabled() const { return !transactionOptions.options.empty(); }
 
@@ -2653,6 +2661,20 @@ Reference<ReadYourWritesTransaction> getTransaction(Database db,
 	return tr;
 }
 
+// TODO: Update this function to get rid of Database and ReadYourWritesTransaction after refactoring
+Reference<ITransaction> getTransaction(Database db,
+                                       Reference<ReadYourWritesTransaction>& tr,
+                                       Reference<ITransaction>& tr2,
+                                       FdbOptions* options,
+                                       bool intrans) {
+	if (!tr || !intrans) {
+		tr = makeReference<ReadYourWritesTransaction>(db);
+		options->apply(tr);
+	}
+	tr2 = Reference<ITransaction>(new ThreadSafeTransaction(tr.getPtr()));
+	return tr2;
+}
+
 std::string newCompletion(const char* base, const char* name) {
 	return format("%s%s ", base, name);
 }
@@ -3124,8 +3146,9 @@ ACTOR Future<int> cli(CLIOptions opt, LineNoise* plinenoise) {
 
 	state Database db;
 	state Reference<ReadYourWritesTransaction> tr;
-	// Note: refactoring work, will replace db when we have all commands through the general fdb interface
+	// TODO: refactoring work, will replace db, tr when we have all commands through the general fdb interface
 	state Reference<IDatabase> db2;
+	state Reference<ITransaction> tr2;
 
 	state bool writeMode = false;
 
@@ -3790,7 +3813,8 @@ ACTOR Future<int> cli(CLIOptions opt, LineNoise* plinenoise) {
 				}
 
 				if (tokencmp(tokens[0], "consistencycheck")) {
-					bool _result = wait(consistencycheckCommand(db2, tokens));
+					getTransaction(db, tr, tr2, options, intrans);
+					bool _result = wait(consistencycheckCommand(tr2, tokens));
 					is_error = !_result;
 					continue;
 				}
diff --git a/fdbcli/fdbcli.h b/fdbcli/fdbcli.h
index a5328191d5..007ee6e5f5 100644
--- a/fdbcli/fdbcli.h
+++ b/fdbcli/fdbcli.h
@@ -62,7 +62,7 @@ void printUsage(StringRef command);
 
 // All fdbcli commands (alphabetically)
 // consistency command
-Future<bool> consistencycheckCommand(Reference<IDatabase> db, std::vector<StringRef> tokens);
+Future<bool> consistencycheckCommand(Reference<ITransaction> tr, std::vector<StringRef> tokens);
 
 } // namespace fdb_cli
 
diff --git a/fdbclient/ThreadSafeTransaction.cpp b/fdbclient/ThreadSafeTransaction.cpp
index c5bf2dce87..b40b4b7a2f 100644
--- a/fdbclient/ThreadSafeTransaction.cpp
+++ b/fdbclient/ThreadSafeTransaction.cpp
@@ -152,6 +152,12 @@ ThreadSafeTransaction::ThreadSafeTransaction(DatabaseContext* cx) {
 	    nullptr);
 }
 
+// This constructor is only used while refactoring fdbcli and only called from the main thread
+ThreadSafeTransaction::ThreadSafeTransaction(ReadYourWritesTransaction* ryw) : tr(ryw) {
+	if (tr)
+		tr->addref();
+}
+
 ThreadSafeTransaction::~ThreadSafeTransaction() {
 	ReadYourWritesTransaction* tr = this->tr;
 	if (tr)
diff --git a/fdbclient/ThreadSafeTransaction.h b/fdbclient/ThreadSafeTransaction.h
index e6360c2a6d..5552837a37 100644
--- a/fdbclient/ThreadSafeTransaction.h
+++ b/fdbclient/ThreadSafeTransaction.h
@@ -72,6 +72,9 @@ public:
 	explicit ThreadSafeTransaction(DatabaseContext* cx);
 	~ThreadSafeTransaction() override;
 
+	// Note: used while refactoring fdbcli, need to be removed later
+	explicit ThreadSafeTransaction(ReadYourWritesTransaction* ryw);
+
 	void cancel() override;
 	void setVersion(Version v) override;
 	ThreadFuture<Version> getReadVersion() override;