tcjlj/flvx
1
0
Fork 0
mirror of https://github.com/Sagit-chu/flvx.git synced 2026-10-05 09:36:37 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: implement max conn limiter dispatching

Browse Source
This commit is contained in:
sagitchu
2026-04-23 15:55:14 +08:00
parent eacbc12d7c
commit 77321d4fa1
4 changed files with 222 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
go-backend/internal/http/handler/control_plane.go
+54 -5
View File
@@ -266,6 +266,22 @@ func (h *Handler) syncForwardServicesWithWarnings(forward *forwardRecord, method
serviceBase := buildForwardServiceBaseWithResolvedUserTunnel(forward.ID, forward.UserID, userTunnelID)
user, err := h.repo.GetUserByID(forward.UserID)
if err != nil {
return nil, err
}
var cLimiterName string
var maxConnToSet int
if forward.MaxConn > 0 {
maxConnToSet = forward.MaxConn
cLimiterName = fmt.Sprintf("rule_conn_limit_%d", forward.ID)
} else if user != nil && user.MaxConn > 0 {
maxConnToSet = user.MaxConn
cLimiterName = fmt.Sprintf("user_conn_limit_%d", user.ID)
}
for _, fp := range ports {
if limiterID != nil && speed != nil {
if err := h.ensureLimiterOnNode(fp.NodeID, *limiterID, *speed); err != nil {
@@ -283,11 +299,17 @@ func (h *Handler) syncForwardServicesWithWarnings(forward *forwardRecord, method
}
}
if cLimiterName != "" {
if err := h.ensureConnLimiterOnNode(fp.NodeID, cLimiterName, maxConnToSet); err != nil {
warnings = append(warnings, fmt.Sprintf("节点 %d 连接限制器下发失败: %v", fp.NodeID, err))
}
}
node, err := h.getNodeRecord(fp.NodeID)
if err != nil {
return nil, err
}
services := buildForwardServiceConfigs(serviceBase, forward, tunnel, node, fp.Port, strings.TrimSpace(fp.InIP), limiterID)
services := buildForwardServiceConfigs(serviceBase, forward, tunnel, node, fp.Port, strings.TrimSpace(fp.InIP), limiterID, cLimiterName)
_, err = h.sendNodeCommand(node.ID, method, services, true, false)
if err != nil && allowFallbackAdd && method == "UpdateService" {
if isNotFoundError(err) {
@@ -302,7 +324,7 @@ func (h *Handler) syncForwardServicesWithWarnings(forward *forwardRecord, method
}
if err != nil && strings.EqualFold(strings.TrimSpace(method), "UpdateService") && isCannotAssignRequestedAddressError(err) {
var warning string
warning, err = h.fallbackForwardPortToDefaultBind(forward, tunnel, node, fp, serviceBase, limiterID)
warning, err = h.fallbackForwardPortToDefaultBind(forward, tunnel, node, fp, serviceBase, limiterID, cLimiterName)
if err == nil && warning != "" {
warnings = append(warnings, warning)
}
@@ -328,7 +350,7 @@ func (h *Handler) syncForwardServicesWithWarnings(forward *forwardRecord, method
return warnings, nil
}
func (h *Handler) fallbackForwardPortToDefaultBind(forward *forwardRecord, tunnel *tunnelRecord, node *nodeRecord, fp forwardPortRecord, serviceBase string, limiterID *int64) (string, error) {
func (h *Handler) fallbackForwardPortToDefaultBind(forward *forwardRecord, tunnel *tunnelRecord, node *nodeRecord, fp forwardPortRecord, serviceBase string, limiterID *int64, cLimiterName string) (string, error) {
if h == nil || forward == nil || tunnel == nil || node == nil {
return "", errors.New("invalid bind fallback context")
}
@@ -345,7 +367,7 @@ func (h *Handler) fallbackForwardPortToDefaultBind(forward *forwardRecord, tunne
}
time.Sleep(150 * time.Millisecond)
defaultServices := buildForwardServiceConfigs(serviceBase, forward, tunnel, node, fp.Port, "", limiterID)
defaultServices := buildForwardServiceConfigs(serviceBase, forward, tunnel, node, fp.Port, "", limiterID, cLimiterName)
if _, err := h.sendNodeCommand(node.ID, "AddService", defaultServices, true, false); err != nil {
return "", err
}
@@ -1637,7 +1659,7 @@ func compactErrorMessage(msg string) string {
return strings.Join(strings.Fields(strings.ToLower(msg)), "")
}
func buildForwardServiceConfigs(baseName string, forward *forwardRecord, tunnel *tunnelRecord, node *nodeRecord, port int, bindIP string, limiterID *int64) []map[string]interface{} {
func buildForwardServiceConfigs(baseName string, forward *forwardRecord, tunnel *tunnelRecord, node *nodeRecord, port int, bindIP string, limiterID *int64, cLimiterName string) []map[string]interface{} {
protocols := []string{"tcp", "udp"}
services := make([]map[string]interface{}, 0, 2)
targets := splitRemoteTargets(forward.RemoteAddr)
@@ -1680,6 +1702,9 @@ func buildForwardServiceConfigs(baseName string, forward *forwardRecord, tunnel
},
},
}
if cLimiterName != "" {
service["climiter"] = cLimiterName
}
if protocol == "udp" {
listenerMetadata := map[string]interface{}{
"keepAlive": true,
@@ -1795,6 +1820,30 @@ func (h *Handler) ensureLimiterOnNode(nodeID int64, limiterID int64, speed int)
return nil
}
func (h *Handler) ensureConnLimiterOnNode(nodeID int64, limiterName string, maxConn int) error {
limitStr := fmt.Sprintf("$ %d", maxConn)
payload := map[string]interface{}{
"name": limiterName,
"limits": []string{limitStr},
}
if _, err := h.sendNodeCommand(nodeID, "AddCLimiters", payload, false, false); err != nil {
if !isAlreadyExistsMessage(err.Error()) {
return fmt.Errorf("连接限制器下发失败: %w", err)
}
updatePayload := map[string]interface{}{
"limiter": limiterName,
"data": payload,
}
if _, updateErr := h.sendNodeCommand(nodeID, "UpdateCLimiters", updatePayload, false, false); updateErr != nil {
return fmt.Errorf("连接限制器更新失败: %w", updateErr)
}
}
return nil
}
func buildLimiterAddPayload(limiterID int64, speed int) (string, map[string]interface{}) {
rate := float64(speed) / 8.0
limitStr := fmt.Sprintf("$ %.1fMB %.1fMB", rate, rate)
go-backend/internal/http/handler/control_plane_test.go
+4 -4
View File
@@ -378,7 +378,7 @@ func TestRetryTunnelServiceAddWithCleanupReturnsCleanupError(t *testing.T) {
func TestBuildForwardServiceConfigs_UsesBindIPForListen(t *testing.T) {
forward := &forwardRecord{RemoteAddr: "1.2.3.4:80", Strategy: "fifo", TunnelID: 7}
node := &nodeRecord{TCPListenAddr: "[::]", UDPListenAddr: "[::]"}
services := buildForwardServiceConfigs("1_2_0", forward, nil, node, 22000, "10.9.8.7", nil)
services := buildForwardServiceConfigs("1_2_0", forward, nil, node, 22000, "10.9.8.7", nil, "")
if len(services) != 2 {
t.Fatalf("expected 2 services, got %d", len(services))
}
@@ -393,7 +393,7 @@ func TestBuildForwardServiceConfigs_UsesBindIPForListen(t *testing.T) {
func TestBuildForwardServiceConfigs_DefaultListenAddrWhenBindIPEmpty(t *testing.T) {
forward := &forwardRecord{RemoteAddr: "1.2.3.4:80", Strategy: "fifo", TunnelID: 7}
node := &nodeRecord{TCPListenAddr: "0.0.0.0", UDPListenAddr: "[::]"}
services := buildForwardServiceConfigs("1_2_0", forward, nil, node, 22001, "", nil)
services := buildForwardServiceConfigs("1_2_0", forward, nil, node, 22001, "", nil, "")
if len(services) != 2 {
t.Fatalf("expected 2 services, got %d", len(services))
}
@@ -409,7 +409,7 @@ func TestBuildForwardServiceConfigs_DefaultListenAddrWhenBindIPEmpty(t *testing.
func TestBuildForwardServiceConfigs_BindIPAlreadyContainsPort(t *testing.T) {
forward := &forwardRecord{RemoteAddr: "1.2.3.4:80", Strategy: "fifo", TunnelID: 7}
node := &nodeRecord{TCPListenAddr: "[::]", UDPListenAddr: "[::]"}
services := buildForwardServiceConfigs("1_2_0", forward, nil, node, 55555, "3.3.3.3:12345", nil)
services := buildForwardServiceConfigs("1_2_0", forward, nil, node, 55555, "3.3.3.3:12345", nil, "")
if len(services) != 2 {
t.Fatalf("expected 2 services, got %d", len(services))
}
@@ -464,7 +464,7 @@ func TestBuildForwardServiceConfigs_IPv6BindIP(t *testing.T) {
t.Run(tt.name, func(t *testing.T) {
forward := &forwardRecord{RemoteAddr: "1.2.3.4:80", Strategy: "fifo", TunnelID: 7}
node := &nodeRecord{TCPListenAddr: "[::]", UDPListenAddr: "[::]"}
services := buildForwardServiceConfigs("1_2_0", forward, nil, node, tt.port, tt.bindIP, nil)
services := buildForwardServiceConfigs("1_2_0", forward, nil, node, tt.port, tt.bindIP, nil, "")
if len(services) != 2 {
t.Fatalf("expected 2 services, got %d", len(services))
}