Skip to content

Commit

Permalink
Merge pull request #87 from renproject/fix/dht
Browse files Browse the repository at this point in the history
Fix a potential panic in dht.
  • Loading branch information
tok-kkk authored Aug 3, 2021
2 parents 2b85600 + 2c8b0df commit 307b3dc
Show file tree
Hide file tree
Showing 2 changed files with 45 additions and 13 deletions.
18 changes: 6 additions & 12 deletions dht/table.go
Original file line number Diff line number Diff line change
Expand Up @@ -67,7 +67,7 @@ type Table interface {
type InMemTable struct {
self id.Signatory

sortedMu *sync.Mutex
sortedMu *sync.RWMutex
sorted []id.Signatory

addrsBySignatoryMu *sync.Mutex
Expand All @@ -86,7 +86,7 @@ func NewInMemTable(self id.Signatory) *InMemTable {
return &InMemTable{
self: self,

sortedMu: new(sync.Mutex),
sortedMu: new(sync.RWMutex),
sorted: []id.Signatory{},

addrsBySignatoryMu: new(sync.Mutex),
Expand Down Expand Up @@ -166,8 +166,8 @@ func (table *InMemTable) PeerAddress(peerID id.Signatory) (wire.Address, bool) {

// Peers returns the n closest peer IDs.
func (table *InMemTable) Peers(n int) []id.Signatory {
table.sortedMu.Lock()
defer table.sortedMu.Unlock()
table.sortedMu.RLock()
defer table.sortedMu.RUnlock()

if n <= 0 {
// For values of n that are less than, or equal to, zero, return an
Expand All @@ -183,9 +183,9 @@ func (table *InMemTable) Peers(n int) []id.Signatory {

// RandomPeers returns n random peer IDs
func (table *InMemTable) RandomPeers(n int) []id.Signatory {
table.sortedMu.Lock()
table.sortedMu.RLock()
defer table.sortedMu.RUnlock()
m := len(table.sorted)
table.sortedMu.Unlock()

if n <= 0 {
// For values of n that are less than, or equal to, zero, return an
Expand All @@ -195,8 +195,6 @@ func (table *InMemTable) RandomPeers(n int) []id.Signatory {
}
if n >= m {
sigs := make([]id.Signatory, m)
table.sortedMu.Lock()
defer table.sortedMu.Unlock()
copy(sigs, table.sorted)
return sigs
}
Expand All @@ -208,8 +206,6 @@ func (table *InMemTable) RandomPeers(n int) []id.Signatory {
if m <= 10000 || n >= m/50.0 {
shuffled := make([]id.Signatory, n)
indexPerm := rand.Perm(m)
table.sortedMu.Lock()
defer table.sortedMu.Unlock()
for i := 0; i < n; i++ {
shuffled[i] = table.sorted[indexPerm[i]]
}
Expand All @@ -219,8 +215,6 @@ func (table *InMemTable) RandomPeers(n int) []id.Signatory {
// Otherwise, use Floyd's sampling algorithm to select n random elements
set := make(map[int]struct{}, n)
randomSelection := make([]id.Signatory, 0, n)
table.sortedMu.Lock()
defer table.sortedMu.Unlock()
for i := m - n; i < m; i++ {
index := table.randObj.Intn(i)
if _, ok := set[index]; !ok {
Expand Down
40 changes: 39 additions & 1 deletion dht/table_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,14 +2,15 @@ package dht_test

import (
"fmt"
"github.com/renproject/aw/wire"
"log"
"math/rand"
"strconv"
"testing/quick"
"time"

"github.com/renproject/aw/dht"
"github.com/renproject/aw/dht/dhtutil"
"github.com/renproject/aw/wire"
"github.com/renproject/id"

. "github.com/onsi/ginkgo"
Expand Down Expand Up @@ -204,6 +205,43 @@ var _ = Describe("DHT", func() {
}
}
})

It("should work while deleting peers from the table", func() {
table, _ := initDHT()
numAddrs := rand.Intn(100)
numRandAddrs := rand.Intn(numAddrs)

// Insert `numAddrs` random addresses into the store.
deletedPeers := make([]id.Signatory, 0, 50)
for i := 0; i < numAddrs; i++ {
privKey := id.NewPrivKey()
sig := privKey.Signatory()
addr := wire.NewUnsignedAddress(wire.TCP, "172.16.254.1:3000", uint64(time.Now().UnixNano()))
table.AddPeer(sig, addr)
if i < numAddrs/2 {
deletedPeers = append(deletedPeers, sig)
}
}

done := make(chan struct{}, 1)
go func() {
defer close(done)

for i := range deletedPeers{
table.DeletePeer(deletedPeers[i])
}
}()

total := time.Duration(0)
for i := 0; i <50 ; i ++ {
start := time.Now()
table.RandomPeers(numRandAddrs)
duration := time.Now().Sub(start)
total += duration
}
log.Printf("RandomPeers takes %v on average", total/50)
<- done
})
})

Context("when querying the number of addresses", func() {
Expand Down

0 comments on commit 307b3dc

Please sign in to comment.