@@ -47,23 +47,37 @@ type NoPeers struct{}
47
47
func (NoPeers ) PickPeer (key string ) (peer ProtoGetter , ok bool ) { return }
48
48
49
49
var (
50
- portPicker func () PeerPicker
50
+ portPicker func (groupName string ) PeerPicker
51
51
)
52
52
53
53
// RegisterPeerPicker registers the peer initialization function.
54
54
// It is called once, when the first group is created.
55
+ // Either RegisterPeerPicker or RegisterPerGroupPeerPicker should be
56
+ // called exactly once, but not both.
55
57
func RegisterPeerPicker (fn func () PeerPicker ) {
58
+ if portPicker != nil {
59
+ panic ("RegisterPeerPicker called more than once" )
60
+ }
61
+ portPicker = func (_ string ) PeerPicker { return fn () }
62
+ }
63
+
64
+ // RegisterPerGroupPeerPicker registers the peer initialization function,
65
+ // which takes the groupName, to be used in choosing a PeerPicker.
66
+ // It is called once, when the first group is created.
67
+ // Either RegisterPeerPicker or RegisterPerGroupPeerPicker should be
68
+ // called exactly once, but not both.
69
+ func RegisterPerGroupPeerPicker (fn func (groupName string ) PeerPicker ) {
56
70
if portPicker != nil {
57
71
panic ("RegisterPeerPicker called more than once" )
58
72
}
59
73
portPicker = fn
60
74
}
61
75
62
- func getPeers () PeerPicker {
76
+ func getPeers (groupName string ) PeerPicker {
63
77
if portPicker == nil {
64
78
return NoPeers {}
65
79
}
66
- pk := portPicker ()
80
+ pk := portPicker (groupName )
67
81
if pk == nil {
68
82
pk = NoPeers {}
69
83
}
0 commit comments