From d334b6484edafb1e66f2c9353c95107dccf8b28e Mon Sep 17 00:00:00 2001
From: Jon Fu <jon.fu@snowflake.com>
Date: Wed, 2 Sep 2020 15:17:54 -0400
Subject: [PATCH] attempt to write to system keys with snapshot

---
 fdbclient/SystemData.cpp             |  2 ++
 fdbclient/SystemData.h               |  1 +
 fdbserver/DataDistribution.actor.cpp | 16 ++++++++++++++++
 3 files changed, 19 insertions(+)

diff --git a/fdbclient/SystemData.cpp b/fdbclient/SystemData.cpp
index cb3b247b4c..0650f580f9 100644
--- a/fdbclient/SystemData.cpp
+++ b/fdbclient/SystemData.cpp
@@ -1062,4 +1062,6 @@ const KeyRangeRef testOnlyTxnStateStorePrefixRange(
 );
 
 const KeyRef writeRecoveryKey = LiteralStringRef("\xff/writeRecovery");
+const ValueRef writeRecoveryKeyTrue = LiteralStringRef("1");
+const ValueRef writeRecoveryKeyFalse = LiteralStringRef("0");
 const KeyRef snapshotEndVersionKey = LiteralStringRef("\xff/snapshotEndVersion");
diff --git a/fdbclient/SystemData.h b/fdbclient/SystemData.h
index 4db3225fb5..4006282708 100644
--- a/fdbclient/SystemData.h
+++ b/fdbclient/SystemData.h
@@ -398,6 +398,7 @@ extern const KeyRangeRef testOnlyTxnStateStorePrefixRange;
 
 // Snapshot + Incremental Restore
 extern const KeyRef writeRecoveryKey;
+extern const ValueRef writeRecoveryKeyTrue, writeRecoveryKeyFalse;
 extern const KeyRef snapshotEndVersionKey;
 
 #pragma clang diagnostic pop
diff --git a/fdbserver/DataDistribution.actor.cpp b/fdbserver/DataDistribution.actor.cpp
index 094e2111de..94a49b8889 100644
--- a/fdbserver/DataDistribution.actor.cpp
+++ b/fdbserver/DataDistribution.actor.cpp
@@ -4728,6 +4728,22 @@ static std::set<int> const& normalDataDistributorErrors() {
 
 ACTOR Future<Void> ddSnapCreateCore(DistributorSnapRequest snapReq, Reference<AsyncVar<struct ServerDBInfo>> db ) {
 	state Database cx = openDBOnServer(db, TaskPriority::DefaultDelay, true, true);
+	state Reference<ReadYourWritesTransaction> tr(new ReadYourWritesTransaction(cx));
+	loop {
+		try {
+			tr->setOption(FDBTransactionOptions::ACCESS_SYSTEM_KEYS);
+			tr->setOption(FDBTransactionOptions::LOCK_AWARE);
+			TraceEvent("SnapDataDistributor_WriteFlagAttempt")
+				.detail("SnapPayload", snapReq.snapPayload)
+				.detail("SnapUID", snapReq.snapUID);
+			tr->set(writeRecoveryKey, writeRecoveryKeyTrue);
+			wait(tr->commit());
+			break;
+		} catch (Error& e) {
+			TraceEvent("SnapDataDistributor_WriteFlagError").error(e);
+			wait(tr->onError(e));
+		}
+	}
 	TraceEvent("SnapDataDistributor_SnapReqEnter")
 		.detail("SnapPayload", snapReq.snapPayload)
 		.detail("SnapUID", snapReq.snapUID);