Merge pull request #220 from mrd0ll4r/storage-scrape-support

Storage scrape support

storage/memory/peer_store.go

@@ -84,10 +84,13 @@ type peerStore struct {
 
 var _ storage.PeerStore = &peerStore{}
 
-func (s *peerStore) shardIndex(infoHash bittorrent.InfoHash, p bittorrent.Peer) uint32 {
+func (s *peerStore) shardIndex(infoHash bittorrent.InfoHash, v6 bool) uint32 {
-	idx := binary.BigEndian.Uint32(infoHash[:4]) % uint32(len(s.shards))
+	// There are twice the amount of shards specified by the user, the first
-	if len(p.IP) == net.IPv6len {
+	// half is dedicated to IPv4 swarms and the second half is dedicated to
-		idx += idx + uint32(len(s.shards)/2)
+	// IPv6 swarms.
+	idx := binary.BigEndian.Uint32(infoHash[:4]) % (uint32(len(s.shards)) / 2)
+	if v6 {
+		idx += uint32(len(s.shards) / 2)
 	}
 	return idx
 }
@@ -118,7 +121,7 @@ func (s *peerStore) PutSeeder(ih bittorrent.InfoHash, p bittorrent.Peer) error {
 
 	pk := newPeerKey(p)
 
-	shard := s.shards[s.shardIndex(ih, p)]
+	shard := s.shards[s.shardIndex(ih, len(p.IP) == net.IPv6len)]
 	shard.Lock()
 
 	if _, ok := shard.swarms[ih]; !ok {
@@ -143,7 +146,7 @@ func (s *peerStore) DeleteSeeder(ih bittorrent.InfoHash, p bittorrent.Peer) erro
 
 	pk := newPeerKey(p)
 
-	shard := s.shards[s.shardIndex(ih, p)]
+	shard := s.shards[s.shardIndex(ih, len(p.IP) == net.IPv6len)]
 	shard.Lock()
 
 	if _, ok := shard.swarms[ih]; !ok {
@@ -175,7 +178,7 @@ func (s *peerStore) PutLeecher(ih bittorrent.InfoHash, p bittorrent.Peer) error
 
 	pk := newPeerKey(p)
 
-	shard := s.shards[s.shardIndex(ih, p)]
+	shard := s.shards[s.shardIndex(ih, len(p.IP) == net.IPv6len)]
 	shard.Lock()
 
 	if _, ok := shard.swarms[ih]; !ok {
@@ -200,7 +203,7 @@ func (s *peerStore) DeleteLeecher(ih bittorrent.InfoHash, p bittorrent.Peer) err
 
 	pk := newPeerKey(p)
 
-	shard := s.shards[s.shardIndex(ih, p)]
+	shard := s.shards[s.shardIndex(ih, len(p.IP) == net.IPv6len)]
 	shard.Lock()
 
 	if _, ok := shard.swarms[ih]; !ok {
@@ -232,7 +235,7 @@ func (s *peerStore) GraduateLeecher(ih bittorrent.InfoHash, p bittorrent.Peer) e
 
 	pk := newPeerKey(p)
 
-	shard := s.shards[s.shardIndex(ih, p)]
+	shard := s.shards[s.shardIndex(ih, len(p.IP) == net.IPv6len)]
 	shard.Lock()
 
 	if _, ok := shard.swarms[ih]; !ok {
@@ -261,7 +264,7 @@ func (s *peerStore) AnnouncePeers(ih bittorrent.InfoHash, seeder bool, numWant i
 		numWant = s.maxNumWant
 	}
 
-	shard := s.shards[s.shardIndex(ih, announcer)]
+	shard := s.shards[s.shardIndex(ih, len(announcer.IP) == net.IPv6len)]
 	shard.RLock()
 
 	if _, ok := shard.swarms[ih]; !ok {
@@ -316,6 +319,28 @@ func (s *peerStore) AnnouncePeers(ih bittorrent.InfoHash, seeder bool, numWant i
 	return
 }
 
+func (s *peerStore) ScrapeSwarm(ih bittorrent.InfoHash, v6 bool) (resp bittorrent.Scrape) {
+	select {
+	case <-s.closed:
+		panic("attempted to interact with stopped memory store")
+	default:
+	}
+
+	shard := s.shards[s.shardIndex(ih, v6)]
+	shard.RLock()
+
+	if _, ok := shard.swarms[ih]; !ok {
+		shard.RUnlock()
+		return
+	}
+
+	resp.Incomplete = uint32(len(shard.swarms[ih].leechers))
+	resp.Complete = uint32(len(shard.swarms[ih].seeders))
+	shard.RUnlock()
+
+	return
+}
+
 // collectGarbage deletes all Peers from the PeerStore which are older than the
 // cutoff time.
 //

storage/storage.go

@@ -12,11 +12,12 @@ var ErrResourceDoesNotExist = bittorrent.ClientError("resource does not exist")
 // PeerStore is an interface that abstracts the interactions of storing and
 // manipulating Peers such that it can be implemented for various data stores.
 type PeerStore interface {
-	// PutSeeder adds a Seeder to the Swarm identified by the provided infoHash.
+	// PutSeeder adds a Seeder to the Swarm identified by the provided
+	// infoHash.
 	PutSeeder(infoHash bittorrent.InfoHash, p bittorrent.Peer) error
 
-	// DeleteSeeder removes a Seeder from the Swarm identified by the provided
+	// DeleteSeeder removes a Seeder from the Swarm identified by the
-	// infoHash.
+	// provided infoHash.
 	//
 	// If the Swarm or Peer does not exist, this function should return
 	// ErrResourceDoesNotExist.
@@ -26,32 +27,44 @@ type PeerStore interface {
 	// infoHash.
 	PutLeecher(infoHash bittorrent.InfoHash, p bittorrent.Peer) error
 
-	// DeleteLeecher removes a Leecher from the Swarm identified by the provided
+	// DeleteLeecher removes a Leecher from the Swarm identified by the
-	// infoHash.
+	// provided infoHash.
 	//
 	// If the Swarm or Peer does not exist, this function should return
 	// ErrResourceDoesNotExist.
 	DeleteLeecher(infoHash bittorrent.InfoHash, p bittorrent.Peer) error
 
-	// GraduateLeecher promotes a Leecher to a Seeder in the Swarm identified by
+	// GraduateLeecher promotes a Leecher to a Seeder in the Swarm
-	// the provided infoHash.
+	// identified by the provided infoHash.
 	//
-	// If the given Peer is not present as a Leecher, add the Peer as a Seeder
+	// If the given Peer is not present as a Leecher, add the Peer as a
-	// and return no error.
+	// Seeder and return no error.
 	GraduateLeecher(infoHash bittorrent.InfoHash, p bittorrent.Peer) error
 
 	// AnnouncePeers is a best effort attempt to return Peers from the Swarm
-	// identified by the provided infoHash. The returned Peers are required to be
+	// identified by the provided infoHash. The returned Peers are required
-	// either all IPv4 or all IPv6.
+	// to be either all IPv4 or all IPv6.
 	//
 	// The returned Peers should strive be:
 	// - as close to length equal to numWant as possible without going over
 	// - all IPv4 or all IPv6 depending on the provided peer
 	// - if seeder is true, should ideally return more leechers than seeders
-	// - if seeder is false, should ideally return more seeders than leechers
+	// - if seeder is false, should ideally return more seeders than
+	//   leechers
 	AnnouncePeers(infoHash bittorrent.InfoHash, seeder bool, numWant int, p bittorrent.Peer) (peers []bittorrent.Peer, err error)
 
-	// Stopper is an interface that expects a Stop method to stops the PeerStore.
+	// ScrapeSwarm returns information required to answer a scrape request
+	// about a swarm identified by the given infohash.
+	// The v6 flag indicates whether or not the IPv6 swarm should be
+	// scraped.
+	// The Complete and Incomplete fields of the Scrape must be filled,
+	// filling the Snatches field is optional.
+	// If the infohash is unknown to the PeerStore, an empty Scrape is
+	// returned.
+	ScrapeSwarm(infoHash bittorrent.InfoHash, v6 bool) bittorrent.Scrape
+
+	// Stopper is an interface that expects a Stop method to stop the
+	// PeerStore.
 	// For more details see the documentation in the stopper package.
 	stopper.Stopper
 }