Add function dht.RefreshBuckets. Use RW mutex for accessing hash table.

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)
+		}
+	}
+}

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
+}

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.