代理批量关联租户接口、关于页运行时与资源指标

internal/service/proxy.go

@@ -229,6 +229,63 @@ func (s *ProxyService) Delete(ctx context.Context, id uint) error {
 	return nil
 }
 
+// SetTenants 设置关联此代理的租户全集:列表内的租户改挂到本代理(含从
+// 其他代理改挂),列表外已关联的解除;整体事务生效,返回新的关联计数。
+func (s *ProxyService) SetTenants(ctx context.Context, id uint, cfgIDs []uint) (int64, error) {
+	if err := s.db.WithContext(ctx).First(&model.Proxy{}, id).Error; err != nil {
+		return 0, fmt.Errorf("find proxy %d: %w", id, err)
+	}
+	ids, err := s.validTenantIDs(ctx, cfgIDs)
+	if err != nil {
+		return 0, err
+	}
+	err = s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+		unbind := tx.Model(&model.OciConfig{}).Where("proxy_id = ?", id)
+		if len(ids) > 0 {
+			unbind = unbind.Where("id NOT IN ?", ids)
+		}
+		if err := unbind.Update("proxy_id", nil).Error; err != nil {
+			return fmt.Errorf("unbind tenants: %w", err)
+		}
+		if len(ids) == 0 {
+			return nil
+		}
+		err := tx.Model(&model.OciConfig{}).Where("id IN ?", ids).Update("proxy_id", id).Error
+		if err != nil {
+			return fmt.Errorf("bind tenants: %w", err)
+		}
+		return nil
+	})
+	if err != nil {
+		return 0, err
+	}
+	return int64(len(ids)), nil
+}
+
+// validTenantIDs 去重并校验租户 ID 均存在;含无效 ID 时报 ErrProxyInvalid。
+func (s *ProxyService) validTenantIDs(ctx context.Context, cfgIDs []uint) ([]uint, error) {
+	seen := map[uint]struct{}{}
+	ids := make([]uint, 0, len(cfgIDs))
+	for _, v := range cfgIDs {
+		if _, ok := seen[v]; !ok {
+			seen[v] = struct{}{}
+			ids = append(ids, v)
+		}
+	}
+	if len(ids) == 0 {
+		return ids, nil
+	}
+	var n int64
+	err := s.db.WithContext(ctx).Model(&model.OciConfig{}).Where("id IN ?", ids).Count(&n).Error
+	if err != nil {
+		return nil, fmt.Errorf("count tenants: %w", err)
+	}
+	if n != int64(len(ids)) {
+		return nil, fmt.Errorf("存在无效的租户 ID: %w", ErrProxyInvalid)
+	}
+	return ids, nil
+}
+
 // SpecOf 解密并组装 SDK 用的代理参数;id 为 nil 返回 nil(直连)。
 func (s *ProxyService) SpecOf(ctx context.Context, id *uint) (*oci.ProxySpec, error) {
 	if id == nil {

internal/service/proxy_test.go

@@ -2,6 +2,7 @@ package service
 
 import (
 	"context"
+	"errors"
 	"net/http"
 	"net/http/httptest"
 	"strings"
@@ -227,3 +228,99 @@ func TestProxyGeoProbe(t *testing.T) {
 		t.Fatalf("probe view = %+v, want 空地区带探测时间", view)
 	}
 }
+
+func TestProxySetTenants(t *testing.T) {
+	svc, db := newProxyEnv(t)
+	ctx := context.Background()
+
+	pa, _ := svc.Create(ctx, ProxyInput{Type: "socks5", Host: "10.0.0.1", Port: 1080})
+	pb, _ := svc.Create(ctx, ProxyInput{Type: "socks5", Host: "10.0.0.2", Port: 1080})
+	// t1 已挂 A,t2 挂 B,t3 直连
+	t1 := model.OciConfig{Alias: "t1", ProxyID: &pa.ID}
+	t2 := model.OciConfig{Alias: "t2", ProxyID: &pb.ID}
+	t3 := model.OciConfig{Alias: "t3"}
+	for _, cfg := range []*model.OciConfig{&t1, &t2, &t3} {
+		if err := db.Create(cfg).Error; err != nil {
+			t.Fatalf("seed config: %v", err)
+		}
+	}
+
+	proxyOf := func(id uint) *uint {
+		var row model.OciConfig
+		if err := db.First(&row, id).Error; err != nil {
+			t.Fatalf("load config %d: %v", id, err)
+		}
+		return row.ProxyID
+	}
+
+	cases := []struct {
+		name    string
+		proxyID uint
+		cfgIDs  []uint
+		wantN   int64
+		wantErr error
+	}{
+		{name: "改挂并新增", proxyID: pa.ID, cfgIDs: []uint{t2.ID, t3.ID, t3.ID}, wantN: 2},
+		{name: "清空解除全部", proxyID: pa.ID, cfgIDs: nil, wantN: 0},
+		{name: "无效租户拒绝", proxyID: pa.ID, cfgIDs: []uint{9999}, wantErr: ErrProxyInvalid},
+		{name: "代理不存在", proxyID: 9999, cfgIDs: nil, wantErr: gorm.ErrRecordNotFound},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			n, err := svc.SetTenants(ctx, tc.proxyID, tc.cfgIDs)
+			if tc.wantErr != nil {
+				if !errors.Is(err, tc.wantErr) {
+					t.Fatalf("err = %v, want %v", err, tc.wantErr)
+				}
+				return
+			}
+			if err != nil || n != tc.wantN {
+				t.Fatalf("SetTenants = %d, %v; want %d, nil", n, err, tc.wantN)
+			}
+		})
+	}
+
+	// 终态:清空后三租户全部直连(t1 在第一步已被解除,t2/t3 在第二步解除)
+	for _, id := range []uint{t1.ID, t2.ID, t3.ID} {
+		if got := proxyOf(id); got != nil {
+			t.Fatalf("config %d proxyID = %v, want nil", id, *got)
+		}
+	}
+}
+
+func TestProxySetTenantsRebind(t *testing.T) {
+	svc, db := newProxyEnv(t)
+	ctx := context.Background()
+
+	pa, _ := svc.Create(ctx, ProxyInput{Type: "socks5", Host: "10.0.1.1", Port: 1080})
+	pb, _ := svc.Create(ctx, ProxyInput{Type: "socks5", Host: "10.0.1.2", Port: 1080})
+	cfg := model.OciConfig{Alias: "steal", ProxyID: &pa.ID}
+	if err := db.Create(&cfg).Error; err != nil {
+		t.Fatalf("seed config: %v", err)
+	}
+
+	if _, err := svc.SetTenants(ctx, pb.ID, []uint{cfg.ID}); err != nil {
+		t.Fatalf("SetTenants: %v", err)
+	}
+	var row model.OciConfig
+	if err := db.First(&row, cfg.ID).Error; err != nil {
+		t.Fatalf("load config: %v", err)
+	}
+	if row.ProxyID == nil || *row.ProxyID != pb.ID {
+		t.Fatalf("proxyID = %v, want %d(改挂到 B)", row.ProxyID, pb.ID)
+	}
+	// A 侧计数应归零,B 侧为 1
+	views, err := svc.List(ctx)
+	if err != nil {
+		t.Fatalf("List: %v", err)
+	}
+	for _, v := range views {
+		want := int64(0)
+		if v.ID == pb.ID {
+			want = 1
+		}
+		if v.UsedBy != want {
+			t.Fatalf("proxy %d UsedBy = %d, want %d", v.ID, v.UsedBy, want)
+		}
+	}
+}