From 9ce6c31021a123beea5f55d7af5df950859dea9f Mon Sep 17 00:00:00 2001
From: Jimmy Zelinskie <jimmy.zelinskie+git@gmail.com>
Date: Wed, 3 May 2017 20:49:20 -0400
Subject: [PATCH] storage/memorysubnet: allocate CIDRMask only once

---
 storage/memorybysubnet/peer_store.go      | 14 +++++++-------
 storage/memorybysubnet/peer_store_test.go |  6 ++++--
 2 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/storage/memorybysubnet/peer_store.go b/storage/memorybysubnet/peer_store.go
index 6e3b2d8..1db414c 100644
--- a/storage/memorybysubnet/peer_store.go
+++ b/storage/memorybysubnet/peer_store.go
@@ -94,8 +94,8 @@ func New(cfg Config) (storage.PeerStore, error) {
 	ps := &peerStore{
 		shards:   make([]*peerShard, shardCount*2),
 		closed:   make(chan struct{}),
-		ipv4Mask: cfg.PreferredIPv4SubnetMaskBits,
-		ipv6Mask: cfg.PreferredIPv6SubnetMaskBits,
+		ipv4Mask: net.CIDRMask(cfg.PreferredIPv4SubnetMaskBits, 32),
+		ipv6Mask: net.CIDRMask(cfg.PreferredIPv6SubnetMaskBits, 128),
 	}
 
 	for i := 0; i < shardCount*2; i++ {
@@ -152,8 +152,8 @@ func (s swarm) lenLeechers() (i int) {
 type peerStore struct {
 	shards   []*peerShard
 	closed   chan struct{}
-	ipv4Mask int
-	ipv6Mask int
+	ipv4Mask net.IPMask
+	ipv6Mask net.IPMask
 }
 
 var _ storage.PeerStore = &peerStore{}
@@ -200,9 +200,9 @@ func (s *peerStore) mask(p bittorrent.Peer) string {
 	var maskedIP net.IP
 	switch p.IP.AddressFamily {
 	case bittorrent.IPv4:
-		maskedIP = p.IP.IP.Mask(net.CIDRMask(s.ipv4Mask, 32))
+		maskedIP = p.IP.IP.Mask(s.ipv4Mask)
 	case bittorrent.IPv6:
-		maskedIP = p.IP.IP.Mask(net.CIDRMask(s.ipv6Mask, 128))
+		maskedIP = p.IP.IP.Mask(s.ipv6Mask)
 	default:
 		panic("IP is neither v4 nor v6")
 	}
@@ -554,7 +554,7 @@ func (s *peerStore) collectGarbage(cutoff time.Time) error {
 			}
 
 			for subnet := range shard.swarms[ih].seeders {
-				for pk, mtime := range shard.swarms[ih].leechers[subnet] {
+				for pk, mtime := range shard.swarms[ih].seeders[subnet] {
 					if mtime <= cutoffUnix {
 						delete(shard.swarms[ih].seeders[subnet], pk)
 					}
diff --git a/storage/memorybysubnet/peer_store_test.go b/storage/memorybysubnet/peer_store_test.go
index 352892c..547a1dc 100644
--- a/storage/memorybysubnet/peer_store_test.go
+++ b/storage/memorybysubnet/peer_store_test.go
@@ -2,14 +2,16 @@ package memorybysubnet
 
 import (
 	"testing"
-
 	"time"
 
 	s "github.com/chihaya/chihaya/storage"
 )
 
 func createNew() s.PeerStore {
-	ps, err := New(Config{ShardCount: 1024, GarbageCollectionInterval: 10 * time.Minute})
+	ps, err := New(Config{
+		ShardCount:                1024,
+		GarbageCollectionInterval: 10 * time.Minute,
+	})
 	if err != nil {
 		panic(err)
 	}