github.com/cloudwego/hertz@v0.9.3/pkg/app/client/loadbalance/weight_random.go (about)

     1  /*
     2   * Copyright 2022 CloudWeGo Authors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package loadbalance
    18  
    19  import (
    20  	"sync"
    21  
    22  	"github.com/bytedance/gopkg/lang/fastrand"
    23  	"github.com/cloudwego/hertz/pkg/app/client/discovery"
    24  	"github.com/cloudwego/hertz/pkg/common/hlog"
    25  	"golang.org/x/sync/singleflight"
    26  )
    27  
    28  type weightedBalancer struct {
    29  	cachedWeightInfo sync.Map
    30  	sfg              singleflight.Group
    31  }
    32  
    33  type weightInfo struct {
    34  	instances []discovery.Instance
    35  	entries   []int
    36  	weightSum int
    37  }
    38  
    39  // NewWeightedBalancer creates a loadbalancer using weighted-random algorithm.
    40  func NewWeightedBalancer() Loadbalancer {
    41  	lb := &weightedBalancer{}
    42  	return lb
    43  }
    44  
    45  func (wb *weightedBalancer) calcWeightInfo(e discovery.Result) *weightInfo {
    46  	w := &weightInfo{
    47  		instances: make([]discovery.Instance, len(e.Instances)),
    48  		weightSum: 0,
    49  		entries:   make([]int, len(e.Instances)),
    50  	}
    51  
    52  	var cnt int
    53  
    54  	for idx := range e.Instances {
    55  		weight := e.Instances[idx].Weight()
    56  		if weight > 0 {
    57  			w.instances[cnt] = e.Instances[idx]
    58  			w.entries[cnt] = weight
    59  			w.weightSum += weight
    60  			cnt++
    61  		} else {
    62  			hlog.SystemLogger().Warnf("Invalid weight=%d on instance address=%s", weight, e.Instances[idx].Address())
    63  		}
    64  	}
    65  
    66  	w.instances = w.instances[:cnt]
    67  
    68  	return w
    69  }
    70  
    71  // Pick implements the Loadbalancer interface.
    72  func (wb *weightedBalancer) Pick(e discovery.Result) discovery.Instance {
    73  	wi, ok := wb.cachedWeightInfo.Load(e.CacheKey)
    74  	if !ok {
    75  		wi, _, _ = wb.sfg.Do(e.CacheKey, func() (interface{}, error) {
    76  			return wb.calcWeightInfo(e), nil
    77  		})
    78  		wb.cachedWeightInfo.Store(e.CacheKey, wi)
    79  	}
    80  
    81  	w := wi.(*weightInfo)
    82  	if w.weightSum <= 0 {
    83  		return nil
    84  	}
    85  
    86  	weight := fastrand.Intn(w.weightSum)
    87  	for i := 0; i < len(w.instances); i++ {
    88  		weight -= w.entries[i]
    89  		if weight < 0 {
    90  			return w.instances[i]
    91  		}
    92  	}
    93  
    94  	return nil
    95  }
    96  
    97  // Rebalance implements the Loadbalancer interface.
    98  func (wb *weightedBalancer) Rebalance(e discovery.Result) {
    99  	wb.cachedWeightInfo.Store(e.CacheKey, wb.calcWeightInfo(e))
   100  }
   101  
   102  // Delete implements the Loadbalancer interface.
   103  func (wb *weightedBalancer) Delete(cacheKey string) {
   104  	wb.cachedWeightInfo.Delete(cacheKey)
   105  }
   106  
   107  func (wb *weightedBalancer) Name() string {
   108  	return "weight_random"
   109  }