Skip to content

Commit 252876c

Browse files
author
auxten
committed
DHT.FindNeighbor add ServerRole filter
1 parent aaee018 commit 252876c

5 files changed

Lines changed: 94 additions & 5 deletions

File tree

‎consistent/consistent.go‎

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -277,7 +277,7 @@ func (c *Consistent) GetTwoNeighbors(name string) (proto.Node, proto.Node, error
277277
}
278278

279279
// GetNeighborsEx returns the N closest distinct nodes to the name input in the circle.
280-
func (c *Consistent) GetNeighborsEx(name string, n int, roles []proto.ServerRole) ([]proto.Node, error) {
280+
func (c *Consistent) GetNeighborsEx(name string, n int, roles proto.ServerRoles) ([]proto.Node, error) {
281281
c.RLock()
282282
defer c.RUnlock()
283283

@@ -296,10 +296,13 @@ func (c *Consistent) GetNeighborsEx(name string, n int, roles []proto.ServerRole
296296
res = make([]proto.Node, 0, n)
297297
elem = *c.circle[c.sortedHashes[i]]
298298
)
299+
var noFilter = roles == nil || len(roles) == 0
299300

300-
res = append(res, elem)
301+
if noFilter || roles.Contains(elem.Role) {
302+
res = append(res, elem)
303+
}
301304

302-
if len(res) == n {
305+
if noFilter && len(res) == n {
303306
return res, nil
304307
}
305308

@@ -308,8 +311,10 @@ func (c *Consistent) GetNeighborsEx(name string, n int, roles []proto.ServerRole
308311
i = 0
309312
}
310313
elem = *c.circle[c.sortedHashes[i]]
311-
if !sliceContainsMember(res, elem) {
312-
res = append(res, elem)
314+
if noFilter || roles.Contains(elem.Role) {
315+
if !sliceContainsMember(res, elem) {
316+
res = append(res, elem)
317+
}
313318
}
314319
if len(res) == n {
315320
break

‎consistent/consistent_test.go‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -454,6 +454,44 @@ func TestGetN(t *testing.T) {
454454
}
455455
}
456456

457+
func TestGetNFilterRole(t *testing.T) {
458+
kms.Unittest = true
459+
os.Remove(testStorePath)
460+
kms.ResetBucket()
461+
462+
x, _ := InitConsistent(testStorePath, new(KMSStorage), false)
463+
defer os.Remove(testStorePath)
464+
n := NewNodeFromID("abcdefg")
465+
n.Role = Leader
466+
x.Add(n)
467+
x.Add(NewNodeFromID("hijklmn"))
468+
x.Add(NewNodeFromID("opqrstu"))
469+
members, err := x.GetNeighborsEx("9999999", 3, ServerRoles{Unknown})
470+
if err != nil {
471+
t.Fatal(err)
472+
}
473+
if len(members) != 2 {
474+
t.Errorf("expected 3 members instead of %d", len(members))
475+
}
476+
if members[0].ID != "opqrstu" {
477+
t.Errorf("wrong members[0]: %v", members[0])
478+
}
479+
if members[1].ID != "hijklmn" {
480+
t.Errorf("wrong members[1]: %v", members[1])
481+
}
482+
483+
members, err = x.GetNeighborsEx("9999999", 3, ServerRoles{Leader, Follower})
484+
if err != nil {
485+
t.Fatal(err)
486+
}
487+
if len(members) != 1 {
488+
t.Errorf("expected 3 members instead of %d", len(members))
489+
}
490+
if members[0].ID != "abcdefg" {
491+
t.Errorf("wrong members[0]: %v", members[0])
492+
}
493+
}
494+
457495
func TestGetNLess(t *testing.T) {
458496
kms.Unittest = true
459497
os.Remove(testStorePath)

‎proto/nodeinfo.go‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -199,3 +199,15 @@ func parseServerRole(roleStr string) (role ServerRole, err error) {
199199

200200
return Unknown, nil
201201
}
202+
203+
type ServerRoles []ServerRole
204+
205+
// Contains returns if given role is in the ServerRoles
206+
func (ss *ServerRoles) Contains(role ServerRole) bool {
207+
for _, s := range *ss {
208+
if s == role {
209+
return true
210+
}
211+
}
212+
return false
213+
}

‎proto/nodeinfo_test.go‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,18 @@ func TestNodeKey_Less(t *testing.T) {
6666
})
6767
}
6868

69+
func TestServerRoles_Contains(t *testing.T) {
70+
Convey("ServerRoles Contains", t, func() {
71+
ss := make(ServerRoles, 0)
72+
73+
So(ss.Contains(Follower), ShouldBeFalse)
74+
So(ss.Contains(Unknown), ShouldBeFalse)
75+
ss = append(ss, Leader)
76+
ss = append(ss, Follower)
77+
So(ss.Contains(Leader), ShouldBeTrue)
78+
})
79+
}
80+
6981
func unmarshalAndMarshal(str string) string {
7082
var role ServerRole
7183
yaml.Unmarshal([]byte(str), &role)

‎route/service_test.go‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,7 @@ func TestDHTService_FindNeighbor_FindNode(t *testing.T) {
9797
node1 := NewNode()
9898
node1.InitNodeCryptoInfo(100 * time.Millisecond)
9999
node1.Addr = "node1 addr"
100+
node1.Role = Miner
100101

101102
reqA := &PingReq{
102103
Node: *node1,
@@ -158,6 +159,27 @@ func TestDHTService_FindNeighbor_FindNode(t *testing.T) {
158159
So(nodeIDList, ShouldContain, string(node2.ID))
159160
})
160161

162+
req = &FindNeighborReq{
163+
NodeID: "123",
164+
Count: 10,
165+
Roles: ServerRoles{Miner},
166+
}
167+
resp = new(FindNeighborResp)
168+
err = client.Call("DHT.FindNeighbor", req, resp)
169+
if err != nil {
170+
log.Error(err)
171+
}
172+
log.Debugf("resp filter: %v", resp)
173+
nodeIDList = make([]string, 0)
174+
for _, n := range resp.Nodes[:] {
175+
nodeIDList = append(nodeIDList, string(n.ID))
176+
}
177+
log.Debugf("nodeIDList filter: %v", nodeIDList)
178+
Convey("test FindNeighbor", t, func() {
179+
So(nodeIDList, ShouldContain, string(node1.ID))
180+
So(nodeIDList, ShouldNotContain, string(node2.ID))
181+
})
182+
161183
reqFN3 := &FindNodeReq{
162184
NodeID: node2.ID,
163185
}

0 commit comments

Comments
 (0)