Repository: go-kratos/kratos
Pull request #2088 (2022-06-07 14:08:24)

feat: add peer for selector
feat: add peer for selector

codecov-commenter: # [Codecov](https://codecov.io/gh/go-kratos/kratos/pull/2088?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos) Report
> Merging [#2088](https://codecov.io/gh/go-kratos/kratos/pull/2088?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos) (e57926e) into [main](https://codecov.io/gh/go-kratos/kratos/commit/eb280903f067d9aaaf19997b587d48d0335a633c?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos) (eb28090) will **decrease** coverage by `0.06%`.
> The diff coverage is `75.00%`.

> :exclamation: Current head e57926e differs from pull request most recent head 72a3f6e. Consider uploading reports for the commit 72a3f6e to get more accurate results

```diff
@@            Coverage Diff             @@
##             main    #2088      +/-   ##
==========================================
- Coverage   79.07%   79.01%   -0.07%     
==========================================
  Files          83       84       +1     
  Lines        3671     3679       +8     
==========================================
+ Hits         2903     2907       +4     
- Misses        558      561       +3     
- Partials      210      211       +1     
```


| [Impacted Files](https://codecov.io/gh/go-kratos/kratos/pull/2088?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos) | Coverage Δ | |
|---|---|---|
| [selector/default\_selector.go](https://codecov.io/gh/go-kratos/kratos/pull/2088/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos#diff-c2VsZWN0b3IvZGVmYXVsdF9zZWxlY3Rvci5nbw==) | `80.95% <33.33%> (-3.67%)` | :arrow_down: |
| [selector/peer.go](https://codecov.io/gh/go-kratos/kratos/pull/2088/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos#diff-c2VsZWN0b3IvcGVlci5nbw==) | `100.00% <100.00%> (ø)` | |
| [internal/context/context.go](https://codecov.io/gh/go-kratos/kratos/pull/2088/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos#diff-aW50ZXJuYWwvY29udGV4dC9jb250ZXh0Lmdv) | `97.01% <0.00%> (-2.99%)` | :arrow_down: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/go-kratos/kratos/pull/2088?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/go-kratos/kratos/pull/2088?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos). Last update [eb28090...72a3f6e](https://codecov.io/gh/go-kratos/kratos/pull/2088?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos).


Patch:
diff --git a/selector/default_selector.go b/selector/default_selector.go
index 7550befd..0e99249e 100644
--- a/selector/default_selector.go
+++ b/selector/default_selector.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
 }
 
diff --git a/selector/peer.go b/selector/peer.go
new file mode 100644
index 00000000..273ba064
--- /dev/null
+++ b/selector/peer.go
@@ -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
+}
diff --git a/transport/grpc/client.go b/transport/grpc/client.go
index a5174b73..923e5267 100644
--- a/transport/grpc/client.go
+++ b/transport/grpc/client.go
@@ -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
 	}
diff --git a/transport/http/client.go b/transport/http/client.go
index 2ae29ffb..f1c21359 100644
--- a/transport/http/client.go
+++ b/transport/http/client.go
@@ -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)
 	}


Test patch:
diff --git a/selector/peer_test.go b/selector/peer_test.go
new file mode 100644
index 00000000..d0eeeb81
--- /dev/null
+++ b/selector/peer_test.go
@@ -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!")
+	}
+}