294 lines
11 KiB
Go
294 lines
11 KiB
Go
package polling
|
|
|
|
import (
|
|
"context"
|
|
"math"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/redis/go-redis/v9"
|
|
"gorm.io/gorm"
|
|
|
|
auditinfra "github.com/break/junhong_cmp_fiber/internal/infrastructure/audit"
|
|
"github.com/break/junhong_cmp_fiber/internal/model"
|
|
"github.com/break/junhong_cmp_fiber/internal/store/postgres"
|
|
"github.com/break/junhong_cmp_fiber/pkg/constants"
|
|
"github.com/break/junhong_cmp_fiber/pkg/errors"
|
|
"github.com/break/junhong_cmp_fiber/pkg/middleware"
|
|
)
|
|
|
|
// ConcurrencyService 并发控制服务
|
|
type ConcurrencyService struct {
|
|
store *postgres.PollingConcurrencyConfigStore
|
|
db *gorm.DB
|
|
auditWriter *auditinfra.Writer
|
|
redis *redis.Client
|
|
}
|
|
|
|
// SetAudit 注入轮询并发配置事务与统一审计 Writer。
|
|
func (s *ConcurrencyService) SetAudit(db *gorm.DB, writer *auditinfra.Writer) {
|
|
s.db = db
|
|
s.auditWriter = writer
|
|
}
|
|
|
|
// NewConcurrencyService 创建并发控制服务实例
|
|
func NewConcurrencyService(store *postgres.PollingConcurrencyConfigStore, redis *redis.Client) *ConcurrencyService {
|
|
return &ConcurrencyService{
|
|
store: store,
|
|
redis: redis,
|
|
}
|
|
}
|
|
|
|
// ConcurrencyStatus 并发状态
|
|
type ConcurrencyStatus struct {
|
|
TaskType string `json:"task_type"`
|
|
TaskTypeName string `json:"task_type_name"`
|
|
MaxConcurrency int `json:"max_concurrency"`
|
|
Current int64 `json:"current"`
|
|
Available int64 `json:"available"`
|
|
Utilization float64 `json:"utilization"`
|
|
}
|
|
|
|
// List 获取所有并发控制配置及当前状态
|
|
func (s *ConcurrencyService) List(ctx context.Context) ([]*ConcurrencyStatus, error) {
|
|
configs, err := s.store.List(ctx)
|
|
if err != nil {
|
|
return nil, errors.Wrap(errors.CodeInternalError, err, "获取并发配置列表失败")
|
|
}
|
|
|
|
result := make([]*ConcurrencyStatus, 0, len(configs))
|
|
for _, cfg := range configs {
|
|
status := &ConcurrencyStatus{
|
|
TaskType: cfg.TaskType,
|
|
TaskTypeName: s.getTaskTypeName(cfg.TaskType),
|
|
MaxConcurrency: cfg.MaxConcurrency,
|
|
}
|
|
|
|
// 从 Redis 获取当前并发数
|
|
currentKey := pollingConcurrencyCurrentKey(cfg.TaskType)
|
|
current, err := s.redis.Get(ctx, currentKey).Int64()
|
|
if err != nil && err != redis.Nil {
|
|
current = 0
|
|
}
|
|
if current < 0 {
|
|
current = 0
|
|
}
|
|
|
|
status.Current = current
|
|
status.Available = int64(cfg.MaxConcurrency) - current
|
|
if status.Available < 0 {
|
|
status.Available = 0
|
|
}
|
|
if cfg.MaxConcurrency > 0 {
|
|
status.Utilization = math.Round(float64(current)/float64(cfg.MaxConcurrency)*1000000) / 100
|
|
}
|
|
|
|
result = append(result, status)
|
|
}
|
|
|
|
return result, nil
|
|
}
|
|
|
|
// GetByTaskType 根据任务类型获取并发配置及状态
|
|
func (s *ConcurrencyService) GetByTaskType(ctx context.Context, taskType string) (*ConcurrencyStatus, error) {
|
|
cfg, err := s.store.GetByTaskType(ctx, taskType)
|
|
if err != nil {
|
|
return nil, errors.Wrap(errors.CodeNotFound, err, "并发配置不存在")
|
|
}
|
|
|
|
status := &ConcurrencyStatus{
|
|
TaskType: cfg.TaskType,
|
|
TaskTypeName: s.getTaskTypeName(cfg.TaskType),
|
|
MaxConcurrency: cfg.MaxConcurrency,
|
|
}
|
|
|
|
// 从 Redis 获取当前并发数
|
|
currentKey := pollingConcurrencyCurrentKey(cfg.TaskType)
|
|
current, err := s.redis.Get(ctx, currentKey).Int64()
|
|
if err != nil && err != redis.Nil {
|
|
current = 0
|
|
}
|
|
if current < 0 {
|
|
current = 0
|
|
}
|
|
|
|
status.Current = current
|
|
status.Available = int64(cfg.MaxConcurrency) - current
|
|
if status.Available < 0 {
|
|
status.Available = 0
|
|
}
|
|
if cfg.MaxConcurrency > 0 {
|
|
status.Utilization = math.Round(float64(current)/float64(cfg.MaxConcurrency)*1000000) / 100
|
|
}
|
|
|
|
return status, nil
|
|
}
|
|
|
|
// UpdateMaxConcurrency 更新最大并发数
|
|
func (s *ConcurrencyService) UpdateMaxConcurrency(ctx context.Context, taskType string, maxConcurrency int, updatedBy uint) error {
|
|
// 验证参数
|
|
if maxConcurrency < 1 || maxConcurrency > constants.PollingMaxConcurrencyLimit {
|
|
return errors.New(errors.CodeInvalidParam, "并发数必须为 1-1000")
|
|
}
|
|
|
|
// 验证任务类型存在
|
|
config, err := s.store.GetByTaskType(ctx, taskType)
|
|
if err != nil {
|
|
return errors.Wrap(errors.CodeNotFound, err, "任务类型不存在")
|
|
}
|
|
|
|
before := config.MaxConcurrency
|
|
err = runPollingTransaction(ctx, s.db, s.auditWriter, func(tx *gorm.DB) error {
|
|
if err := s.store.WithTx(tx).UpdateMaxConcurrency(ctx, taskType, maxConcurrency, updatedBy); err != nil {
|
|
return err
|
|
}
|
|
config.MaxConcurrency = maxConcurrency
|
|
return writePollingAudit(ctx, tx, s.auditWriter, auditinfra.PollingInput{
|
|
ActionCode: constants.AuditActionPollingConcurrencyUpdated, Summary: "更新轮询并发配置",
|
|
ResourceType: constants.AuditResourcePollingConcurrencyConfig, ResourceID: config.ID,
|
|
ResourceKey: config.TaskType, DisplayName: s.getTaskTypeName(config.TaskType), OperatorID: updatedBy,
|
|
IdentitySnapshot: pollingConcurrencyIdentity(config),
|
|
BeforeData: map[string]any{"max_concurrency": before}, AfterData: map[string]any{"max_concurrency": maxConcurrency},
|
|
})
|
|
})
|
|
if err != nil {
|
|
appErr := errors.Wrap(errors.CodeInternalError, err, "更新并发配置失败")
|
|
recordPollingFailure(ctx, s.db, s.auditWriter, auditinfra.PollingInput{
|
|
ActionCode: constants.AuditActionPollingConcurrencyUpdated, Summary: "更新轮询并发配置失败",
|
|
ResourceType: constants.AuditResourcePollingConcurrencyConfig, ResourceID: config.ID,
|
|
ResourceKey: config.TaskType, DisplayName: s.getTaskTypeName(config.TaskType), OperatorID: updatedBy,
|
|
IdentitySnapshot: pollingConcurrencyIdentity(config), BeforeData: map[string]any{"max_concurrency": before},
|
|
}, appErr)
|
|
return appErr
|
|
}
|
|
|
|
// 同步更新 Redis 配置缓存
|
|
configKey := constants.RedisPollingConcurrencyConfigKey(taskType)
|
|
if err := s.redis.Set(ctx, configKey, maxConcurrency, 24*time.Hour).Err(); err != nil {
|
|
// Redis 更新失败不影响主流程，下次读取会从数据库重新加载
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// ResetConcurrency 重置并发计数（用于信号量修复）
|
|
func (s *ConcurrencyService) ResetConcurrency(ctx context.Context, taskType string) error {
|
|
operatorID := middleware.GetUserIDFromContext(ctx)
|
|
if operatorID == 0 {
|
|
return errors.New(errors.CodeUnauthorized, "未授权访问")
|
|
}
|
|
// 验证任务类型存在
|
|
config, err := s.store.GetByTaskType(ctx, taskType)
|
|
if err != nil {
|
|
return errors.Wrap(errors.CodeNotFound, err, "任务类型不存在")
|
|
}
|
|
|
|
// 重置 Redis 当前计数为 0
|
|
currentKey := pollingConcurrencyCurrentKey(config.TaskType)
|
|
before, getErr := s.redis.Get(ctx, currentKey).Int64()
|
|
beforeExists := getErr == nil
|
|
if getErr != nil && getErr != redis.Nil {
|
|
appErr := errors.Wrap(errors.CodeInternalError, getErr, "读取并发计数失败")
|
|
recordPollingFailure(ctx, s.db, s.auditWriter, auditinfra.PollingInput{
|
|
ActionCode: constants.AuditActionPollingConcurrencyReset, Summary: "重置轮询并发计数失败",
|
|
ResourceType: constants.AuditResourcePollingConcurrencyConfig, ResourceID: config.ID,
|
|
ResourceKey: config.TaskType, DisplayName: s.getTaskTypeName(config.TaskType), OperatorID: operatorID,
|
|
IdentitySnapshot: pollingConcurrencyIdentity(config),
|
|
}, appErr)
|
|
return appErr
|
|
}
|
|
beforeTTL := time.Duration(0)
|
|
if beforeExists {
|
|
beforeTTL, _ = s.redis.PTTL(ctx, currentKey).Result()
|
|
if beforeTTL < 0 {
|
|
beforeTTL = 0
|
|
}
|
|
}
|
|
if err := s.redis.Set(ctx, currentKey, 0, 24*time.Hour).Err(); err != nil {
|
|
appErr := errors.Wrap(errors.CodeInternalError, err, "重置并发计数失败")
|
|
recordPollingFailure(ctx, s.db, s.auditWriter, auditinfra.PollingInput{
|
|
ActionCode: constants.AuditActionPollingConcurrencyReset, Summary: "重置轮询并发计数失败",
|
|
ResourceType: constants.AuditResourcePollingConcurrencyConfig, ResourceID: config.ID,
|
|
ResourceKey: config.TaskType, DisplayName: s.getTaskTypeName(config.TaskType), OperatorID: operatorID,
|
|
IdentitySnapshot: pollingConcurrencyIdentity(config), BeforeData: map[string]any{"current": before},
|
|
}, appErr)
|
|
return appErr
|
|
}
|
|
err = runPollingTransaction(ctx, s.db, s.auditWriter, func(tx *gorm.DB) error {
|
|
return writePollingAudit(ctx, tx, s.auditWriter, auditinfra.PollingInput{
|
|
ActionCode: constants.AuditActionPollingConcurrencyReset, Summary: "重置轮询并发计数",
|
|
ResourceType: constants.AuditResourcePollingConcurrencyConfig, ResourceID: config.ID,
|
|
ResourceKey: config.TaskType, DisplayName: s.getTaskTypeName(config.TaskType), OperatorID: operatorID,
|
|
IdentitySnapshot: pollingConcurrencyIdentity(config),
|
|
BeforeData: map[string]any{"current": before}, AfterData: map[string]any{"current": int64(0)},
|
|
})
|
|
})
|
|
if err == nil {
|
|
return nil
|
|
}
|
|
if beforeExists {
|
|
_ = s.redis.Set(ctx, currentKey, before, beforeTTL).Err()
|
|
} else {
|
|
_ = s.redis.Del(ctx, currentKey).Err()
|
|
}
|
|
appErr := errors.Wrap(errors.CodeInternalError, err, "记录重置并发计数审计失败")
|
|
recordPollingFailure(ctx, s.db, s.auditWriter, auditinfra.PollingInput{
|
|
ActionCode: constants.AuditActionPollingConcurrencyReset, Summary: "重置轮询并发计数失败",
|
|
ResourceType: constants.AuditResourcePollingConcurrencyConfig, ResourceID: config.ID,
|
|
ResourceKey: config.TaskType, DisplayName: s.getTaskTypeName(config.TaskType), OperatorID: operatorID,
|
|
IdentitySnapshot: pollingConcurrencyIdentity(config), BeforeData: map[string]any{"current": before},
|
|
}, appErr)
|
|
return appErr
|
|
}
|
|
|
|
// InitFromDB 从数据库初始化 Redis 并发配置
|
|
func (s *ConcurrencyService) InitFromDB(ctx context.Context) error {
|
|
configs, err := s.store.List(ctx)
|
|
if err != nil {
|
|
return errors.Wrap(errors.CodeInternalError, err, "获取并发配置失败")
|
|
}
|
|
|
|
for _, cfg := range configs {
|
|
configKey := constants.RedisPollingConcurrencyConfigKey(cfg.TaskType)
|
|
if err := s.redis.Set(ctx, configKey, cfg.MaxConcurrency, 24*time.Hour).Err(); err != nil {
|
|
// 忽略单个配置同步失败
|
|
continue
|
|
}
|
|
}
|
|
_ = s.redis.Del(ctx, constants.RedisPollingConcurrencyConfigKey("stop_start")).Err()
|
|
|
|
return nil
|
|
}
|
|
|
|
// SyncConfigToRedis 同步单个配置到 Redis
|
|
func (s *ConcurrencyService) SyncConfigToRedis(ctx context.Context, config *model.PollingConcurrencyConfig) error {
|
|
configKey := constants.RedisPollingConcurrencyConfigKey(config.TaskType)
|
|
return s.redis.Set(ctx, configKey, config.MaxConcurrency, 24*time.Hour).Err()
|
|
}
|
|
|
|
// pollingConcurrencyCurrentKey 将配置中的短任务类型转换为 Worker 使用的完整计数键。
|
|
func pollingConcurrencyCurrentKey(taskType string) string {
|
|
if !strings.HasPrefix(taskType, "polling:") {
|
|
taskType = "polling:" + taskType
|
|
}
|
|
return constants.RedisPollingConcurrencyCurrentKey(taskType)
|
|
}
|
|
|
|
// getTaskTypeName 获取任务类型的中文名称
|
|
func (s *ConcurrencyService) getTaskTypeName(taskType string) string {
|
|
switch taskType {
|
|
case "realname":
|
|
return "实名检查"
|
|
case "carddata":
|
|
return "流量检查"
|
|
case "package":
|
|
return "套餐检查"
|
|
case "protect":
|
|
return "保护期检查"
|
|
case "card_status":
|
|
return "卡状态检查"
|
|
default:
|
|
return taskType
|
|
}
|
|
}
|