grpclb_picker.go 5.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170
  1. /*
  2. *
  3. * Copyright 2017 gRPC authors.
  4. *
  5. * Licensed under the Apache License, Version 2.0 (the "License");
  6. * you may not use this file except in compliance with the License.
  7. * You may obtain a copy of the License at
  8. *
  9. * http://www.apache.org/licenses/LICENSE-2.0
  10. *
  11. * Unless required by applicable law or agreed to in writing, software
  12. * distributed under the License is distributed on an "AS IS" BASIS,
  13. * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14. * See the License for the specific language governing permissions and
  15. * limitations under the License.
  16. *
  17. */
  18. package grpclb
  19. import (
  20. "context"
  21. "sync"
  22. "sync/atomic"
  23. "google.golang.org/grpc/balancer"
  24. lbpb "google.golang.org/grpc/balancer/grpclb/grpc_lb_v1"
  25. "google.golang.org/grpc/codes"
  26. "google.golang.org/grpc/status"
  27. )
  28. // rpcStats is same as lbmpb.ClientStats, except that numCallsDropped is a map
  29. // instead of a slice.
  30. type rpcStats struct {
  31. // Only access the following fields atomically.
  32. numCallsStarted int64
  33. numCallsFinished int64
  34. numCallsFinishedWithClientFailedToSend int64
  35. numCallsFinishedKnownReceived int64
  36. mu sync.Mutex
  37. // map load_balance_token -> num_calls_dropped
  38. numCallsDropped map[string]int64
  39. }
  40. func newRPCStats() *rpcStats {
  41. return &rpcStats{
  42. numCallsDropped: make(map[string]int64),
  43. }
  44. }
  45. // toClientStats converts rpcStats to lbpb.ClientStats, and clears rpcStats.
  46. func (s *rpcStats) toClientStats() *lbpb.ClientStats {
  47. stats := &lbpb.ClientStats{
  48. NumCallsStarted: atomic.SwapInt64(&s.numCallsStarted, 0),
  49. NumCallsFinished: atomic.SwapInt64(&s.numCallsFinished, 0),
  50. NumCallsFinishedWithClientFailedToSend: atomic.SwapInt64(&s.numCallsFinishedWithClientFailedToSend, 0),
  51. NumCallsFinishedKnownReceived: atomic.SwapInt64(&s.numCallsFinishedKnownReceived, 0),
  52. }
  53. s.mu.Lock()
  54. dropped := s.numCallsDropped
  55. s.numCallsDropped = make(map[string]int64)
  56. s.mu.Unlock()
  57. for token, count := range dropped {
  58. stats.CallsFinishedWithDrop = append(stats.CallsFinishedWithDrop, &lbpb.ClientStatsPerToken{
  59. LoadBalanceToken: token,
  60. NumCalls: count,
  61. })
  62. }
  63. return stats
  64. }
  65. func (s *rpcStats) drop(token string) {
  66. atomic.AddInt64(&s.numCallsStarted, 1)
  67. s.mu.Lock()
  68. s.numCallsDropped[token]++
  69. s.mu.Unlock()
  70. atomic.AddInt64(&s.numCallsFinished, 1)
  71. }
  72. func (s *rpcStats) failedToSend() {
  73. atomic.AddInt64(&s.numCallsStarted, 1)
  74. atomic.AddInt64(&s.numCallsFinishedWithClientFailedToSend, 1)
  75. atomic.AddInt64(&s.numCallsFinished, 1)
  76. }
  77. func (s *rpcStats) knownReceived() {
  78. atomic.AddInt64(&s.numCallsStarted, 1)
  79. atomic.AddInt64(&s.numCallsFinishedKnownReceived, 1)
  80. atomic.AddInt64(&s.numCallsFinished, 1)
  81. }
  82. type errPicker struct {
  83. // Pick always returns this err.
  84. err error
  85. }
  86. func (p *errPicker) Pick(ctx context.Context, opts balancer.PickOptions) (balancer.SubConn, func(balancer.DoneInfo), error) {
  87. return nil, nil, p.err
  88. }
  89. // rrPicker does roundrobin on subConns. It's typically used when there's no
  90. // response from remote balancer, and grpclb falls back to the resolved
  91. // backends.
  92. //
  93. // It guaranteed that len(subConns) > 0.
  94. type rrPicker struct {
  95. mu sync.Mutex
  96. subConns []balancer.SubConn // The subConns that were READY when taking the snapshot.
  97. subConnsNext int
  98. }
  99. func (p *rrPicker) Pick(ctx context.Context, opts balancer.PickOptions) (balancer.SubConn, func(balancer.DoneInfo), error) {
  100. p.mu.Lock()
  101. defer p.mu.Unlock()
  102. sc := p.subConns[p.subConnsNext]
  103. p.subConnsNext = (p.subConnsNext + 1) % len(p.subConns)
  104. return sc, nil, nil
  105. }
  106. // lbPicker does two layers of picks:
  107. //
  108. // First layer: roundrobin on all servers in serverList, including drops and backends.
  109. // - If it picks a drop, the RPC will fail as being dropped.
  110. // - If it picks a backend, do a second layer pick to pick the real backend.
  111. //
  112. // Second layer: roundrobin on all READY backends.
  113. //
  114. // It's guaranteed that len(serverList) > 0.
  115. type lbPicker struct {
  116. mu sync.Mutex
  117. serverList []*lbpb.Server
  118. serverListNext int
  119. subConns []balancer.SubConn // The subConns that were READY when taking the snapshot.
  120. subConnsNext int
  121. stats *rpcStats
  122. }
  123. func (p *lbPicker) Pick(ctx context.Context, opts balancer.PickOptions) (balancer.SubConn, func(balancer.DoneInfo), error) {
  124. p.mu.Lock()
  125. defer p.mu.Unlock()
  126. // Layer one roundrobin on serverList.
  127. s := p.serverList[p.serverListNext]
  128. p.serverListNext = (p.serverListNext + 1) % len(p.serverList)
  129. // If it's a drop, return an error and fail the RPC.
  130. if s.Drop {
  131. p.stats.drop(s.LoadBalanceToken)
  132. return nil, nil, status.Errorf(codes.Unavailable, "request dropped by grpclb")
  133. }
  134. // If not a drop but there's no ready subConns.
  135. if len(p.subConns) <= 0 {
  136. return nil, nil, balancer.ErrNoSubConnAvailable
  137. }
  138. // Return the next ready subConn in the list, also collect rpc stats.
  139. sc := p.subConns[p.subConnsNext]
  140. p.subConnsNext = (p.subConnsNext + 1) % len(p.subConns)
  141. done := func(info balancer.DoneInfo) {
  142. if !info.BytesSent {
  143. p.stats.failedToSend()
  144. } else if info.BytesReceived {
  145. p.stats.knownReceived()
  146. }
  147. }
  148. return sc, done, nil
  149. }