Coverage for qreward / client / load_balancer.py: 100.00%

1"""Load balancing strategies for proxy selection.""" 

2 

3import threading (empty)

4from enum import Enum (empty)

5from typing import List, Optional (empty)

6 

7 

8class LoadBalanceStrategy(Enum): (empty)

9 """Load balancing strategy enumeration.""" 

10 

11 ROUND_ROBIN = "round_robin" (empty)

12 WEIGHTED_ROUND_ROBIN = "weighted_round_robin" (empty)

13 LEAST_CONNECTIONS = "least_connections" (empty)

14 

15 

16class RoundRobinSelector: (empty)

17 """Thread-safe round-robin selector over a list of keys. 

18 

19 Skips keys that are marked unhealthy. 

20 """ 

21 

22 def __init__(self) -> None: (empty)

23 self._index = 0 (empty)

24 self._lock = threading.Lock() (empty)

25 

26 def select( (empty)

27 self, 

28 keys: List[str], 

29 healthy_keys: set[str], 

30 ) -> Optional[str]: 

31 """Select the next healthy key in round-robin order. 

32 

33 Args: 

34 keys: Ordered list of all proxy keys. 

35 healthy_keys: Set of keys currently considered healthy. 

36 

37 Returns: 

38 The selected key, or None if no healthy key exists. 

39 """ 

40 if not keys or not healthy_keys: (empty)

41 return None (empty)

42 

43 with self._lock: (empty)

44 total = len(keys) (empty)

45 for _ in range(total): (empty)

46 key = keys[self._index % total] (empty)

47 self._index += 1 (empty)

48 if key in healthy_keys: (empty)

49 return key (empty)

50 

51 return None (empty)

52 

53 

54class WeightedRoundRobinSelector: (empty)

55 """Thread-safe weighted round-robin selector. 

56 

57 Keys with higher weights are selected more frequently. 

58 Uses the smooth weighted round-robin algorithm (Nginx-style). 

59 """ 

60 

61 def __init__(self) -> None: (empty)

62 self._current_weights: dict[str, int] = {} (empty)

63 self._lock = threading.Lock() (empty)

64 

65 def update_weights(self, weights: dict[str, int]) -> None: (empty)

66 """Update the weight configuration. 

67 

68 Args: 

69 weights: Mapping of key to weight (positive integer). 

70 """ 

71 with self._lock: (empty)

72 for key in weights: (empty)

73 if key not in self._current_weights: (empty)

74 self._current_weights[key] = 0 (empty)

75 removed = set(self._current_weights) - set(weights) (empty)

76 for key in removed: (empty)

77 del self._current_weights[key] (empty)

78 

79 def select( (empty)

80 self, 

81 keys: List[str], 

82 healthy_keys: set[str], 

83 weights: dict[str, int], 

84 ) -> Optional[str]: 

85 """Select the next healthy key using smooth weighted round-robin. 

86 

87 Args: 

88 keys: Ordered list of all proxy keys. 

89 healthy_keys: Set of keys currently considered healthy. 

90 weights: Mapping of key to weight. 

91 

92 Returns: 

93 The selected key, or None if no healthy key exists. 

94 """ 

95 if not keys or not healthy_keys: (empty)

96 return None (empty)

97 

98 with self._lock: (empty)

99 eligible = [k for k in keys if k in healthy_keys and k in weights] (empty)

100 if not eligible: (empty)

101 return None (empty)

102 

103 total_weight = sum(weights[k] for k in eligible) (empty)

104 

105 # Increase current weights by effective weight 

106 for key in eligible: (empty)

107 self._current_weights.setdefault(key, 0) (empty)

108 self._current_weights[key] += weights[key] (empty)

109 

110 # Select the key with the highest current weight 

111 best_key = max(eligible, key=lambda k: self._current_weights[k]) (empty)

112 

113 # Decrease the selected key's current weight 

114 self._current_weights[best_key] -= total_weight (empty)

115 

116 return best_key (empty)