feat: add peer for selector (#2088)

* add peer for selector
pull/2090/head
longxboy 2 years ago committed by GitHub
parent eb280903f0
commit 123fc1e6c8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
  1. 4
      selector/default_selector.go
  2. 25
      selector/peer.go
  3. 24
      selector/peer_test.go
  4. 2
      transport/grpc/client.go
  5. 2
      transport/http/client.go

@ -53,6 +53,10 @@ func (d *Default) Select(ctx context.Context, opts ...SelectOption) (selected No
if err != nil {
return nil, nil, err
}
p, ok := FromPeerContext(ctx)
if ok {
p.Node = wn.Raw()
}
return wn.Raw(), done, nil
}

@ -0,0 +1,25 @@
package selector
import (
"context"
)
type peerKey struct{}
// Peer contains the information of the peer for an RPC, such as the address
// and authentication information.
type Peer struct {
// node is the peer node.
Node Node
}
// NewPeerContext creates a new context with peer information attached.
func NewPeerContext(ctx context.Context, p *Peer) context.Context {
return context.WithValue(ctx, peerKey{}, p)
}
// FromPeerContext returns the peer information in ctx if it exists.
func FromPeerContext(ctx context.Context) (p *Peer, ok bool) {
p, ok = ctx.Value(peerKey{}).(*Peer)
return
}

@ -0,0 +1,24 @@
package selector
import (
"context"
"testing"
)
func TestPeer(t *testing.T) {
p := Peer{
Node: mockWeightedNode{},
}
ctx := NewPeerContext(context.Background(), &p)
p2, ok := FromPeerContext(ctx)
if !ok || p2.Node == nil {
t.Fatalf(" no peer found!")
}
}
func TestNotPeer(t *testing.T) {
_, ok := FromPeerContext(context.Background())
if ok {
t.Fatalf("test no peer found peer!")
}
}

@ -184,6 +184,8 @@ func unaryClientInterceptor(ms []middleware.Middleware, timeout time.Duration, f
if len(ms) > 0 {
h = middleware.Chain(ms...)(h)
}
var p selector.Peer
ctx = selector.NewPeerContext(ctx, &p)
_, err := h(ctx, req)
return err
}

@ -247,6 +247,8 @@ func (client *Client) invoke(ctx context.Context, req *http.Request, args interf
}
return reply, nil
}
var p selector.Peer
ctx = selector.NewPeerContext(ctx, &p)
if len(client.opts.middleware) > 0 {
h = middleware.Chain(client.opts.middleware...)(h)
}

Loading…
Cancel
Save