refac(bitswap, util) extract KeySet

exchange/bitswap/bitswap.go

@@ -2,7 +2,6 @@ package bitswap
 
 import (
 	"errors"
-	"sync"
 
 	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"
 	ds "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/datastore.go"
@@ -29,9 +28,7 @@ func NetMessageSession(parent context.Context, p *peer.Peer, s bsnet.NetMessageS
 		strategy:      strategy.New(),
 		routing:       directory,
 		sender:        networkAdapter,
-		wantlist: WantList{
+		wantlist:      u.NewKeySet(),
-			data: make(map[u.Key]struct{}),
-		},
 	}
 	networkAdapter.SetDelegate(bs)
 
@@ -58,38 +55,7 @@ type bitswap struct {
 	// TODO(brian): save the strategy's state to the datastore
 	strategy strategy.Strategy
 
-	wantlist WantList
+	wantlist u.KeySet
-}
-
-type WantList struct {
-	lock sync.RWMutex
-	data map[u.Key]struct{}
-}
-
-func (wl *WantList) Add(k u.Key) {
-	u.DOut("Adding %v to Wantlist\n", k.Pretty())
-	wl.lock.Lock()
-	defer wl.lock.Unlock()
-
-	wl.data[k] = struct{}{}
-}
-
-func (wl *WantList) Remove(k u.Key) {
-	u.DOut("Removing %v from Wantlist\n", k.Pretty())
-	wl.lock.Lock()
-	defer wl.lock.Unlock()
-
-	delete(wl.data, k)
-}
-
-func (wl *WantList) Keys() []u.Key {
-	wl.lock.RLock()
-	defer wl.lock.RUnlock()
-	keys := make([]u.Key, 0)
-	for k, _ := range wl.data {
-		keys = append(keys, k)
-	}
-	return keys
 }
 
 // GetBlock attempts to retrieve a particular block from peers within the

exchange/bitswap/bitswap_test.go

@@ -289,9 +289,7 @@ func session(net tn.Network, rs tn.RoutingServer, id peer.ID) instance {
 		strategy:      strategy.New(),
 		routing:       htc,
 		sender:        adapter,
-		wantlist: WantList{
+		wantlist:      util.NewKeySet(),
-			data: make(map[util.Key]struct{}),
-		},
 	}
 	adapter.SetDelegate(bs)
 	return instance{

exchange/bitswap/strategy/interface.go

@@ -30,21 +30,3 @@ type Strategy interface {
 
 	NumBytesReceivedFrom(*peer.Peer) uint64
 }
-
-type WantList interface {
-	// Peer returns the owner of the WantList
-	Peer() *peer.Peer
-
-	// Intersection returns the keys common to both WantLists
-	Intersection(WantList) WantList
-
-	KeySet
-}
-
-// TODO(brian): potentially move this somewhere more generic. For now, it's
-// useful in BitSwap operations.
-
-type KeySet interface {
-	Contains(u.Key) bool
-	Keys() []u.Key
-}

util/key_set.go

@@ -0,0 +1,46 @@
+package util
+
+import (
+	"sync"
+)
+
+type KeySet interface {
+	Add(Key)
+	Remove(Key)
+	Keys() []Key
+}
+
+type ks struct {
+	lock sync.RWMutex
+	data map[Key]struct{}
+}
+
+func NewKeySet() KeySet {
+	return &ks{
+		data: make(map[Key]struct{}),
+	}
+}
+
+func (wl *ks) Add(k Key) {
+	wl.lock.Lock()
+	defer wl.lock.Unlock()
+
+	wl.data[k] = struct{}{}
+}
+
+func (wl *ks) Remove(k Key) {
+	wl.lock.Lock()
+	defer wl.lock.Unlock()
+
+	delete(wl.data, k)
+}
+
+func (wl *ks) Keys() []Key {
+	wl.lock.RLock()
+	defer wl.lock.RUnlock()
+	keys := make([]Key, 0)
+	for k, _ := range wl.data {
+		keys = append(keys, k)
+	}
+	return keys
+}