From d654beff0d99028ba68c644ff09a629a490f2ed1 Mon Sep 17 00:00:00 2001
From: decanus <7621705+decanus@users.noreply.github.com>
Date: Tue, 16 Jun 2020 20:10:56 +0200
Subject: [PATCH] working on close

---
 dht/dht.go        | 15 +++++++++++++++
 node/node.go      |  9 ++++++++-
 node/node_test.go |  2 +-
 3 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/dht/dht.go b/dht/dht.go
index 9ab0a37..5434bed 100644
--- a/dht/dht.go
+++ b/dht/dht.go
@@ -3,6 +3,7 @@ package dht
 import (
 	"bytes"
 	"context"
+	"log"
 	"sync"
 
 	"github.com/decanus/bureka/dht/state"
@@ -154,6 +155,20 @@ func (d *DHT) Heartbeat(id state.Peer) {
 	}
 }
 
+// Close sends a message to all neighbors that a peer is exiting the DHT.
+func (d *DHT) Close()  {
+	d.MapNeighbors(func(peer state.Peer) {
+		err := d.Send(
+			context.Background(),
+			&pb.Message{Key: peer, Type: pb.Message_NODE_EXIT, Sender: d.ID},
+		)
+
+		if err != nil {
+			log.Println(err)
+		}
+	})
+}
+
 // MapNeighbors iterates over the NeighborhoodSet and calls the process for every peer.
 func (d *DHT) MapNeighbors(process MapFunc) {
 	d.RLock()
diff --git a/node/node.go b/node/node.go
index cac3adc..22a71fb 100644
--- a/node/node.go
+++ b/node/node.go
@@ -2,6 +2,7 @@ package node
 
 import (
 	"context"
+	"errors"
 
 	logging "github.com/ipfs/go-log"
 	"github.com/libp2p/go-libp2p-core/event"
@@ -73,7 +74,7 @@ func (n *Node) FindPeer(ctx context.Context, id peer.ID) (peer.AddrInfo, error)
 	b := []byte(id)
 	p := n.dht.Find(b)
 	if p == nil {
-		return peer.AddrInfo{}, nil // @todo error
+		return peer.AddrInfo{}, errors.New("failed to find peer")
 	}
 
 	id, err := peer.IDFromBytes(p)
@@ -87,3 +88,9 @@ func (n *Node) FindPeer(ctx context.Context, id peer.ID) (peer.AddrInfo, error)
 func (n *Node) Send(ctx context.Context, msg *pb.Message) error {
 	return n.dht.Send(ctx, msg)
 }
+
+// Close closes the node and stops the DHT.
+func (n *Node) Close() {
+	n.dht.Close()
+	n.host.Close()
+}
diff --git a/node/node_test.go b/node/node_test.go
index 1dafd76..6bc5a7c 100644
--- a/node/node_test.go
+++ b/node/node_test.go
@@ -105,7 +105,7 @@ func TestFindPeer(t *testing.T) {
 	dhts := setupNodes(t, ctx, 4)
 	defer func() {
 		for i := 0; i < 4; i++ {
-			dhts[i].host.Close()
+			dhts[i].Close()
 		}
 	}()