Addressed comments

Author: akankshamahajan15

fdbserver/datadistributor/DataDistribution.cpp

@@ -313,40 +313,71 @@ Future<Void> remoteRecovered(Reference<AsyncVar<ServerDBInfo> const> db) {
 // On value=2, clears backupPartitionListKey.
 // In both cases the request key is cleared in the same commit.
 Future<Void> monitorBackupPartitionRequired(Database cx, KeyRangeMap<ShardTrackedData>* shards, UID ddId) {
+	// The partition computation can wait arbitrarily long on shard-metrics tracking, so it runs OUTSIDE any
+	// transaction to avoid transaction_too_old. A short re-read in the write transaction protects against
+	// the race where a value=2 (cleanup) arrives while we are computing for a value=1.
 	while (true) {
-		ReadYourWritesTransaction tr(cx);
-		while (true) {
+		// Phase 1: peek the request key in a loop. If nothing pending, park on watch and wait, then re-read.
+		int8_t requestType = 0;
+		while (requestType == 0) {
+			ReadYourWritesTransaction tr(cx);
 			Error err;
 			try {
 				tr.setOption(FDBTransactionOptions::ACCESS_SYSTEM_KEYS);
 				tr.setOption(FDBTransactionOptions::LOCK_AWARE);
 				Optional<Value> value = co_await tr.get(backupPartitionRequiredKey);
-				int8_t requestType = value.present() ? decodeBackupPartitionRequiredValue(value.get()) : 0;
-
-				if (requestType == 1) {
-					std::vector<KeyRange> partitions = co_await calculateBackupPartitionKeyRanges(shards);
-					tr.set(backupPartitionListKey, encodeBackupPartitionListValue(partitions));
-					tr.clear(backupPartitionRequiredKey);
-					co_await tr.commit();
-					TraceEvent("DDBackupPartitionsComputed", ddId).detail("NumPartitions", partitions.size());
-					break;
-				} else if (requestType == 2) {
-					tr.clear(backupPartitionListKey);
-					tr.clear(backupPartitionRequiredKey);
+				requestType = value.present() ? decodeBackupPartitionRequiredValue(value.get()) : 0;
+				if (requestType == 0) {
+					Future<Void> watchFuture = tr.watch(backupPartitionRequiredKey);
 					co_await tr.commit();
-					TraceEvent("DDBackupPartitionsCleared", ddId);
-					break;
+					co_await watchFuture;
 				}
-
-				Future<Void> watchFuture = tr.watch(backupPartitionRequiredKey);
-				co_await tr.commit();
-				co_await watchFuture;
-				break;
+				continue;
 			} catch (Error& e) {
 				err = e;
 			}
 			co_await tr.onError(err);
 		}
+
+		// Phase 2: compute outside any transaction (may wait long on shard metrics).
+		std::vector<KeyRange> partitions;
+		if (requestType == 1) {
+			partitions = co_await calculateBackupPartitionKeyRanges(shards);
+		}
+
+		// Phase 3: short txn to re-check the request value and write the result.
+		{
+			ReadYourWritesTransaction tr(cx);
+			while (true) {
+				Error err;
+				try {
+					tr.setOption(FDBTransactionOptions::ACCESS_SYSTEM_KEYS);
+					tr.setOption(FDBTransactionOptions::LOCK_AWARE);
+					Optional<Value> value = co_await tr.get(backupPartitionRequiredKey);
+					int8_t currentType = value.present() ? decodeBackupPartitionRequiredValue(value.get()) : 0;
+					if (currentType != requestType) {
+						// Someone wrote a new request while we were computing partitions; restart the outer loop
+						// so the next iteration acts on the new value.
+						break;
+					}
+					if (requestType == 1) {
+						tr.set(backupPartitionListKey, encodeBackupPartitionListValue(partitions));
+						tr.clear(backupPartitionRequiredKey);
+						co_await tr.commit();
+						TraceEvent("DDBackupPartitionsComputed", ddId).detail("NumPartitions", partitions.size());
+					} else {
+						tr.clear(backupPartitionListKey);
+						tr.clear(backupPartitionRequiredKey);
+						co_await tr.commit();
+						TraceEvent("DDBackupPartitionsCleared", ddId);
+					}
+					break;
+				} catch (Error& e) {
+					err = e;
+				}
+				co_await tr.onError(err);
+			}
+		}
 	}
 }