Merge pull request #70 from skyportsystems/perGroupPeerPicker4.1

portPicker function now takes group name

2 files changed, 18 insertions, 4 deletions

diff --git a/groupcache.go b/groupcache.go
index 9499dbb..c227237 100644
--- a/groupcache.go
+++ b/groupcache.go
@@ -198,7 +198,7 @@ func (g *Group) Name() string {
 
 func (g *Group) initPeers() {
 	if g.peers == nil {
-		g.peers = getPeers()
+		g.peers = getPeers(g.name)
 	}
 }
 
diff --git a/peers.go b/peers.go
index a74a79b..1625ff0 100644
--- a/peers.go
+++ b/peers.go
@@ -47,23 +47,37 @@ type NoPeers struct{}
 func (NoPeers) PickPeer(key string) (peer ProtoGetter, ok bool) { return }
 
 var (
-	portPicker func() PeerPicker
+	portPicker func(groupName string) PeerPicker
 )
 
 // RegisterPeerPicker registers the peer initialization function.
 // It is called once, when the first group is created.
+// Either RegisterPeerPicker or RegisterPerGroupPeerPicker should be
+// called exactly once, but not both.
 func RegisterPeerPicker(fn func() PeerPicker) {
 	if portPicker != nil {
 		panic("RegisterPeerPicker called more than once")
 	}
+	portPicker = func(_ string) PeerPicker { return fn() }
+}
+
+// RegisterPerGroupPeerPicker registers the peer initialization function,
+// which takes the groupName, to be used in choosing a PeerPicker.
+// It is called once, when the first group is created.
+// Either RegisterPeerPicker or RegisterPerGroupPeerPicker should be
+// called exactly once, but not both.
+func RegisterPerGroupPeerPicker(fn func(groupName string) PeerPicker) {
+	if portPicker != nil {
+		panic("RegisterPeerPicker called more than once")
+	}
 	portPicker = fn
 }
 
-func getPeers() PeerPicker {
+func getPeers(groupName string) PeerPicker {
 	if portPicker == nil {
 		return NoPeers{}
 	}
-	pk := portPicker()
+	pk := portPicker(groupName)
 	if pk == nil {
 		pk = NoPeers{}
 	}