From e2116b051ff852b340e2a7913a969e551f020145 Mon Sep 17 00:00:00 2001
From: Charles Beauville <charles@flower.dev>
Date: Tue, 14 Nov 2023 19:52:52 +0100
Subject: [PATCH] C++ SDK: Delete local node for gRPC-rere (#2597)

---
 src/cc/flwr/src/grpc_rere.cc | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/src/cc/flwr/src/grpc_rere.cc b/src/cc/flwr/src/grpc_rere.cc
index c4920e986b1..267874a7a0e 100644
--- a/src/cc/flwr/src/grpc_rere.cc
+++ b/src/cc/flwr/src/grpc_rere.cc
@@ -19,6 +19,14 @@ std::optional<flwr::proto::Node> get_node_from_store() {
   return node->second;
 }
 
+void delete_node_from_store() {
+  std::lock_guard<std::mutex> lock(node_store_mutex);
+  auto node = node_store.find(KEY_NODE);
+  if (node == node_store.end() || !node->second.has_value()) {
+    node_store.erase(node);
+  }
+}
+
 std::optional<flwr::proto::TaskIns> get_current_task_ins() {
   std::lock_guard<std::mutex> state_lock(state_mutex);
   auto current_task_ins = state.find(KEY_TASK_INS);
@@ -80,8 +88,7 @@ void delete_node(const std::unique_ptr<flwr::proto::Fleet::Stub> &stub) {
     delete_node_request.release_node(); // Release if status is ok
   }
 
-  // TODO: Check if Node needs to be removed from local map
-  // node_store.erase(node);
+  delete_node_from_store();
 }
 
 std::optional<flwr::proto::TaskIns>