From 6693d442a6e20df196738989df3db355cc0357ba Mon Sep 17 00:00:00 2001
From: Kleissner <peter@kleissner.investments>
Date: Sat, 24 Apr 2021 22:42:39 +0200
Subject: [PATCH] Add function dht.RefreshBuckets. Use RW mutex for accessing
 hash table.

---
 dht/DHT Lite.go   | 12 ++++++++++++
 dht/Hash Table.go | 44 ++++++++++++++++++++++++++++----------------
 dht/README.md     |  6 +++---
 3 files changed, 43 insertions(+), 19 deletions(-)

diff --git a/dht/DHT Lite.go b/dht/DHT Lite.go
index 36b2063..4d1cdad 100644
--- a/dht/DHT Lite.go	
+++ b/dht/DHT Lite.go	
@@ -233,3 +233,15 @@ func (dht *DHT) FindNode(key []byte) (value []byte, found bool, err error) {
 		closestNode = sl.Nodes[0]
 	}
 }
+
+// ---- DHT Health ----
+
+// RefreshBuckets refreshes all buckets not meeting the target node number. 0 to refresh all.
+func (dht *DHT) RefreshBuckets(target int) {
+	for bucket, total := range dht.ht.getTotalNodesPerBucket() {
+		if target == 0 || total < target {
+			nodeR := dht.ht.getRandomIDFromBucket(bucket)
+			dht.FindNode(nodeR)
+		}
+	}
+}
diff --git a/dht/Hash Table.go b/dht/Hash Table.go
index 3dc48ca..c3ea3b2 100644
--- a/dht/Hash Table.go	
+++ b/dht/Hash Table.go	
@@ -34,14 +34,14 @@ type hashTable struct {
 	//  └ Least recently seen                    Most recently seen ┘
 	RoutingTable [][]*Node // bBits x bSize
 
-	mutex *sync.Mutex
+	mutex *sync.RWMutex
 }
 
 func newHashTable(self *Node, bits, bucketSize int) *hashTable {
 	ht := &hashTable{
 		bBits: bits,
 		bSize: bucketSize,
-		mutex: &sync.Mutex{},
+		mutex: &sync.RWMutex{},
 		Self:  self,
 	}
 
@@ -75,8 +75,8 @@ func (ht *hashTable) markNodeAsSeen(index int, ID []byte) {
 }
 
 func (ht *hashTable) doesNodeExistInBucket(bucket int, node []byte) bool {
-	ht.mutex.Lock()
-	defer ht.mutex.Unlock()
+	ht.mutex.RLock()
+	defer ht.mutex.RUnlock()
 	for _, v := range ht.RoutingTable[bucket] {
 		if bytes.Compare(v.ID, node) == 0 {
 			return true
@@ -87,8 +87,8 @@ func (ht *hashTable) doesNodeExistInBucket(bucket int, node []byte) bool {
 
 // getClosestContacts returns the closest nodes to the target. filterFunc is optional and allows the caller to filter the nodes.
 func (ht *hashTable) getClosestContacts(num int, target []byte, filterFunc NodeFilterFunc, ignoredNodes ...[]byte) *shortList {
-	ht.mutex.Lock()
-	defer ht.mutex.Unlock()
+	ht.mutex.RLock()
+	defer ht.mutex.RUnlock()
 
 	// First we need to build the list of adjacent indices to our target in order
 	index := ht.getBucketIndexFromDifferingBit(target)
@@ -181,14 +181,14 @@ func (ht *hashTable) removeNode(ID []byte) {
 }
 
 func (ht *hashTable) getTotalNodesInBucket(bucket int) int {
-	ht.mutex.Lock()
-	defer ht.mutex.Unlock()
+	ht.mutex.RLock()
+	defer ht.mutex.RUnlock()
 	return len(ht.RoutingTable[bucket])
 }
 
 func (ht *hashTable) getRandomIDFromBucket(bucket int) []byte {
-	ht.mutex.Lock()
-	defer ht.mutex.Unlock()
+	ht.mutex.RLock()
+	defer ht.mutex.RUnlock()
 	// Set the new ID to to be equal in every byte up to
 	// the byte of the first differing bit in the bucket
 
@@ -228,8 +228,8 @@ func (ht *hashTable) getRandomIDFromBucket(bucket int) []byte {
 }
 
 func (ht *hashTable) lastSeenBefore(cutoff time.Time) (nodes []*Node) {
-	ht.mutex.Lock()
-	defer ht.mutex.Unlock()
+	ht.mutex.RLock()
+	defer ht.mutex.RUnlock()
 	nodes = make([]*Node, 0, ht.bSize)
 	for _, v := range ht.RoutingTable {
 		for _, n := range v {
@@ -266,8 +266,8 @@ func (ht *hashTable) getBucketIndexFromDifferingBit(id1 []byte) int {
 }
 
 func (ht *hashTable) totalNodes() int {
-	ht.mutex.Lock()
-	defer ht.mutex.Unlock()
+	ht.mutex.RLock()
+	defer ht.mutex.RUnlock()
 	var total int
 	for _, v := range ht.RoutingTable {
 		total += len(v)
@@ -276,8 +276,8 @@ func (ht *hashTable) totalNodes() int {
 }
 
 func (ht *hashTable) Nodes() (nodes []*Node) {
-	ht.mutex.Lock()
-	defer ht.mutex.Unlock()
+	ht.mutex.RLock()
+	defer ht.mutex.RUnlock()
 	nodes = make([]*Node, 0, ht.bSize)
 	for _, v := range ht.RoutingTable {
 		nodes = append(nodes, v...)
@@ -296,3 +296,15 @@ func hasBit(n byte, pos uint) bool {
 	val := n & (1 << pos)
 	return (val > 0)
 }
+
+// getTotalNodesPerBucket returns the count of nodes in all buckets
+func (ht *hashTable) getTotalNodesPerBucket() (total []int) {
+	ht.mutex.RLock()
+	defer ht.mutex.RUnlock()
+
+	for n, _ := range ht.RoutingTable {
+		total = append(total, len(ht.RoutingTable[n]))
+	}
+
+	return total
+}
diff --git a/dht/README.md b/dht/README.md
index fbc94e2..491adbc 100644
--- a/dht/README.md
+++ b/dht/README.md
@@ -3,7 +3,7 @@
 This code is a fork from https://github.com/james-lawrence/kademlia and https://github.com/prettymuchbryce/kademlia with modifications for proper abstraction. All networking code was removed from the original one. This package shall only provide DHT fuctionality.
 
 The following functions are not handled here and must be done by the caller, if desired:
-* Remove nodes that are deemed inactive via `dht.RemoveNode`
-* Provide a function `ShouldEvict` to determine if a node shall be evicted in favor of another one
-* Refreshing buckets: Pick random node in the unrefreshed bucket, ask closest alpha neighbors. This makes sure the bucket remains accessible?
+* Remove nodes that are deemed inactive via `dht.RemoveNode`.
+* Provide a function `ShouldEvict` to determine if a node shall be evicted in favor of another one.
+* Refresh buckets via `dht.RefreshBuckets`.
 * The actual store data functions (and associated replication/expiration) are not provided, only the functionality to traverse through the network.