Repository: go-kratos/kratos
Pull request #1611 (2021-11-05 07:40:06)

feat: add raw node
return raw node to improve performance 

codecov-commenter: # [Codecov](https://codecov.io/gh/go-kratos/kratos/pull/1611?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos) Report
> Merging [#1611](https://codecov.io/gh/go-kratos/kratos/pull/1611?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos) (90f0d85) into [main](https://codecov.io/gh/go-kratos/kratos/commit/9d2719f9f493b84372719c73e1a7f8a83fbab557?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos) (9d2719f) will **decrease** coverage by `0.13%`.
> The diff coverage is `25.00%`.

[![Impacted file tree graph](https://codecov.io/gh/go-kratos/kratos/pull/1611/graphs/tree.svg?width=650&height=150&src=pr&token=XwnmBTz78x&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos)](https://codecov.io/gh/go-kratos/kratos/pull/1611?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos)

```diff
@@            Coverage Diff             @@
##             main    #1611      +/-   ##
==========================================
- Coverage   81.55%   81.41%   -0.14%     
==========================================
  Files          78       78              
  Lines        3389     3396       +7     
==========================================
+ Hits         2764     2765       +1     
- Misses        426      431       +5     
- Partials      199      200       +1     
```


| [Impacted Files](https://codecov.io/gh/go-kratos/kratos/pull/1611?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos) | Coverage Δ | |
|---|---|---|
| [selector/node/direct/direct.go](https://codecov.io/gh/go-kratos/kratos/pull/1611/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos#diff-c2VsZWN0b3Ivbm9kZS9kaXJlY3QvZGlyZWN0Lmdv) | `85.71% <0.00%> (-14.29%)` | :arrow_down: |
| [selector/node/ewma/node.go](https://codecov.io/gh/go-kratos/kratos/pull/1611/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos#diff-c2VsZWN0b3Ivbm9kZS9ld21hL25vZGUuZ28=) | `87.09% <0.00%> (-1.92%)` | :arrow_down: |
| [selector/default\_selector.go](https://codecov.io/gh/go-kratos/kratos/pull/1611/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos#diff-c2VsZWN0b3IvZGVmYXVsdF9zZWxlY3Rvci5nbw==) | `93.75% <50.00%> (-6.25%)` | :arrow_down: |
| [internal/context/context.go](https://codecov.io/gh/go-kratos/kratos/pull/1611/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%> (ø)` | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/go-kratos/kratos/pull/1611?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/1611?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos). Last update [9d2719f...90f0d85](https://codecov.io/gh/go-kratos/kratos/pull/1611?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/balancer.go b/selector/balancer.go
index 9e5accbd..dbd64bbb 100644
--- a/selector/balancer.go
+++ b/selector/balancer.go
@@ -19,6 +19,9 @@ type BalancerBuilder interface {
 type WeightedNode interface {
 	Node
 
+	// Raw returns the original node
+	Raw() Node
+
 	// Weight is the runtime calculated weight
 	Weight() float64
 
diff --git a/selector/default_selector.go b/selector/default_selector.go
index 0e08d615..af4455f6 100644
--- a/selector/default_selector.go
+++ b/selector/default_selector.go
@@ -38,7 +38,11 @@ func (d *Default) Select(ctx context.Context, opts ...SelectOption) (selected No
 	if len(candidates) == 0 {
 		return nil, nil, ErrNoAvailable
 	}
-	return d.Balancer.Pick(ctx, candidates)
+	wn, done, err := d.Balancer.Pick(ctx, candidates)
+	if err != nil {
+		return nil, nil, err
+	}
+	return wn.Raw(), done, nil
 }
 
 // Apply update nodes info.
diff --git a/selector/node/direct/direct.go b/selector/node/direct/direct.go
index e3bf139d..62b30ed7 100644
--- a/selector/node/direct/direct.go
+++ b/selector/node/direct/direct.go
@@ -50,3 +50,7 @@ func (n *Node) Weight() float64 {
 func (n *Node) PickElapsed() time.Duration {
 	return time.Duration(time.Now().UnixNano() - atomic.LoadInt64(&n.lastPick))
 }
+
+func (n *Node) Raw() selector.Node {
+	return n.Node
+}
diff --git a/selector/node/ewma/node.go b/selector/node/ewma/node.go
index 444d8018..2f2ff7e6 100644
--- a/selector/node/ewma/node.go
+++ b/selector/node/ewma/node.go
@@ -178,3 +178,7 @@ func (n *Node) Weight() (weight float64) {
 func (n *Node) PickElapsed() time.Duration {
 	return time.Duration(time.Now().UnixNano() - atomic.LoadInt64(&n.lastPick))
 }
+
+func (n *Node) Raw() selector.Node {
+	return n.Node
+}
diff --git a/transport/grpc/balancer.go b/transport/grpc/balancer.go
index 94b3d29b..9d0da64d 100644
--- a/transport/grpc/balancer.go
+++ b/transport/grpc/balancer.go
@@ -50,19 +50,15 @@ type Builder struct {
 // Build creates a grpc Picker.
 func (b *Builder) Build(info base.PickerBuildInfo) gBalancer.Picker {
 	nodes := make([]selector.Node, 0)
-	subConns := make(map[string]gBalancer.SubConn)
 	for conn, info := range info.ReadySCs {
-		if _, ok := subConns[info.Address.Addr]; ok {
-			continue
-		}
-		subConns[info.Address.Addr] = conn
-
 		ins, _ := info.Address.Attributes.Value("rawServiceInstance").(*registry.ServiceInstance)
-		nodes = append(nodes, selector.NewNode(info.Address.Addr, ins))
+		nodes = append(nodes, &grpcNode{
+			Node:    selector.NewNode(info.Address.Addr, ins),
+			subConn: conn,
+		})
 	}
 	p := &Picker{
 		selector: b.builder.Build(),
-		subConns: subConns,
 	}
 	p.selector.Apply(nodes)
 	return p
@@ -70,7 +66,6 @@ func (b *Builder) Build(info base.PickerBuildInfo) gBalancer.Picker {
 
 // Picker is a grpc picker.
 type Picker struct {
-	subConns map[string]gBalancer.SubConn
 	selector selector.Selector
 }
 
@@ -87,10 +82,9 @@ func (p *Picker) Pick(info gBalancer.PickInfo) (gBalancer.PickResult, error) {
 	if err != nil {
 		return gBalancer.PickResult{}, err
 	}
-	sub := p.subConns[n.Address()]
 
 	return gBalancer.PickResult{
-		SubConn: sub,
+		SubConn: n.(*grpcNode).subConn,
 		Done: func(di gBalancer.DoneInfo) {
 			done(info.Ctx, selector.DoneInfo{
 				Err:           di.Err,
@@ -113,3 +107,8 @@ func (t Trailer) Get(k string) string {
 	}
 	return ""
 }
+
+type grpcNode struct {
+	selector.Node
+	subConn gBalancer.SubConn
+}


Test patch:
diff --git a/selector/node/direct/direct_test.go b/selector/node/direct/direct_test.go
index cca73368..51689f52 100644
--- a/selector/node/direct/direct_test.go
+++ b/selector/node/direct/direct_test.go
@@ -48,6 +48,6 @@ func TestDirectDefaultWeight(t *testing.T) {
 	time.Sleep(time.Millisecond * 10)
 	done(context.Background(), selector.DoneInfo{})
 	assert.Equal(t, float64(100), wn.Weight())
-	assert.Greater(t, time.Millisecond*15, wn.PickElapsed())
+	assert.Greater(t, time.Millisecond*20, wn.PickElapsed())
 	assert.Less(t, time.Millisecond*5, wn.PickElapsed())
 }
diff --git a/selector/selector_test.go b/selector/selector_test.go
index 8168c9bb..3c0fa13e 100644
--- a/selector/selector_test.go
+++ b/selector/selector_test.go
@@ -17,6 +17,11 @@ type mockWeightedNode struct {
 	lastPick int64
 }
 
+// Raw returns the original node
+func (n *mockWeightedNode) Raw() Node {
+	return n.Node
+}
+
 // Weight is the runtime calculated weight
 func (n *mockWeightedNode) Weight() float64 {
 	if n.InitialWeight() != nil {