dubbo.apache.org/dubbo-go/v3@v3.1.1/cluster/loadbalance/roundrobin/loadbalance.go (about)

     1  /*
     2   * Licensed to the Apache Software Foundation (ASF) under one or more
     3   * contributor license agreements.  See the NOTICE file distributed with
     4   * this work for additional information regarding copyright ownership.
     5   * The ASF licenses this file to You under the Apache License, Version 2.0
     6   * (the "License"); you may not use this file except in compliance with
     7   * the License.  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 roundrobin
    19  
    20  import (
    21  	"math"
    22  	"sync"
    23  	"sync/atomic"
    24  	"time"
    25  )
    26  
    27  import (
    28  	"dubbo.apache.org/dubbo-go/v3/cluster/loadbalance"
    29  	"dubbo.apache.org/dubbo-go/v3/common/constant"
    30  	"dubbo.apache.org/dubbo-go/v3/common/extension"
    31  	"dubbo.apache.org/dubbo-go/v3/protocol"
    32  )
    33  
    34  const (
    35  	Complete = 0
    36  	Updating = 1
    37  )
    38  
    39  var (
    40  	methodWeightMap sync.Map          // [string]invokers
    41  	state           = int32(Complete) // update lock acquired ?
    42  	recyclePeriod   = 60 * time.Second.Nanoseconds()
    43  )
    44  
    45  func init() {
    46  	extension.SetLoadbalance(constant.LoadBalanceKeyRoundRobin, NewRRLoadBalance)
    47  }
    48  
    49  type rrLoadBalance struct{}
    50  
    51  // NewRRLoadBalance returns a round robin load balance
    52  //
    53  // Use the weight's common advisory to determine round robin ratio
    54  func NewRRLoadBalance() loadbalance.LoadBalance {
    55  	return &rrLoadBalance{}
    56  }
    57  
    58  // Select gets invoker based on round robin load balancing strategy
    59  func (lb *rrLoadBalance) Select(invokers []protocol.Invoker, invocation protocol.Invocation) protocol.Invoker {
    60  	count := len(invokers)
    61  	if count == 0 {
    62  		return nil
    63  	}
    64  	if count == 1 {
    65  		return invokers[0]
    66  	}
    67  
    68  	key := invokers[0].GetURL().Path + "." + invocation.MethodName()
    69  	cache, _ := methodWeightMap.LoadOrStore(key, &cachedInvokers{})
    70  	cachedInvokers := cache.(*cachedInvokers)
    71  
    72  	var (
    73  		clean               = false
    74  		totalWeight         = int64(0)
    75  		maxCurrentWeight    = int64(math.MinInt64)
    76  		now                 = time.Now()
    77  		selectedInvoker     protocol.Invoker
    78  		selectedWeightRobin *weightedRoundRobin
    79  	)
    80  
    81  	for _, invoker := range invokers {
    82  		weight := loadbalance.GetWeight(invoker, invocation)
    83  		if weight < 0 {
    84  			weight = 0
    85  		}
    86  
    87  		identifier := invoker.GetURL().Key()
    88  		loaded, found := cachedInvokers.LoadOrStore(identifier, &weightedRoundRobin{weight: weight})
    89  		weightRobin := loaded.(*weightedRoundRobin)
    90  		if !found {
    91  			clean = true
    92  		}
    93  
    94  		if weightRobin.Weight() != weight {
    95  			weightRobin.setWeight(weight)
    96  		}
    97  
    98  		currentWeight := weightRobin.increaseCurrent()
    99  		weightRobin.lastUpdate = &now
   100  
   101  		if currentWeight > maxCurrentWeight {
   102  			maxCurrentWeight = currentWeight
   103  			selectedInvoker = invoker
   104  			selectedWeightRobin = weightRobin
   105  		}
   106  		totalWeight += weight
   107  	}
   108  
   109  	cleanIfRequired(clean, cachedInvokers, &now)
   110  
   111  	if selectedWeightRobin != nil {
   112  		selectedWeightRobin.Current(totalWeight)
   113  		return selectedInvoker
   114  	}
   115  
   116  	// should never happen
   117  	return invokers[0]
   118  }
   119  
   120  func cleanIfRequired(clean bool, invokers *cachedInvokers, now *time.Time) {
   121  	if clean && atomic.CompareAndSwapInt32(&state, Complete, Updating) {
   122  		defer atomic.CompareAndSwapInt32(&state, Updating, Complete)
   123  		invokers.Range(func(identify, robin interface{}) bool {
   124  			weightedRoundRobin := robin.(*weightedRoundRobin)
   125  			elapsed := now.Sub(*weightedRoundRobin.lastUpdate).Nanoseconds()
   126  			if elapsed > recyclePeriod {
   127  				invokers.Delete(identify)
   128  			}
   129  			return true
   130  		})
   131  	}
   132  }
   133  
   134  // Record the weight of the invoker
   135  type weightedRoundRobin struct {
   136  	weight     int64
   137  	current    int64
   138  	lastUpdate *time.Time
   139  }
   140  
   141  func (robin *weightedRoundRobin) Weight() int64 {
   142  	return atomic.LoadInt64(&robin.weight)
   143  }
   144  
   145  func (robin *weightedRoundRobin) setWeight(weight int64) {
   146  	robin.weight = weight
   147  	robin.current = 0
   148  }
   149  
   150  func (robin *weightedRoundRobin) increaseCurrent() int64 {
   151  	return atomic.AddInt64(&robin.current, robin.weight)
   152  }
   153  
   154  func (robin *weightedRoundRobin) Current(delta int64) {
   155  	atomic.AddInt64(&robin.current, -1*delta)
   156  }
   157  
   158  type cachedInvokers struct {
   159  	sync.Map /*[string]weightedRoundRobin*/
   160  }