github.com/xmplusdev/xmcore@v1.8.11-0.20240412132628-5518b55526af/app/router/weight.go (about)

     1  package router
     2  
     3  import (
     4  	"regexp"
     5  	"strconv"
     6  	"strings"
     7  	"sync"
     8  )
     9  
    10  type weightScaler func(value, weight float64) float64
    11  
    12  var numberFinder = regexp.MustCompile(`\d+(\.\d+)?`)
    13  
    14  // NewWeightManager creates a new WeightManager with settings
    15  func NewWeightManager(s []*StrategyWeight, defaultWeight float64, scaler weightScaler) *WeightManager {
    16  	return &WeightManager{
    17  		settings:      s,
    18  		cache:         make(map[string]float64),
    19  		scaler:        scaler,
    20  		defaultWeight: defaultWeight,
    21  	}
    22  }
    23  
    24  // WeightManager manages weights for specific settings
    25  type WeightManager struct {
    26  	settings      []*StrategyWeight
    27  	cache         map[string]float64
    28  	scaler        weightScaler
    29  	defaultWeight float64
    30  	mu            sync.Mutex
    31  }
    32  
    33  // Get get the weight of specified tag
    34  func (s *WeightManager) Get(tag string) float64 {
    35  	s.mu.Lock()
    36  	defer s.mu.Unlock()
    37  	weight, ok := s.cache[tag]
    38  	if ok {
    39  		return weight
    40  	}
    41  	weight = s.findValue(tag)
    42  	s.cache[tag] = weight
    43  	return weight
    44  }
    45  
    46  // Apply applies weight to the value
    47  func (s *WeightManager) Apply(tag string, value float64) float64 {
    48  	return s.scaler(value, s.Get(tag))
    49  }
    50  
    51  func (s *WeightManager) findValue(tag string) float64 {
    52  	for _, w := range s.settings {
    53  		matched := s.getMatch(tag, w.Match, w.Regexp)
    54  		if matched == "" {
    55  			continue
    56  		}
    57  		if w.Value > 0 {
    58  			return float64(w.Value)
    59  		}
    60  		// auto weight from matched
    61  		numStr := numberFinder.FindString(matched)
    62  		if numStr == "" {
    63  			return s.defaultWeight
    64  		}
    65  		weight, err := strconv.ParseFloat(numStr, 64)
    66  		if err != nil {
    67  			newError("unexpected error from ParseFloat: ", err).AtError().WriteToLog()
    68  			return s.defaultWeight
    69  		}
    70  		return weight
    71  	}
    72  	return s.defaultWeight
    73  }
    74  
    75  func (s *WeightManager) getMatch(tag, find string, isRegexp bool) string {
    76  	if !isRegexp {
    77  		idx := strings.Index(tag, find)
    78  		if idx < 0 {
    79  			return ""
    80  		}
    81  		return find
    82  	}
    83  	r, err := regexp.Compile(find)
    84  	if err != nil {
    85  		newError("invalid regexp: ", find, "err: ", err).AtError().WriteToLog()
    86  		return ""
    87  	}
    88  	return r.FindString(tag)
    89  }