Merge pull request #147 from joshdekock/middleware-refactor-peer-seed-leech-info

Add peer/seed info helper functions, and tests
This commit is contained in:
Jimmy Zelinskie 2016-04-02 17:06:45 -04:00
commit f0a5d78036
5 changed files with 300 additions and 7 deletions

View file

@ -72,3 +72,11 @@ type Peer struct {
type Params interface {
String(key string) (string, error)
}
// Equal reports whether peer and x are the same.
func (peer *Peer) Equal(x *Peer) bool {
if peer.ID == x.ID && peer.Port == x.Port && peer.IP.Equal(x.IP) {
return true
}
return false
}

45
chihaya_test.go Normal file
View file

@ -0,0 +1,45 @@
// 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 chihaya
import (
"net"
"testing"
"github.com/stretchr/testify/assert"
)
var (
peers = []struct {
peerID string
ip string
port uint16
}{
{"-AZ3034-6wfG2wk6wWLc", "250.183.81.177", 5720},
{"-BS5820-oy4La2MWGEFj", "fd45:7856:3dae::48", 2878},
{"-TR0960-6ep6svaa61r4", "fd45:7856:3dae::48", 2878},
{"-BS5820-oy4La2MWGEFj", "fd0a:29a8:8445::38", 2878},
{"-BS5820-oy4La2MWGEFj", "fd45:7856:3dae::48", 8999},
}
)
func TestPeerEquality(t *testing.T) {
// Build peers from test data.
var builtPeers []*Peer
for _, peer := range peers {
builtPeers = append(builtPeers, &Peer{
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]))
assert.False(t, builtPeers[1].Equal(builtPeers[2]))
assert.False(t, builtPeers[1].Equal(builtPeers[3]))
assert.False(t, builtPeers[1].Equal(builtPeers[4]))
}

View file

@ -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])
}

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

View file

@ -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