@@ -47,23 +47,37 @@ type NoPeers struct{}
4747func (NoPeers ) PickPeer (key string ) (peer ProtoGetter , ok bool ) { return }
4848
4949var (
50- portPicker func () PeerPicker
50+ portPicker func (groupName string ) PeerPicker
5151)
5252
5353// RegisterPeerPicker registers the peer initialization function.
5454// It is called once, when the first group is created.
55+ // Either RegisterPeerPicker or RegisterPerGroupPeerPicker should be
56+ // called exactly once, but not both.
5557func 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 ) {
5670 if portPicker != nil {
5771 panic ("RegisterPeerPicker called more than once" )
5872 }
5973 portPicker = fn
6074}
6175
62- func getPeers () PeerPicker {
76+ func getPeers (groupName string ) PeerPicker {
6377 if portPicker == nil {
6478 return NoPeers {}
6579 }
66- pk := portPicker ()
80+ pk := portPicker (groupName )
6781 if pk == nil {
6882 pk = NoPeers {}
6983 }
0 commit comments