fix portpool allocator

imroc · imroc · commit 915456e173df · 2025-06-21T17:06:38.000+08:00
Signed-off-by: roc &lt;roc@imroc.cc&gt;
diff --git a/internal/portpool/portpool.go b/internal/portpool/portpool.go
@@ -2,6 +2,7 @@ package portpool
 
 import (
 	"context"
+	"iter"
 	"slices"
 	"sync"
 
@@ -112,68 +113,105 @@ type PortPool struct {
 	lbList   []LBKey
 }
 
+func (pp *PortPool) getCache() iter.Seq2[LBKey, map[ProtocolPort]struct{}] {
+	return func(yield func(LBKey, map[ProtocolPort]struct{}) bool) {
+		switch pp.LbPolicy {
+		case constant.LbPolicyInOrder, constant.LbPolicyUniform: // 有序分配 或 均匀分配
+			if pp.LbPolicy == constant.LbPolicyUniform { // 如果是均匀分配，则需要按已分配数量排序，找分配数量最小的 lb 分配
+				slices.SortFunc(pp.lbList, func(a, b LBKey) int {
+					return len(pp.cache[a]) - len(pp.cache[b])
+				})
+			}
+			for _, lbKey := range pp.lbList {
+				if !yield(lbKey, pp.cache[lbKey]) { // 若 yield 返回 false 则中断
+					return
+				}
+			}
+		default: // 默认用 Random，按 map 的 key 顺序遍历（golang map 的 key 是无序的，每次遍历顺序随机）
+			for lbKey, allocated := range pp.cache {
+				if !yield(lbKey, allocated) { // 若 yield 返回 false 则中断
+					return
+				}
+			}
+		}
+	}
+}
+
 func (pp *PortPool) IsLbExists(key LBKey) bool {
 	pp.mu.Lock()
 	defer pp.mu.Unlock()
 	_, exists := pp.cache[key]
 	return exists
 }
 
-// 分配指定端口
-func (pp *PortPool) AllocatePort(ctx context.Context, quota int64, ports ...ProtocolPort) ([]PortAllocation, bool) {
+func (pp *PortPool) AllocatePortFromRange(ctx context.Context, startPort, endPort, quota, segmentLength uint16, protocol string) ([]PortAllocation, bool) {
 	pp.mu.Lock()
 	defer pp.mu.Unlock()
 	if len(pp.cache) == 0 {
 		return nil, true
 	}
+	portNum := 1
+	if protocol == constant.ProtocolTCPUDP {
+		portNum = 2
+	}
 	quotaExceeded := true
-	tryAllocate := func(lbKey LBKey, allocated map[ProtocolPort]struct{}) []PortAllocation {
-		if int64(len(allocated)+len(ports)) > quota { // 监听器数量已满，换下个 lb
-			return nil
+	for lb, allocated := range pp.getCache() {
+		if uint16(len(allocated)+portNum) > quota { // 监听器数量已满，换下个 lb
+			continue
 		}
 		quotaExceeded = false
-		canAllocate := true
-		for _, port := range ports { // 确保所有待分配的端口都未被分配
-			if _, exists := allocated[port.Key()]; exists { // 有端口已被占用，标记无法分配
-				canAllocate = false
-				break
+		for port := startPort; port <= endPort; port += segmentLength { // 遍历该端口池的所有端口号
+			endPort := uint16(0)
+			if segmentLength > 1 {
+				endPort = port + segmentLength - 1
 			}
-		}
-		if canAllocate { // 找到有 lb 可分配端口，分配端口并返回
-			result := []PortAllocation{}
-			for _, port := range ports {
-				allocated[port.Key()] = struct{}{}
-				pa := PortAllocation{
-					PortPool:     pp,
-					ProtocolPort: port,
-					LBKey:        lbKey,
-				}
-				result = append(result, pa)
+			if result := pp.tryAllocateFromLb(lb, allocated, port, endPort, protocol); len(result) > 0 {
+				return result, false
 			}
-			return result
 		}
-		return nil
-	}
-	switch pp.LbPolicy {
-	case constant.LbPolicyInOrder, constant.LbPolicyUniform: // 有序分配 或 均匀分配
-		if pp.LbPolicy == constant.LbPolicyUniform { // 如果是均匀分配，则需要按已分配数量排序，找分配数量最小的 lb 分配
-			slices.SortFunc(pp.lbList, func(a, b LBKey) int {
-				return len(pp.cache[a]) - len(pp.cache[b])
-			})
+	}
+	return nil, quotaExceeded
+}
+
+func (pp *PortPool) tryAllocateFromLb(lbKey LBKey, allocated map[ProtocolPort]struct{}, port, endPort uint16, protocol string) []PortAllocation {
+	ports := portsToAllocate(port, endPort, protocol)
+	for _, port := range ports { // 确保所有待分配的端口都未被分配
+		if _, exists := allocated[port.Key()]; exists { // 有端口已被占用，标记无法分配
+			return nil
+		}
+	}
+	result := []PortAllocation{}
+	for _, port := range ports {
+		allocated[port.Key()] = struct{}{}
+		pa := PortAllocation{
+			PortPool:     pp,
+			ProtocolPort: port,
+			LBKey:        lbKey,
 		}
-		for _, lbKey := range pp.lbList {
-			allocated := pp.cache[lbKey]
-			result := tryAllocate(lbKey, allocated)
-			if len(result) > 0 {
-				return result, false
-			}
+		result = append(result, pa)
+	}
+	return result
+}
+
+// 分配指定端口
+func (pp *PortPool) AllocatePort(ctx context.Context, quota int64, port, endPort uint16, protocol string) ([]PortAllocation, bool) {
+	pp.mu.Lock()
+	defer pp.mu.Unlock()
+	if len(pp.cache) == 0 {
+		return nil, true
+	}
+	quotaExceeded := true
+	portNum := 1
+	if protocol == constant.ProtocolTCPUDP {
+		portNum = 2
+	}
+	for lbKey, allocated := range pp.getCache() {
+		if int64(len(allocated)+portNum) > quota { // 监听器数量已满，换下个 lb
+			continue
 		}
-	default: // 默认用 Random，按 map 的 key 顺序遍历（golang map 的 key 是无序的，每次遍历顺序随机）
-		for lbKey, allocated := range pp.cache {
-			result := tryAllocate(lbKey, allocated)
-			if len(result) > 0 {
-				return result, false
-			}
+		quotaExceeded = false
+		if result := pp.tryAllocateFromLb(lbKey, allocated, port, endPort, protocol); len(result) > 0 {
+			return result, false
 		}
 	}
 	// 所有 lb 都无法分配此端口，返回空结果
diff --git a/internal/portpool/portpools.go b/internal/portpool/portpools.go
@@ -38,82 +38,47 @@ func (pp PortPools) Names() string {
 // 从所有端口池中都分配出指定端口，不同端口池可分配不同端口
 func (pp PortPools) allocatePortAcrossPools(
 	ctx context.Context,
-	startPort, endPort, quota, segmentLength uint16,
-	getPortsToAllocate func(port, endPort uint16) (ports []ProtocolPort),
-) (PortAllocations, error) {
+	startPort, endPort, quota, segmentLength uint16, protocol string,
+) PortAllocations {
 	log.FromContext(ctx).V(10).Info("allocatePortAcrossPools", "pools", pp.Names(), "startPort", startPort, "endPort", endPort, "segmentLength", segmentLength)
 	var allocatedPorts PortAllocations
-LOOP_POOL:
 	for _, pool := range pp { // 遍历所有端口池（由于不需要保证所有端口池的端口号相同，因此外层循环直接遍历端口池）
-		for port := startPort; port <= endPort; port += segmentLength { // 遍历该端口池的所有端口号
-			select {
-			case <-ctx.Done():
-				allocatedPorts.Release()
-				if err := ctx.Err(); err != nil {
-					return nil, errors.WithStack(err)
-				}
-				return nil, nil
-			default:
-			}
-			endPort := uint16(0)
-			if segmentLength > 1 {
-				endPort = port + segmentLength - 1
-			}
-			portsToAllocate := getPortsToAllocate(port, endPort)
-			// 尝试分配端口
-			result, quotaExceeded := pool.AllocatePort(ctx, int64(quota), portsToAllocate...)
-			if quotaExceeded { // 超配额，不可能分配成功，不再继续尝试
-				allocatedPorts.Release()
-				return nil, nil
-			}
-			if len(result) > 0 { // 该端口池分配到了端口，追加到结果中
-				allocatedPorts = append(allocatedPorts, result...)
-				log.FromContext(ctx).V(10).Info("allocated port", "pool", pool.Name, "port", port)
-				continue LOOP_POOL
-			} else {
-				// 该端口池中无法分配此端口，尝试下一个端口
-				log.FromContext(ctx).V(10).Info("no available port can be allocated, try next port", "pool", pool.Name, "port", port)
-				continue
-			}
+		// 尝试分配端口
+		result, quotaExceeded := pool.AllocatePortFromRange(ctx, startPort, endPort, quota, segmentLength, protocol)
+		if quotaExceeded { // 超配额，不可能分配成功，不再继续尝试
+			allocatedPorts.Release()
+			return nil
+		}
+		if len(result) > 0 { // 该端口池分配到了端口，追加到结果中
+			allocatedPorts = append(allocatedPorts, result...)
+			log.FromContext(ctx).V(10).Info("allocated port", "pool", pool.Name, "ports", allocatedPorts)
+		} else { // 只要有一个端口池分配失败，则认为无法分配，释放已分配端口，等待 lb 扩容
+			allocatedPorts.Release()
+			return nil
 		}
-		// 该端口池所有端口都无法分配，不再尝试
-		allocatedPorts.Release()
-		return nil, nil
 	}
 	// 所有端口池都分配成功，返回结果
-	return allocatedPorts, nil
+	return allocatedPorts
 }
 
 // 从所有端口池中都分配出指定端口，不同端口池必须分配相同端口
 func (pp PortPools) allocateSamePortAcrossPools(
 	ctx context.Context,
-	startPort, endPort, quota, segmentLength uint16,
-	getPortsToAllocate func(port, endPort uint16) (ports []ProtocolPort),
-) (PortAllocations, error) {
+	startPort, endPort, quota, segmentLength uint16, protocol string,
+) PortAllocations {
 	log.FromContext(ctx).Info("allocateSamePortAcrossPools", "pools", pp.Names(), "startPort", startPort, "endPort", endPort, "segmentLength", segmentLength)
 LOOP_PORT:
 	for port := startPort; port <= endPort; port += segmentLength { // 遍历所有端口号，确保所有端口池都能分配到相同端口号
 		endPort := uint16(0)
 		if segmentLength > 1 {
 			endPort = port + segmentLength - 1
 		}
-		portsToAllocate := getPortsToAllocate(port, endPort)
 		var allocatedPorts PortAllocations
 		for _, pool := range pp { // 在所有端口池中查找可用端口，TCP 和 UDP 端口号相同且都未被分配，则分配此端口号
-			select {
-			case <-ctx.Done():
-				allocatedPorts.Release()
-				if err := ctx.Err(); err != nil {
-					return nil, errors.WithStack(err)
-				}
-				return nil, nil
-			default:
-			}
-
-			results, quotaExceeded := pool.AllocatePort(ctx, int64(quota), portsToAllocate...)
+			results, quotaExceeded := pool.AllocatePort(ctx, int64(quota), port, endPort, protocol)
 			if quotaExceeded {
 				allocatedPorts.Release()
-				return nil, nil
+				return nil
 			}
 			if len(results) > 0 {
 				// 此端口池分配到了此端口，追加到结果中
@@ -124,10 +89,10 @@ LOOP_PORT:
 			}
 		}
 		// 分配结束，返回结果可能为空）
-		return allocatedPorts, nil
+		return allocatedPorts
 	}
 	// 所有端口池都无法分配，返回空结果
-	return nil, nil
+	return nil
 }
 
 var (
@@ -136,6 +101,28 @@ var (
 	ErrPortPoolNotAllocatable = errors.New("port pool not allocatable")
 )
 
+func portsToAllocate(port, endPort uint16, protocol string) (ports []ProtocolPort) {
+	if protocol == constant.ProtocolTCPUDP {
+		ports = append(ports, ProtocolPort{
+			Port:     port,
+			Protocol: constant.ProtocolTCP,
+			EndPort:  endPort,
+		})
+		ports = append(ports, ProtocolPort{
+			Port:     port,
+			Protocol: constant.ProtocolUDP,
+			EndPort:  endPort,
+		})
+	} else {
+		ports = append(ports, ProtocolPort{
+			Port:     port,
+			Protocol: protocol,
+			EndPort:  endPort,
+		})
+	}
+	return
+}
+
 // 从一个或多个端口池中分配一个指定协议的端口，分配成功返回端口号，失败返回错误
 func (pp PortPools) AllocatePort(ctx context.Context, protocol string, useSamePortAcrossPools bool) (ports PortAllocations, err error) {
 	startPort := uint16(0)
@@ -187,66 +174,10 @@ func (pp PortPools) AllocatePort(ctx context.Context, protocol string, useSamePo
 		segmentLength = 1
 	}
 
-	getPortsToAllocate := func(port, endPort uint16) (ports []ProtocolPort) {
-		if protocol == constant.ProtocolTCPUDP {
-			ports = append(ports, ProtocolPort{
-				Port:     port,
-				Protocol: constant.ProtocolTCP,
-				EndPort:  endPort,
-			})
-			ports = append(ports, ProtocolPort{
-				Port:     port,
-				Protocol: constant.ProtocolUDP,
-				EndPort:  endPort,
-			})
-		} else {
-			ports = append(ports, ProtocolPort{
-				Port:     port,
-				Protocol: protocol,
-				EndPort:  endPort,
-			})
-		}
-		return
-	}
-
 	if useSamePortAcrossPools {
-		ports, err = pp.allocateSamePortAcrossPools(ctx, startPort, endPort, quota, segmentLength, getPortsToAllocate)
+		ports = pp.allocateSamePortAcrossPools(ctx, startPort, endPort, quota, segmentLength, protocol)
 	} else {
-		ports, err = pp.allocatePortAcrossPools(ctx, startPort, endPort, quota, segmentLength, getPortsToAllocate)
-	}
-	if err != nil {
-		return nil, errors.WithStack(err)
+		ports = pp.allocatePortAcrossPools(ctx, startPort, endPort, quota, segmentLength, protocol)
 	}
 	return ports, nil
 }
-
-// func (pp PortPools) ReleasePort(lbId string, port uint16, protocol string) {
-// 	if protocol == "TCPUDP" {
-// 		tcpPort := ProtocolPort{
-// 			Port:     port,
-// 			Protocol: "TCP",
-// 		}
-// 		udpPort := ProtocolPort{
-// 			Port:     port,
-// 			Protocol: "UDP",
-// 		}
-// 		for _, portPool := range pp {
-// 			if portCache := portPool.cache[lbId]; portCache != nil {
-// 				delete(portCache, tcpPort)
-// 				delete(portCache, udpPort)
-// 				break
-// 			}
-// 		}
-// 	} else {
-// 		port := ProtocolPort{
-// 			Port:     port,
-// 			Protocol: protocol,
-// 		}
-// 		for _, portPool := range pp {
-// 			if portCache := portPool.cache[lbId]; portCache != nil {
-// 				delete(portCache, port)
-// 				break
-// 			}
-// 		}
-// 	}
-// }
