From 8a772520a50abf10c901ab9d5ad81b9d58bcf0ae Mon Sep 17 00:00:00 2001
From: realbigsean <sean@sigmaprime.io>
Date: Mon, 24 Feb 2025 21:00:51 -0800
Subject: [PATCH] Cache validator registration only after successful publish
 (#7034)

---
 .../src/preparation_service.rs                | 29 ++++++++++---------
 1 file changed, 16 insertions(+), 13 deletions(-)

diff --git a/validator_client/validator_services/src/preparation_service.rs b/validator_client/validator_services/src/preparation_service.rs
index fe6eab3a8a7..1f63f932a81 100644
--- a/validator_client/validator_services/src/preparation_service.rs
+++ b/validator_client/validator_services/src/preparation_service.rs
@@ -428,7 +428,7 @@ impl<T: SlotClock + 'static, E: EthSpec> PreparationService<T, E> {
                     pubkey,
                 } = key.clone();
 
-                let signed_data = match self
+                match self
                     .validator_store
                     .sign_validator_registration_data(ValidatorRegistrationData {
                         fee_recipient,
@@ -458,13 +458,7 @@ impl<T: SlotClock + 'static, E: EthSpec> PreparationService<T, E> {
                         );
                         continue;
                     }
-                };
-
-                self.validator_registration_cache
-                    .write()
-                    .insert(key, signed_data.clone());
-
-                signed_data
+                }
             };
             signed.push(signed_data);
         }
@@ -478,11 +472,20 @@ impl<T: SlotClock + 'static, E: EthSpec> PreparationService<T, E> {
                     })
                     .await
                 {
-                    Ok(()) => info!(
-                        log,
-                        "Published validator registrations to the builder network";
-                        "count" => batch.len(),
-                    ),
+                    Ok(()) => {
+                        info!(
+                            log,
+                            "Published validator registrations to the builder network";
+                            "count" => batch.len(),
+                        );
+                        let mut guard = self.validator_registration_cache.write();
+                        for signed_data in batch {
+                            guard.insert(
+                                ValidatorRegistrationKey::from(signed_data.message.clone()),
+                                signed_data.clone(),
+                            );
+                        }
+                    }
                     Err(e) => warn!(
                         log,
                         "Unable to publish validator registrations to the builder network";