398 lines
11 KiB
Go
398 lines
11 KiB
Go
package check
|
||
|
||
import (
|
||
consts "cid/consts/check"
|
||
dao "cid/dao/check"
|
||
entity "cid/model/entity/check"
|
||
"context"
|
||
"encoding/json"
|
||
"sync"
|
||
"time"
|
||
|
||
"gitea.redpowerfuture.com/red-future/common/beans"
|
||
"github.com/gogf/gf/v2/frame/g"
|
||
"github.com/gogf/gf/v2/os/grpool"
|
||
)
|
||
|
||
// ContentCheckConfig 送检配置
|
||
type ContentCheckConfig struct {
|
||
BatchSize int `json:"batch_size"`
|
||
ImageEnabled bool `json:"image_enabled"`
|
||
VideoEnabled bool `json:"video_enabled"`
|
||
IntervalSeconds int `json:"interval_seconds"`
|
||
PollInterval int `json:"poll_interval"` // 自动轮询检测结果间隔（秒）
|
||
}
|
||
|
||
// DefaultConfig 默认配置
|
||
var DefaultConfig = ContentCheckConfig{
|
||
BatchSize: 10,
|
||
ImageEnabled: true,
|
||
VideoEnabled: true,
|
||
IntervalSeconds: 30,
|
||
PollInterval: 60,
|
||
}
|
||
|
||
// TencentContentCheckService 腾讯内容送检服务
|
||
type TencentContentCheckService struct {
|
||
mu sync.RWMutex
|
||
config ContentCheckConfig
|
||
isRunning bool
|
||
pool *grpool.Pool
|
||
cancel context.CancelFunc
|
||
}
|
||
|
||
// TencentContentCheck 送检服务单例
|
||
var TencentContentCheck = &TencentContentCheckService{
|
||
config: DefaultConfig,
|
||
}
|
||
|
||
// SetConfig 设置配置
|
||
func (s *TencentContentCheckService) SetConfig(config ContentCheckConfig) {
|
||
s.mu.Lock()
|
||
s.config = config
|
||
s.mu.Unlock()
|
||
}
|
||
|
||
// Start 启动定时任务
|
||
func (s *TencentContentCheckService) Start(ctx context.Context) error {
|
||
s.mu.Lock()
|
||
if s.isRunning {
|
||
s.mu.Unlock()
|
||
g.Log().Info(ctx, "送检服务已在运行中，跳过启动")
|
||
return nil
|
||
}
|
||
|
||
s.isRunning = true
|
||
config := s.config
|
||
s.pool = grpool.New(5)
|
||
s.mu.Unlock()
|
||
|
||
g.Log().Infof(ctx, "启动内容送检服务，配置: batch_size=%d, interval=%ds, poll=%ds, image=%v, video=%v",
|
||
config.BatchSize, config.IntervalSeconds, config.PollInterval, config.ImageEnabled, config.VideoEnabled)
|
||
|
||
schedCtx, cancel := context.WithCancel(context.Background())
|
||
s.cancel = cancel
|
||
|
||
// 定时送检协程
|
||
g.Go(schedCtx, func(ctx context.Context) {
|
||
ticker := time.NewTicker(time.Duration(config.IntervalSeconds) * time.Second)
|
||
defer ticker.Stop()
|
||
|
||
// 启动时先执行一次
|
||
_ = s.pool.Add(ctx, func(jobCtx context.Context) {
|
||
s.processAll(jobCtx)
|
||
})
|
||
|
||
for {
|
||
select {
|
||
case <-ticker.C:
|
||
_ = s.pool.Add(ctx, func(jobCtx context.Context) {
|
||
s.processAll(jobCtx)
|
||
})
|
||
case <-ctx.Done():
|
||
s.pool.Close()
|
||
return
|
||
}
|
||
}
|
||
}, nil)
|
||
|
||
// 自动轮询检测结果协程（无论回调/轮询模式，都定期查一次未处理的结果）
|
||
pollInterval := config.PollInterval
|
||
if pollInterval <= 0 {
|
||
pollInterval = 60
|
||
}
|
||
g.Go(schedCtx, func(ctx context.Context) {
|
||
pollTicker := time.NewTicker(time.Duration(pollInterval) * time.Second)
|
||
defer pollTicker.Stop()
|
||
|
||
g.Log().Infof(ctx, "启动自动轮询检测结果, 间隔=%ds", pollInterval)
|
||
|
||
for {
|
||
select {
|
||
case <-pollTicker.C:
|
||
_, _, _ = MaterialVerify.PollPendingResults(ctx)
|
||
case <-ctx.Done():
|
||
return
|
||
}
|
||
}
|
||
}, nil)
|
||
|
||
return nil
|
||
}
|
||
|
||
// Stop 停止定时任务
|
||
func (s *TencentContentCheckService) Stop(ctx context.Context) {
|
||
s.mu.Lock()
|
||
if !s.isRunning {
|
||
s.mu.Unlock()
|
||
return
|
||
}
|
||
s.isRunning = false
|
||
s.mu.Unlock()
|
||
|
||
if s.cancel != nil {
|
||
s.cancel()
|
||
}
|
||
g.Log().Info(ctx, "停止内容送检服务")
|
||
}
|
||
|
||
// processAll 处理所有待送检数据
|
||
func (s *TencentContentCheckService) processAll(ctx context.Context) {
|
||
ctx = context.WithValue(ctx, "user", &beans.User{UserName: "system", TenantId: 1})
|
||
|
||
startTime := time.Now()
|
||
g.Log().Info(ctx, "开始处理待送检数据...")
|
||
|
||
var totalProcessed int
|
||
|
||
s.mu.RLock()
|
||
imageEnabled := s.config.ImageEnabled
|
||
videoEnabled := s.config.VideoEnabled
|
||
s.mu.RUnlock()
|
||
|
||
if imageEnabled {
|
||
imageCount, err := dao.TencentImage.CountPending(ctx)
|
||
if err != nil {
|
||
g.Log().Errorf(ctx, "统计待送检图片数量失败: %v", err)
|
||
} else if imageCount > 0 {
|
||
count, procErr := s.processImages(ctx)
|
||
if procErr != nil {
|
||
g.Log().Errorf(ctx, "图片送检处理失败: %v", procErr)
|
||
}
|
||
totalProcessed += count
|
||
}
|
||
}
|
||
|
||
if videoEnabled {
|
||
videoCount, err := dao.TencentVideo.CountPending(ctx)
|
||
if err != nil {
|
||
g.Log().Errorf(ctx, "统计待送检视频数量失败: %v", err)
|
||
} else if videoCount > 0 {
|
||
count, procErr := s.processVideos(ctx)
|
||
if procErr != nil {
|
||
g.Log().Errorf(ctx, "视频送检处理失败: %v", procErr)
|
||
}
|
||
totalProcessed += count
|
||
}
|
||
}
|
||
|
||
duration := time.Since(startTime).Milliseconds()
|
||
g.Log().Infof(ctx, "处理完成，共处理 %d 条数据，耗时 %dms", totalProcessed, duration)
|
||
}
|
||
|
||
// processImages 处理图片送检（统一走 MaterialVerify 系统）
|
||
func (s *TencentContentCheckService) processImages(ctx context.Context) (int, error) {
|
||
s.mu.RLock()
|
||
batchSize := s.config.BatchSize
|
||
s.mu.RUnlock()
|
||
|
||
images, err := dao.TencentImage.GetPendingList(ctx, batchSize)
|
||
if err != nil {
|
||
g.Log().Errorf(ctx, "获取待送检图片失败: %v", err)
|
||
return 0, err
|
||
}
|
||
|
||
if len(images) == 0 {
|
||
return 0, nil
|
||
}
|
||
|
||
g.Log().Infof(ctx, "开始送检 %d 张图片", len(images))
|
||
|
||
successCount := 0
|
||
failedCount := 0
|
||
|
||
for _, img := range images {
|
||
_, err := MaterialVerify.VerifyImageByID(ctx, img.ImageID)
|
||
if err != nil {
|
||
failedCount++
|
||
g.Log().Errorf(ctx, "图片送检失败, imageId=%s, error=%v", img.ImageID, err)
|
||
} else {
|
||
successCount++
|
||
}
|
||
|
||
time.Sleep(100 * time.Millisecond)
|
||
}
|
||
|
||
g.Log().Infof(ctx, "图片送检完成，成功: %d，失败: %d", successCount, failedCount)
|
||
return len(images), nil
|
||
}
|
||
|
||
// processVideos 处理视频送检（统一走 MaterialVerify 系统）
|
||
func (s *TencentContentCheckService) processVideos(ctx context.Context) (int, error) {
|
||
s.mu.RLock()
|
||
batchSize := s.config.BatchSize
|
||
s.mu.RUnlock()
|
||
|
||
videos, err := dao.TencentVideo.GetPendingList(ctx, batchSize)
|
||
if err != nil {
|
||
g.Log().Errorf(ctx, "获取待送检视频失败: %v", err)
|
||
return 0, err
|
||
}
|
||
|
||
if len(videos) == 0 {
|
||
return 0, nil
|
||
}
|
||
|
||
g.Log().Infof(ctx, "开始送检 %d 个视频", len(videos))
|
||
|
||
successCount := 0
|
||
failedCount := 0
|
||
|
||
for _, video := range videos {
|
||
_, err := MaterialVerify.VerifyVideoByID(ctx, video.VideoID)
|
||
if err != nil {
|
||
failedCount++
|
||
g.Log().Errorf(ctx, "视频送检失败, videoId=%s, error=%v", video.VideoID, err)
|
||
} else {
|
||
successCount++
|
||
}
|
||
|
||
time.Sleep(100 * time.Millisecond)
|
||
}
|
||
|
||
g.Log().Infof(ctx, "视频送检完成，成功: %d，失败: %d", successCount, failedCount)
|
||
return len(videos), nil
|
||
}
|
||
|
||
// writeAuditLog 写入审计日志（tencent_content_check_log）
|
||
// 当 suggestion<0 时记录为 SUBMITTING（已提交等待结果），否则记录为 COMPLETED（检测完成）
|
||
func (s *TencentContentCheckService) writeAuditLog(ctx context.Context, sourceTable string, sourceID int64, mediaID string, mediaURL string, taskID string, suggestion, label, resultType int, responseData string, duration int64) {
|
||
requestParam := map[string]interface{}{
|
||
"media_id": mediaID,
|
||
"url": mediaURL,
|
||
}
|
||
requestParamJSON, _ := json.Marshal(requestParam)
|
||
|
||
status := consts.CheckStatusSubmitting
|
||
if suggestion >= 0 {
|
||
status = consts.CheckStatusCompleted
|
||
}
|
||
|
||
log := &entity.TencentContentCheckLog{
|
||
SourceTable: sourceTable,
|
||
SourceID: sourceID,
|
||
RequestURL: "易盾内容安全检测接口",
|
||
RequestParam: string(requestParamJSON),
|
||
Status: status,
|
||
CheckTime: time.Now().UnixMilli(),
|
||
TaskID: taskID,
|
||
Suggestion: suggestion,
|
||
Label: label,
|
||
ResultType: resultType,
|
||
ResponseData: responseData,
|
||
Duration: duration,
|
||
}
|
||
|
||
id, err := dao.TencentContentCheckLog.Create(ctx, log)
|
||
if err != nil {
|
||
g.Log().Errorf(ctx, "创建送检审计日志失败: %v", err)
|
||
return
|
||
}
|
||
|
||
g.Log().Debugf(ctx, "创建送检审计日志成功, id=%d, sourceTable=%s, sourceID=%d, taskId=%s", id, sourceTable, sourceID, taskID)
|
||
}
|
||
|
||
// SubmitImageByID 根据图片ID手动提交送检（统一走 MaterialVerify 系统）
|
||
func (s *TencentContentCheckService) SubmitImageByID(ctx context.Context, imageID string) (*ImageSubmitResult, error) {
|
||
mLog, err := MaterialVerify.VerifyImageByID(ctx, imageID)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
image, err := dao.TencentImage.GetByImageID(ctx, imageID)
|
||
if err == nil && image != nil {
|
||
s.writeAuditLog(ctx, consts.SourceTableTencentImage, image.Id, imageID, image.PreviewURL, mLog.TaskID, -1, 0, 0, "", 0)
|
||
}
|
||
|
||
return &ImageSubmitResult{
|
||
TaskID: mLog.TaskID,
|
||
}, nil
|
||
}
|
||
|
||
// SubmitVideoByID 根据视频ID手动提交送检（统一走 MaterialVerify 系统）
|
||
func (s *TencentContentCheckService) SubmitVideoByID(ctx context.Context, videoID string) (*VideoSubmitResult, error) {
|
||
mLog, err := MaterialVerify.VerifyVideoByID(ctx, videoID)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
video, err := dao.TencentVideo.GetByVideoID(ctx, videoID)
|
||
if err == nil && video != nil {
|
||
s.writeAuditLog(ctx, consts.SourceTableTencentVideo, video.Id, videoID, video.PreviewURL, mLog.TaskID, -1, 0, 0, "", 0)
|
||
}
|
||
|
||
return &VideoSubmitResult{
|
||
TaskID: mLog.TaskID,
|
||
}, nil
|
||
}
|
||
|
||
// GetPendingStats 获取待送检统计
|
||
func (s *TencentContentCheckService) GetPendingStats(ctx context.Context) map[string]int {
|
||
stats := make(map[string]int)
|
||
|
||
s.mu.RLock()
|
||
imageEnabled := s.config.ImageEnabled
|
||
videoEnabled := s.config.VideoEnabled
|
||
s.mu.RUnlock()
|
||
|
||
if imageEnabled {
|
||
count, err := dao.TencentImage.CountPending(ctx)
|
||
if err != nil {
|
||
g.Log().Errorf(ctx, "统计待送检图片数量失败: %v", err)
|
||
}
|
||
stats["image_pending"] = count
|
||
}
|
||
|
||
if videoEnabled {
|
||
count, err := dao.TencentVideo.CountPending(ctx)
|
||
if err != nil {
|
||
g.Log().Errorf(ctx, "统计待送检视频数量失败: %v", err)
|
||
}
|
||
stats["video_pending"] = count
|
||
}
|
||
|
||
return stats
|
||
}
|
||
|
||
// IsRunning 获取运行状态
|
||
func (s *TencentContentCheckService) IsRunning() bool {
|
||
s.mu.RLock()
|
||
defer s.mu.RUnlock()
|
||
return s.isRunning
|
||
}
|
||
|
||
// GetConfig 获取当前配置
|
||
func (s *TencentContentCheckService) GetConfig() ContentCheckConfig {
|
||
s.mu.RLock()
|
||
defer s.mu.RUnlock()
|
||
return s.config
|
||
}
|
||
|
||
// GetImageSubmitResult 获取图片提交结果（从送检日志解析）
|
||
func (s *TencentContentCheckService) GetImageSubmitResult(ctx context.Context, id int64) (*ImageSubmitResult, error) {
|
||
log, err := dao.TencentContentCheckLog.GetByID(ctx, id)
|
||
if err != nil || log == nil {
|
||
return nil, err
|
||
}
|
||
|
||
var result ImageSubmitResult
|
||
if err := json.Unmarshal([]byte(log.ResponseData), &result); err != nil {
|
||
return nil, err
|
||
}
|
||
return &result, nil
|
||
}
|
||
|
||
// GetVideoSubmitResult 获取视频提交结果（从送检日志解析）
|
||
func (s *TencentContentCheckService) GetVideoSubmitResult(ctx context.Context, id int64) (*VideoSubmitResult, error) {
|
||
log, err := dao.TencentContentCheckLog.GetByID(ctx, id)
|
||
if err != nil || log == nil {
|
||
return nil, err
|
||
}
|
||
|
||
var result VideoSubmitResult
|
||
if err := json.Unmarshal([]byte(log.ResponseData), &result); err != nil {
|
||
return nil, err
|
||
}
|
||
return &result, nil
|
||
}
|