mirror of
https://github.com/sot-tech/mochi.git
synced 2026-04-28 08:29:59 -07:00
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
@@ -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)
|
||||
}
|
||||
Reference in New Issue
Block a user