From 09ec58acb2eed0f616151b69bcb125cbf777c3b7 Mon Sep 17 00:00:00 2001
From: Neha Oudin <neha@oudin.red>
Date: Tue, 17 Sep 2024 21:36:51 +0200
Subject: [PATCH] fix: make reconfigure safe

---
 src/charm.py | 71 ++++++++++++++++++++++++++--------------------------
 1 file changed, 36 insertions(+), 35 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index e01c13266..169d1e193 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -639,43 +639,39 @@ def _on_relation_handler(self, event: RelationEvent) -> None:
         if not (self.unit.is_leader() and event.unit) or not self.db_initialised:
             return
 
-        # Reconfigure replica set if needed.
-        self._reconfigure_replica_set()
+        # If needed, reconfigure replicaset (eg : network cut)
+        try:
+            self._reconfigure_replica_set()
+            self.__add_members_to_replicaset(event)
+        except NotReadyError:
+            self.status.set_and_share_status(WaitingStatus("waiting to reconfigure replica set"))
+            logger.error("Deferring reconfigure: another member doing sync right now")
+            event.defer()
+        except PyMongoError as e:
+            self.status.set_and_share_status(WaitingStatus("waiting to reconfigure replica set"))
+            logger.error("Deferring reconfigure: error=%r", e)
+            event.defer()
 
+    def __add_members_to_replicaset(self, event: RelationEvent):
         with MongoDBConnection(self.mongodb_config) as mongo:
-            try:
-                replset_members = mongo.get_replset_members()
-                # compare set of mongod replica set members and juju hosts to avoid the unnecessary
-                # reconfiguration.
-                if replset_members == self.mongodb_config.hosts:
-                    return
+            replset_members = mongo.get_replset_members()
+            # compare set of mongod replica set members and juju hosts to avoid the unnecessary
+            # reconfiguration.
+            if replset_members == self.mongodb_config.hosts:
+                return
 
-                for member in self.mongodb_config.hosts - replset_members:
-                    logger.debug("Adding %s to replica set", member)
-                    with MongoDBConnection(
-                        self.mongodb_config, member, direct=True
-                    ) as direct_mongo:
-                        if not direct_mongo.is_ready:
-                            self.status.set_and_share_status(
-                                WaitingStatus("waiting to reconfigure replica set")
-                            )
-                            logger.debug("Deferring reconfigure: %s is not ready yet.", member)
-                            event.defer()
-                            return
-                    mongo.add_replset_member(member)
-                    self.status.set_and_share_status(ActiveStatus())
-            except NotReadyError:
-                self.status.set_and_share_status(
-                    WaitingStatus("waiting to reconfigure replica set")
-                )
-                logger.error("Deferring reconfigure: another member doing sync right now")
-                event.defer()
-            except PyMongoError as e:
-                self.status.set_and_share_status(
-                    WaitingStatus("waiting to reconfigure replica set")
-                )
-                logger.error("Deferring reconfigure: error=%r", e)
-                event.defer()
+            for member in self.mongodb_config.hosts - replset_members:
+                logger.debug("Adding %s to replica set", member)
+                with MongoDBConnection(self.mongodb_config, member, direct=True) as direct_mongo:
+                    if not direct_mongo.is_ready:
+                        self.status.set_and_share_status(
+                            WaitingStatus("waiting to reconfigure replica set")
+                        )
+                        logger.debug("Deferring reconfigure: %s is not ready yet.", member)
+                        event.defer()
+                        return
+                mongo.add_replset_member(member)
+                self.status.set_and_share_status(ActiveStatus())
 
     def _on_leader_elected(self, event: LeaderElectedEvent) -> None:
         """Generates necessary keyfile and updates replica hosts."""
@@ -777,7 +773,12 @@ def _on_update_status(self, event: UpdateStatusEvent):
             return
 
         # If needed, reconfigure replicaset (eg : network cut)
-        self._reconfigure_replica_set()
+        try:
+            self._reconfigure_replica_set()
+        except ServerSelectionTimeoutError as e:
+            self.status.set_and_share_status(WaitingStatus("waiting to reconfigure replica set"))
+            logger.error("Deferring reconfigure: error=%r", e)
+            return
 
         # Cannot check more advanced MongoDB statuses if mongod hasn't started.
         with MongoDBConnection(self.mongodb_config, "localhost", direct=True) as direct_mongo: