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

Add peer/seed info helper functions, and tests

chihaya.go

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

chihaya_test.go

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

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

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

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