From 2143dd311ecd392762100f749c3a7ff8c63784a1 Mon Sep 17 00:00:00 2001
From: Adam <mimol091@gmail.com>
Date: Sat, 9 Aug 2025 13:33:28 +0400
Subject: [PATCH] removal of unavailable peers

---
 node/peerdiscovery.go      |  9 +++++++++
 node/peerdiscovery_test.go | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/node/peerdiscovery.go b/node/peerdiscovery.go
index 2d3afe8..2218bc1 100644
--- a/node/peerdiscovery.go
+++ b/node/peerdiscovery.go
@@ -8,7 +8,9 @@ import (
 
 	"github.com/volodymyrprokopyuk/go-blockchain/node/rpc"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/status"
 )
 
 type PeerReader interface {
@@ -103,6 +105,13 @@ func (d *PeerDiscovery) DiscoverPeers(period time.Duration) {
       for _, peer := range d.Peers() {
         if peer != d.cfg.NodeAddr {
           peers, err := d.grpcPeerDiscover(peer)
+          errorStatus, assert := status.FromError(err)
+          if assert && errorStatus.Code() == codes.Unavailable {
+            d.mtx.Lock()
+            delete(d.peers, peer)
+            d.mtx.Unlock()
+            continue
+          }
           if err != nil {
             fmt.Println(err)
             continue
diff --git a/node/peerdiscovery_test.go b/node/peerdiscovery_test.go
index 5dcf6e9..26a430c 100644
--- a/node/peerdiscovery_test.go
+++ b/node/peerdiscovery_test.go
@@ -52,3 +52,40 @@ func TestPeerDiscovery(t *testing.T) {
     t.Errorf("bootstrap address %v is not in node known peers", bootAddr)
   }
 }
+
+func TestRemovalOfUnhealthyPeer(t *testing.T) {
+  wg := new(sync.WaitGroup)
+  // Create the peer discovery without staring for the bootstrap node
+  bootPeerDisc := createPeerDiscovery(t.Context(), wg, true, false)
+  // Start the gRPC server on the bootstrap node
+  waitForPeersDiscovery := make(chan struct{})
+  grpcStartSvr(t, bootAddr, func(grpcSrv *grpc.Server) {
+    node := rpc.NewNodeSrv(bootPeerDisc, nil)
+    rpc.RegisterNodeServer(grpcSrv, node)
+
+    go func() {
+      <-waitForPeersDiscovery
+      grpcSrv.Stop()
+    }()
+  })
+  // Create and start the peer discovery for the new node
+  nodePeerDisc := createPeerDiscovery(t.Context(), wg, false, true)
+  // Wait for the peer discovery to discover peers
+  time.Sleep(150 * time.Millisecond)
+  // Verify that the bootstrap node and the new node have discovered each other
+
+  if !slices.Contains(bootPeerDisc.Peers(), nodeAddr) {
+    t.Errorf("node address %v is not in bootstrap known peers", nodeAddr)
+  }
+  if !slices.Contains(nodePeerDisc.Peers(), bootAddr) {
+    t.Errorf("bootstrap address %v is not in node known peers", bootAddr)
+  }
+
+  waitForPeersDiscovery <- struct{}{}
+  // Wait for the peer discovery to try to discover peers
+  time.Sleep(250 * time.Millisecond)
+
+  if len(nodePeerDisc.Peers()) != 0 {
+    t.Errorf("node peer discovery should not have any peers after the bootstrap node is stopped")
+  }
+}