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