From a9d3c2e071fe39ca669a6e122ec4da7cee0b883c Mon Sep 17 00:00:00 2001 From: Josh de Kock Date: Thu, 31 Mar 2016 01:45:28 +0100 Subject: [PATCH] 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. --- chihaya_test.go | 11 +- server/store/memory/peer_store.go | 68 +++++++++-- server/store/memory/peer_store_test.go | 159 +++++++++++++++++++++++++ server/store/peer_store.go | 27 +++++ 4 files changed, 253 insertions(+), 12 deletions(-) create mode 100644 server/store/memory/peer_store_test.go diff --git a/chihaya_test.go b/chihaya_test.go index 0186c52..20805e3 100644 --- a/chihaya_test.go +++ b/chihaya_test.go @@ -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])) diff --git a/server/store/memory/peer_store.go b/server/store/memory/peer_store.go index e6004bc..846256e 100644 --- a/server/store/memory/peer_store.go +++ b/server/store/memory/peer_store.go @@ -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]) +} diff --git a/server/store/memory/peer_store_test.go b/server/store/memory/peer_store_test.go new file mode 100644 index 0000000..62e350f --- /dev/null +++ b/server/store/memory/peer_store_test.go @@ -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) +} diff --git a/server/store/peer_store.go b/server/store/peer_store.go index e2d4936..8aadd47 100644 --- a/server/store/peer_store.go +++ b/server/store/peer_store.go @@ -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