From 469ccf9aea85c52f06879901b675994530aa108a Mon Sep 17 00:00:00 2001
From: Trevor Clinkenbeard <tclinkenbeard@openai.com>
Date: Fri, 24 Jul 2026 17:45:30 -0700
Subject: [PATCH] Register Attrition injectors only after selection

---
 fdbserver/workloads/MachineAttrition.cpp | 45 +++++++++++++++++++++---
 1 file changed, 41 insertions(+), 4 deletions(-)

diff --git a/fdbserver/workloads/MachineAttrition.cpp b/fdbserver/workloads/MachineAttrition.cpp
index 0642d3c571..dfd65c9981 100644
--- a/fdbserver/workloads/MachineAttrition.cpp
+++ b/fdbserver/workloads/MachineAttrition.cpp
@@ -165,10 +165,6 @@ struct MachineAttritionWorkload : FailureInjectionWorkload {
 		enabled = !clientId && g_network->isSimulated() && faultInjectionActivated;
 		suspendDuration = 10.0;
 		iterate = true;
-		if (enabled) {
-			phase = MachineAttritionPhase::get(wcx.sharedRandomNumber);
-			phaseOwner = phase->registerWorkload(machinesToLeave);
-		}
 	}
 
 	explicit MachineAttritionWorkload(WorkloadContext const& wcx) : FailureInjectionWorkload(wcx) {
@@ -202,6 +198,13 @@ struct MachineAttritionWorkload : FailureInjectionWorkload {
 		}
 	}
 
+	void initFailureInjectionMode(DeterministicRandom&) override {
+		if (enabled) {
+			phase = MachineAttritionPhase::get(sharedRandomNumber);
+			phaseOwner = phase->registerWorkload(machinesToLeave);
+		}
+	}
+
 	bool shouldInject(DeterministicRandom& random,
 	                  const WorkloadRequest& work,
 	                  const unsigned alreadyAdded) const override {
@@ -627,5 +630,39 @@ TEST_CASE("/fdbserver/workloads/MachineAttrition/phaseCoordination") {
 	return Void();
 }
 
+TEST_CASE("/fdbserver/workloads/MachineAttrition/failureInjectorSelection") {
+	WorkloadContext wcx;
+	wcx.clientId = 0;
+	wcx.clientCount = 1;
+	wcx.sharedRandomNumber = -3;
+
+	auto phase = MachineAttritionPhase::get(wcx.sharedRandomNumber);
+	auto explicitOwner = phase->registerWorkload(0);
+	std::vector<LocalityData> machines = { LocalityData(Optional<Standalone<StringRef>>(),
+		                                                Standalone<StringRef>("zone0"_sr),
+		                                                Standalone<StringRef>("zone0"_sr),
+		                                                Optional<Standalone<StringRef>>()) };
+	phase->observe(machines);
+
+	TestWorkloadImpl<MachineAttritionWorkload, true> unselected(wcx, NoOptions{});
+	ASSERT(!unselected.phase);
+	ASSERT(unselected.phaseOwner == 0);
+	auto unselectedTargets = machines;
+	ASSERT(phase->prepareTargets(unselectedTargets, explicitOwner));
+	ASSERT(unselectedTargets.size() == 1);
+
+	TestWorkloadImpl<MachineAttritionWorkload, true> selected(wcx, NoOptions{});
+	ASSERT(!selected.phase);
+	selected.enabled = true;
+	DeterministicRandom random(wcx.sharedRandomNumber);
+	selected.initFailureInjectionMode(random);
+	ASSERT(selected.phase == phase);
+	ASSERT(selected.phaseOwner != 0);
+	auto selectedTargets = machines;
+	ASSERT(!phase->prepareTargets(selectedTargets, explicitOwner));
+
+	return Void();
+}
+
 WorkloadFactory<MachineAttritionWorkload> MachineAttritionWorkloadFactory;
 FailureInjectorFactory<MachineAttritionWorkload> MachineAttritionFailureWorkloadFactory;