store/peer_store: Add tests & functions for getting info about peers
This commit adds four functions: GetSeeders, GetLeechers, NumSeeders, and NumLeechers to the store/peer_store API. The first two functions are for getting a list of all the seeders/leechers by an info hash, the latter two are helper functions which use the Get functions, combine the ipv4, and ipv4, and then len() them. It also adds some minimal tests for memory/peer_store.
This commit is contained in:
Josh de Kock
parent
36e533ba02
commit
a9d3c2e071
4 changed files with 253 additions and 12 deletions
|
|
@ -23,18 +23,19 @@ var (
|
|||
{"-BS5820-oy4La2MWGEFj", "fd0a:29a8:8445::38", 2878},
|
||||
{"-BS5820-oy4La2MWGEFj", "fd45:7856:3dae::48", 8999},
|
||||
}
|
||||
builtPeers []*Peer
|
||||
)
|
||||
|
||||
func TestPeerEquality(t *testing.T) {
|
||||
// Build peers from test data.
|
||||
for i := 0; i < len(peers); i++ {
|
||||
var builtPeers []*Peer
|
||||
for _, peer := range peers {
|
||||
builtPeers = append(builtPeers, &Peer{
|
||||
ID: PeerID(peers[i].peerID),
|
||||
IP: net.ParseIP(peers[i].ip),
|
||||
Port: peers[i].port,
|
||||
ID: PeerID(peer.peerID),
|
||||
IP: net.ParseIP(peer.ip),
|
||||
Port: peer.port,
|
||||
})
|
||||
}
|
||||
|
||||
assert.True(t, builtPeers[0].Equal(builtPeers[0]))
|
||||
assert.False(t, builtPeers[0].Equal(builtPeers[1]))
|
||||
assert.True(t, builtPeers[1].Equal(builtPeers[1]))
|
||||
|
|
|
|||
|
|
@ -28,8 +28,13 @@ func (d *peerStoreDriver) New(storecfg *store.DriverConfig) (store.PeerStore, er
|
|||
return nil, err
|
||||
}
|
||||
|
||||
shards := make([]*peerShard, cfg.Shards)
|
||||
for i := 0; i < cfg.Shards; i++ {
|
||||
shards[i] = &peerShard{}
|
||||
shards[i].peers = make(map[string]map[string]peer)
|
||||
}
|
||||
return &peerStore{
|
||||
shards: make([]*peerShard, cfg.Shards),
|
||||
shards: shards,
|
||||
}, nil
|
||||
}
|
||||
|
||||
|
|
@ -49,6 +54,9 @@ func newPeerStoreConfig(storecfg *store.DriverConfig) (*peerStoreConfig, error)
|
|||
return nil, err
|
||||
}
|
||||
|
||||
if cfg.Shards < 1 {
|
||||
cfg.Shards = 1
|
||||
}
|
||||
return &cfg, nil
|
||||
}
|
||||
|
||||
|
|
@ -88,7 +96,6 @@ func leechersKey(infoHash chihaya.InfoHash) string {
|
|||
|
||||
func (s *peerStore) PutSeeder(infoHash chihaya.InfoHash, p chihaya.Peer) error {
|
||||
key := seedersKey(infoHash)
|
||||
|
||||
shard := s.shards[s.shardIndex(infoHash)]
|
||||
shard.Lock()
|
||||
defer shard.Unlock()
|
||||
|
|
@ -107,7 +114,6 @@ func (s *peerStore) PutSeeder(infoHash chihaya.InfoHash, p chihaya.Peer) error {
|
|||
|
||||
func (s *peerStore) DeleteSeeder(infoHash chihaya.InfoHash, p chihaya.Peer) error {
|
||||
key := seedersKey(infoHash)
|
||||
|
||||
shard := s.shards[s.shardIndex(infoHash)]
|
||||
shard.Lock()
|
||||
defer shard.Unlock()
|
||||
|
|
@ -127,7 +133,6 @@ func (s *peerStore) DeleteSeeder(infoHash chihaya.InfoHash, p chihaya.Peer) erro
|
|||
|
||||
func (s *peerStore) PutLeecher(infoHash chihaya.InfoHash, p chihaya.Peer) error {
|
||||
key := leechersKey(infoHash)
|
||||
|
||||
shard := s.shards[s.shardIndex(infoHash)]
|
||||
shard.Lock()
|
||||
defer shard.Unlock()
|
||||
|
|
@ -146,7 +151,6 @@ func (s *peerStore) PutLeecher(infoHash chihaya.InfoHash, p chihaya.Peer) error
|
|||
|
||||
func (s *peerStore) DeleteLeecher(infoHash chihaya.InfoHash, p chihaya.Peer) error {
|
||||
key := leechersKey(infoHash)
|
||||
|
||||
shard := s.shards[s.shardIndex(infoHash)]
|
||||
shard.Lock()
|
||||
defer shard.Unlock()
|
||||
|
|
@ -167,7 +171,6 @@ func (s *peerStore) DeleteLeecher(infoHash chihaya.InfoHash, p chihaya.Peer) err
|
|||
func (s *peerStore) GraduateLeecher(infoHash chihaya.InfoHash, p chihaya.Peer) error {
|
||||
lkey := leechersKey(infoHash)
|
||||
skey := seedersKey(infoHash)
|
||||
|
||||
shard := s.shards[s.shardIndex(infoHash)]
|
||||
shard.Lock()
|
||||
defer shard.Unlock()
|
||||
|
|
@ -224,7 +227,6 @@ func (s *peerStore) CollectGarbage(cutoff time.Time) error {
|
|||
func (s *peerStore) AnnouncePeers(infoHash chihaya.InfoHash, seeder bool, numWant int) (peers, peers6 []chihaya.Peer, err error) {
|
||||
lkey := leechersKey(infoHash)
|
||||
skey := seedersKey(infoHash)
|
||||
|
||||
shard := s.shards[s.shardIndex(infoHash)]
|
||||
shard.RLock()
|
||||
defer shard.RUnlock()
|
||||
|
|
@ -280,3 +282,55 @@ func (s *peerStore) AnnouncePeers(infoHash chihaya.InfoHash, seeder bool, numWan
|
|||
|
||||
return
|
||||
}
|
||||
|
||||
func (s *peerStore) GetSeeders(infoHash chihaya.InfoHash) (peers, peers6 []chihaya.Peer, err error) {
|
||||
key := seedersKey(infoHash)
|
||||
shard := s.shards[s.shardIndex(infoHash)]
|
||||
shard.RLock()
|
||||
defer shard.RUnlock()
|
||||
|
||||
seeders := shard.peers[key]
|
||||
for _, p := range seeders {
|
||||
if p.IP.To4() == nil {
|
||||
peers6 = append(peers6, p.Peer)
|
||||
} else {
|
||||
peers = append(peers, p.Peer)
|
||||
}
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func (s *peerStore) GetLeechers(infoHash chihaya.InfoHash) (peers, peers6 []chihaya.Peer, err error) {
|
||||
key := leechersKey(infoHash)
|
||||
shard := s.shards[s.shardIndex(infoHash)]
|
||||
shard.RLock()
|
||||
defer shard.RUnlock()
|
||||
|
||||
leechers := shard.peers[key]
|
||||
for _, p := range leechers {
|
||||
if p.IP.To4() == nil {
|
||||
peers6 = append(peers6, p.Peer)
|
||||
} else {
|
||||
peers = append(peers, p.Peer)
|
||||
}
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func (s *peerStore) NumSeeders(infoHash chihaya.InfoHash) int {
|
||||
key := seedersKey(infoHash)
|
||||
shard := s.shards[s.shardIndex(infoHash)]
|
||||
shard.RLock()
|
||||
defer shard.RUnlock()
|
||||
|
||||
return len(shard.peers[key])
|
||||
}
|
||||
|
||||
func (s *peerStore) NumLeechers(infoHash chihaya.InfoHash) int {
|
||||
key := leechersKey(infoHash)
|
||||
shard := s.shards[s.shardIndex(infoHash)]
|
||||
shard.RLock()
|
||||
defer shard.RUnlock()
|
||||
|
||||
return len(shard.peers[key])
|
||||
}
|
||||
|
|
|
|||
159
server/store/memory/peer_store_test.go
Normal file
159
server/store/memory/peer_store_test.go
Normal file
|
|
@ -0,0 +1,159 @@
|
|||
// Copyright 2016 The Chihaya Authors. All rights reserved.
|
||||
// Use of this source code is governed by the BSD 2-Clause license,
|
||||
// which can be found in the LICENSE file.
|
||||
|
||||
package memory
|
||||
|
||||
import (
|
||||
"net"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/chihaya/chihaya"
|
||||
"github.com/chihaya/chihaya/server/store"
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func peerInSlice(peer chihaya.Peer, peers []chihaya.Peer) bool {
|
||||
for _, v := range peers {
|
||||
if v.Equal(&peer) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func TestPeerStoreAPI(t *testing.T) {
|
||||
var (
|
||||
hash = chihaya.InfoHash("11111111111111111111")
|
||||
|
||||
peers = []struct {
|
||||
seeder bool
|
||||
peerID string
|
||||
ip string
|
||||
port uint16
|
||||
}{
|
||||
{false, "-AZ3034-6wfG2wk6wWLc", "250.183.81.177", 5720},
|
||||
{false, "-AZ3042-6ozMq5q6Q3NX", "38.241.13.19", 4833},
|
||||
{false, "-BS5820-oy4La2MWGEFj", "fd45:7856:3dae::48", 2878},
|
||||
{false, "-AR6360-6oZyyMWoOOBe", "fd0a:29a8:8445::38", 3167},
|
||||
{true, "-AG2083-s1hiF8vGAAg0", "231.231.49.173", 1453},
|
||||
{true, "-AG3003-lEl2Mm4NEO4n", "254.99.84.77", 7032},
|
||||
{true, "-MR1100-00HS~T7*65rm", "211.229.100.129", 2614},
|
||||
{true, "-LK0140-ATIV~nbEQAMr", "fdad:c435:bf79::12", 4114},
|
||||
{true, "-KT2210-347143496631", "fdda:1b35:7d6e::9", 6179},
|
||||
{true, "-TR0960-6ep6svaa61r4", "fd7f:78f0:4c77::55", 4727},
|
||||
}
|
||||
unmarshalledConfig = struct {
|
||||
Shards int
|
||||
}{
|
||||
1,
|
||||
}
|
||||
config = store.DriverConfig{
|
||||
"memory",
|
||||
unmarshalledConfig,
|
||||
}
|
||||
d = &peerStoreDriver{}
|
||||
)
|
||||
s, err := d.New(&config)
|
||||
assert.Nil(t, err)
|
||||
assert.NotNil(t, s)
|
||||
|
||||
for _, p := range peers {
|
||||
// Construct chihaya.Peer from test data.
|
||||
peer := chihaya.Peer{
|
||||
chihaya.PeerID(p.peerID),
|
||||
net.ParseIP(p.ip),
|
||||
p.port,
|
||||
}
|
||||
|
||||
if p.seeder {
|
||||
err = s.PutSeeder(hash, peer)
|
||||
} else {
|
||||
err = s.PutLeecher(hash, peer)
|
||||
}
|
||||
assert.Nil(t, err)
|
||||
}
|
||||
|
||||
leechers1, leechers61, err := s.GetLeechers(hash)
|
||||
assert.Nil(t, err)
|
||||
assert.NotEmpty(t, leechers1)
|
||||
assert.NotEmpty(t, leechers61)
|
||||
num := s.NumLeechers(hash)
|
||||
assert.Equal(t, len(leechers1)+len(leechers61), num)
|
||||
|
||||
seeders1, seeders61, err := s.GetSeeders(hash)
|
||||
assert.Nil(t, err)
|
||||
assert.NotEmpty(t, seeders1)
|
||||
assert.NotEmpty(t, seeders61)
|
||||
num = s.NumSeeders(hash)
|
||||
assert.Equal(t, len(seeders1)+len(seeders61), num)
|
||||
|
||||
leechers := append(leechers1, leechers61...)
|
||||
seeders := append(seeders1, seeders61...)
|
||||
|
||||
for _, p := range peers {
|
||||
// Construct chihaya.Peer from test data.
|
||||
peer := chihaya.Peer{
|
||||
chihaya.PeerID(p.peerID),
|
||||
net.ParseIP(p.ip),
|
||||
p.port,
|
||||
}
|
||||
|
||||
if p.seeder {
|
||||
assert.True(t, peerInSlice(peer, seeders))
|
||||
} else {
|
||||
assert.True(t, peerInSlice(peer, leechers))
|
||||
}
|
||||
|
||||
if p.seeder {
|
||||
err = s.DeleteSeeder(hash, peer)
|
||||
} else {
|
||||
err = s.DeleteLeecher(hash, peer)
|
||||
}
|
||||
assert.Nil(t, err)
|
||||
}
|
||||
|
||||
assert.Zero(t, s.NumLeechers(hash))
|
||||
assert.Zero(t, s.NumSeeders(hash))
|
||||
|
||||
// Re-add all the peers to the peerStore.
|
||||
for _, p := range peers {
|
||||
// Construct chihaya.Peer from test data.
|
||||
peer := chihaya.Peer{
|
||||
chihaya.PeerID(p.peerID),
|
||||
net.ParseIP(p.ip),
|
||||
p.port,
|
||||
}
|
||||
if p.seeder {
|
||||
s.PutSeeder(hash, peer)
|
||||
} else {
|
||||
s.PutLeecher(hash, peer)
|
||||
}
|
||||
}
|
||||
|
||||
// Check that there are 6 seeders, and 4 leechers.
|
||||
assert.Equal(t, 6, s.NumSeeders(hash))
|
||||
assert.Equal(t, 4, s.NumLeechers(hash))
|
||||
peer := chihaya.Peer{
|
||||
chihaya.PeerID(peers[0].peerID),
|
||||
net.ParseIP(peers[0].ip),
|
||||
peers[0].port,
|
||||
}
|
||||
err = s.GraduateLeecher(hash, peer)
|
||||
assert.Nil(t, err)
|
||||
// Check that there are 7 seeders, and 3 leechers after graduating a
|
||||
// leecher to a seeder.
|
||||
assert.Equal(t, 7, s.NumSeeders(hash))
|
||||
assert.Equal(t, 3, s.NumLeechers(hash))
|
||||
|
||||
peers1, peers61, err := s.AnnouncePeers(hash, true, 5)
|
||||
assert.Nil(t, err)
|
||||
assert.NotNil(t, peers1)
|
||||
assert.NotNil(t, peers61)
|
||||
|
||||
err = s.CollectGarbage(time.Now())
|
||||
assert.Nil(t, err)
|
||||
assert.Equal(t, s.NumLeechers(hash), 0)
|
||||
assert.Equal(t, s.NumSeeders(hash), 0)
|
||||
}
|
||||
|
|
@ -15,15 +15,42 @@ var peerStoreDrivers = make(map[string]PeerStoreDriver)
|
|||
|
||||
// PeerStore represents an interface for manipulating peers.
|
||||
type PeerStore interface {
|
||||
// PutSeeder adds a seeder for the infoHash to the PeerStore.
|
||||
PutSeeder(infoHash chihaya.InfoHash, p chihaya.Peer) error
|
||||
// DeleteSeeder removes a seeder for the infoHash from the PeerStore.
|
||||
DeleteSeeder(infoHash chihaya.InfoHash, p chihaya.Peer) error
|
||||
|
||||
// PutLeecher adds a leecher for the infoHash to the PeerStore.
|
||||
PutLeecher(infoHash chihaya.InfoHash, p chihaya.Peer) error
|
||||
// DeleteLeecher removes a leecher for the infoHash from the PeerStore.
|
||||
DeleteLeecher(infoHash chihaya.InfoHash, p chihaya.Peer) error
|
||||
|
||||
// GraduateLeecher promotes a peer from a leecher to a seeder for the
|
||||
// infoHash within the PeerStore.
|
||||
GraduateLeecher(infoHash chihaya.InfoHash, p chihaya.Peer) error
|
||||
// AnnouncePeers returns a list of both IPv4, and IPv6 peers for an
|
||||
// announce.
|
||||
//
|
||||
// If seeder is true then the peers returned will only be leechers, the
|
||||
// ammount of leechers returned will be the smaller value of numWant or the
|
||||
// available leechers.
|
||||
// If it is false then seeders will be returned up until numWant or the
|
||||
// available seeders, whichever is smaller. If the available seeders is less
|
||||
// than numWant then peers are returned until numWant or they run out.
|
||||
AnnouncePeers(infoHash chihaya.InfoHash, seeder bool, numWant int) (peers, peers6 []chihaya.Peer, err error)
|
||||
// CollectGarbage deletes peers from the peerStore which are older than the
|
||||
// cutoff time.
|
||||
CollectGarbage(cutoff time.Time) error
|
||||
|
||||
// GetSeeders gets all the seeders for a particular infoHash.
|
||||
GetSeeders(infoHash chihaya.InfoHash) (peers, peers6 []chihaya.Peer, err error)
|
||||
// GetLeechers gets all the leechers for a particular infoHash.
|
||||
GetLeechers(infoHash chihaya.InfoHash) (peers, peers6 []chihaya.Peer, err error)
|
||||
|
||||
// NumSeeders gets the amount of seeders for a particular infoHash.
|
||||
NumSeeders(infoHash chihaya.InfoHash) int
|
||||
// NumLeechers gets the amount of leechers for a particular infoHash.
|
||||
NumLeechers(infoHash chihaya.InfoHash) int
|
||||
}
|
||||
|
||||
// PeerStoreDriver represents an interface for creating a handle to the storage
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue