feat: add channel async polling delay toggle

Fixes #5717 Fixes #4244
2026-06-24 20:40:10 +08:00
parent 5d943281b5
commit d2dcbc3137
13 changed files with 424 additions and 15 deletions
@@ -8,6 +8,7 @@ import (
 	"net/http"
 	"sort"
 	"strings"
+	"sync"
 	"time"

 	"github.com/QuantumNous/new-api/common"
@@ -18,6 +19,7 @@ import (
 	"github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
 	relaycommon "github.com/QuantumNous/new-api/relay/common"

+	"github.com/bytedance/gopkg/util/gopool"
 	"github.com/samber/lo"
 )

@@ -338,19 +340,40 @@ func taskNeedsUpdate(oldTask *model.Task, newTask dto.SunoDataResponse) bool {

 // UpdateVideoTasks 按渠道更新所有视频任务
 func UpdateVideoTasks(ctx context.Context, platform constant.TaskPlatform, taskChannelM map[int][]string, taskM map[string]*model.Task) error {
-	for channelId, taskIds := range taskChannelM {
-		if ctx.Err() != nil {
-			return ctx.Err()
-		}
-		if err := updateVideoTasks(ctx, platform, channelId, taskIds, taskM); err != nil {
-			logger.LogError(ctx, fmt.Sprintf("Channel #%d failed to update video async tasks: %s", channelId, err.Error()))
+	channelIDs := make([]int, 0, len(taskChannelM))
+	for channelID := range taskChannelM {
+		channelIDs = append(channelIDs, channelID)
+	}
+	sort.Ints(channelIDs)
+
+	var wg sync.WaitGroup
+	for _, channelId := range channelIDs {
+		taskIds := taskChannelM[channelId]
+		if len(taskIds) == 0 {
+			continue
 		}
+		taskIds = append([]string(nil), taskIds...)
+
+		wg.Add(1)
+		gopool.Go(func() {
+			defer wg.Done()
+			if err := updateVideoTasks(ctx, platform, channelId, taskIds, taskM); err != nil {
+				logger.LogError(ctx, fmt.Sprintf("Channel #%d failed to update video async tasks: %s", channelId, err.Error()))
+			}
+		})
+	}
+	wg.Wait()
+	if ctx.Err() != nil {
+		return ctx.Err()
 	}
 	return nil
 }

 func updateVideoTasks(ctx context.Context, platform constant.TaskPlatform, channelId int, taskIds []string, taskM map[string]*model.Task) error {
 	logger.LogInfo(ctx, fmt.Sprintf("Channel #%d pending video tasks: %d", channelId, len(taskIds)))
+	if ctx.Err() != nil {
+		return ctx.Err()
+	}
 	if len(taskIds) == 0 {
 		return nil
 	}
@@ -383,14 +406,19 @@ func updateVideoTasks(ctx context.Context, platform constant.TaskPlatform, chann
 	}
 	info.ApiKey = cacheGetChannel.Key
 	adaptor.Init(info)
-	for _, taskId := range taskIds {
+	disablePollingSleep := cacheGetChannel.GetOtherSettings().DisableTaskPollingSleep
+	for i, taskId := range taskIds {
 		if ctx.Err() != nil {
 			return ctx.Err()
 		}
 		if err := updateVideoSingleTask(ctx, adaptor, cacheGetChannel, taskId, taskM); err != nil {
 			logger.LogError(ctx, fmt.Sprintf("Failed to update video task %s: %s", taskId, err.Error()))
 		}
-		// sleep 1 second between each task to avoid hitting rate limits of upstream platforms
+		if disablePollingSleep || i == len(taskIds)-1 {
+			continue
+		}
+
+		// sleep 1 second between tasks for this channel only.
 		select {
 		case <-ctx.Done():
 			return ctx.Err()
@@ -0,0 +1,333 @@
+package service
+
+import (
+	"bytes"
+	"context"
+	"io"
+	"net/http"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/QuantumNous/new-api/common"
+	"github.com/QuantumNous/new-api/constant"
+	"github.com/QuantumNous/new-api/dto"
+	"github.com/QuantumNous/new-api/model"
+	relaycommon "github.com/QuantumNous/new-api/relay/common"
+	"github.com/bytedance/gopkg/util/gopool"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+type taskPollingFetchAdaptor struct {
+	mu           sync.Mutex
+	taskIDs      []string
+	fetched      chan string
+	blockTaskID  string
+	blockStarted chan struct{}
+	releaseBlock chan struct{}
+	blockOnce    sync.Once
+}
+
+func (a *taskPollingFetchAdaptor) Init(_ *relaycommon.RelayInfo) {}
+
+func (a *taskPollingFetchAdaptor) FetchTask(_ string, _ string, body map[string]any, _ string) (*http.Response, error) {
+	taskID, _ := body["task_id"].(string)
+	if taskID == a.blockTaskID && a.releaseBlock != nil {
+		a.blockOnce.Do(func() {
+			if a.blockStarted != nil {
+				close(a.blockStarted)
+			}
+		})
+		<-a.releaseBlock
+	}
+
+	a.mu.Lock()
+	a.taskIDs = append(a.taskIDs, taskID)
+	a.mu.Unlock()
+	if a.fetched != nil {
+		select {
+		case a.fetched <- taskID:
+		default:
+		}
+	}
+
+	response := dto.TaskResponse[model.Task]{
+		Code: dto.TaskSuccessCode,
+		Data: model.Task{
+			TaskID:   taskID,
+			Status:   model.TaskStatusInProgress,
+			Progress: "30%",
+		},
+	}
+	responseBody, err := common.Marshal(response)
+	if err != nil {
+		return nil, err
+	}
+	return &http.Response{
+		StatusCode: http.StatusOK,
+		Body:       io.NopCloser(bytes.NewReader(responseBody)),
+	}, nil
+}
+
+func (a *taskPollingFetchAdaptor) ParseTaskResult([]byte) (*relaycommon.TaskInfo, error) {
+	return &relaycommon.TaskInfo{Status: model.TaskStatusInProgress}, nil
+}
+
+func (a *taskPollingFetchAdaptor) AdjustBillingOnComplete(_ *model.Task, _ *relaycommon.TaskInfo) int {
+	return 0
+}
+
+func (a *taskPollingFetchAdaptor) fetchCount() int {
+	a.mu.Lock()
+	defer a.mu.Unlock()
+	return len(a.taskIDs)
+}
+
+func (a *taskPollingFetchAdaptor) fetchedTaskIDs() []string {
+	a.mu.Lock()
+	defer a.mu.Unlock()
+	return append([]string(nil), a.taskIDs...)
+}
+
+func seedTaskPollingChannel(t *testing.T, id int, disableSleep bool) {
+	t.Helper()
+	ch := &model.Channel{
+		Id:     id,
+		Type:   constant.ChannelTypeKling,
+		Name:   "polling_channel",
+		Key:    "sk-test",
+		Status: common.ChannelStatusEnabled,
+	}
+	if disableSleep {
+		ch.SetOtherSettings(dto.ChannelOtherSettings{DisableTaskPollingSleep: true})
+	}
+	require.NoError(t, model.DB.Create(ch).Error)
+}
+
+func seedPollingTask(t *testing.T, channelID int, publicID string, upstreamID string) *model.Task {
+	t.Helper()
+	task := &model.Task{
+		TaskID:    publicID,
+		Platform:  constant.TaskPlatform("kling"),
+		UserId:    1,
+		ChannelId: channelID,
+		Action:    constant.TaskActionGenerate,
+		Status:    model.TaskStatusInProgress,
+		Progress:  "30%",
+		CreatedAt: time.Now().Unix(),
+		UpdatedAt: time.Now().Unix(),
+		PrivateData: model.TaskPrivateData{
+			UpstreamTaskID: upstreamID,
+		},
+	}
+	require.NoError(t, model.DB.Create(task).Error)
+	return task
+}
+
+func TestUpdateVideoTasksDefaultSleepWaitsBetweenTasks(t *testing.T) {
+	truncate(t)
+
+	const channelID = 101
+	seedTaskPollingChannel(t, channelID, false)
+	first := seedPollingTask(t, channelID, "task_public_1", "upstream_1")
+	second := seedPollingTask(t, channelID, "task_public_2", "upstream_2")
+
+	adaptor := &taskPollingFetchAdaptor{}
+	previousFactory := GetTaskAdaptorFunc
+	GetTaskAdaptorFunc = func(constant.TaskPlatform) TaskPollingAdaptor { return adaptor }
+	t.Cleanup(func() { GetTaskAdaptorFunc = previousFactory })
+
+	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+	defer cancel()
+
+	err := UpdateVideoTasks(ctx, constant.TaskPlatform("kling"), map[int][]string{
+		channelID: {
+			first.GetUpstreamTaskID(),
+			second.GetUpstreamTaskID(),
+		},
+	}, map[string]*model.Task{
+		first.GetUpstreamTaskID():  first,
+		second.GetUpstreamTaskID(): second,
+	})
+
+	require.ErrorIs(t, err, context.DeadlineExceeded)
+	assert.Equal(t, 1, adaptor.fetchCount())
+}
+
+func TestUpdateVideoTasksCanSkipPollingSleepPerChannel(t *testing.T) {
+	truncate(t)
+
+	const channelID = 102
+	seedTaskPollingChannel(t, channelID, true)
+	first := seedPollingTask(t, channelID, "task_public_3", "upstream_3")
+	second := seedPollingTask(t, channelID, "task_public_4", "upstream_4")
+
+	adaptor := &taskPollingFetchAdaptor{}
+	previousFactory := GetTaskAdaptorFunc
+	GetTaskAdaptorFunc = func(constant.TaskPlatform) TaskPollingAdaptor { return adaptor }
+	t.Cleanup(func() { GetTaskAdaptorFunc = previousFactory })
+
+	ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
+	defer cancel()
+
+	err := UpdateVideoTasks(ctx, constant.TaskPlatform("kling"), map[int][]string{
+		channelID: {
+			first.GetUpstreamTaskID(),
+			second.GetUpstreamTaskID(),
+		},
+	}, map[string]*model.Task{
+		first.GetUpstreamTaskID():  first,
+		second.GetUpstreamTaskID(): second,
+	})
+
+	require.NoError(t, err)
+	assert.Equal(t, 2, adaptor.fetchCount())
+}
+
+func TestUpdateVideoTasksDefaultSleepDoesNotBlockOtherChannels(t *testing.T) {
+	truncate(t)
+
+	const firstChannelID = 201
+	const secondChannelID = 202
+	seedTaskPollingChannel(t, firstChannelID, false)
+	seedTaskPollingChannel(t, secondChannelID, false)
+	firstChannelFirst := seedPollingTask(t, firstChannelID, "task_public_5", "upstream_a_1")
+	firstChannelSecond := seedPollingTask(t, firstChannelID, "task_public_6", "upstream_a_2")
+	secondChannelFirst := seedPollingTask(t, secondChannelID, "task_public_7", "upstream_b_1")
+	secondChannelSecond := seedPollingTask(t, secondChannelID, "task_public_8", "upstream_b_2")
+
+	adaptor := &taskPollingFetchAdaptor{}
+	previousFactory := GetTaskAdaptorFunc
+	GetTaskAdaptorFunc = func(constant.TaskPlatform) TaskPollingAdaptor { return adaptor }
+	t.Cleanup(func() { GetTaskAdaptorFunc = previousFactory })
+
+	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
+	defer cancel()
+
+	err := UpdateVideoTasks(ctx, constant.TaskPlatform("kling"), map[int][]string{
+		firstChannelID: {
+			firstChannelFirst.GetUpstreamTaskID(),
+			firstChannelSecond.GetUpstreamTaskID(),
+		},
+		secondChannelID: {
+			secondChannelFirst.GetUpstreamTaskID(),
+			secondChannelSecond.GetUpstreamTaskID(),
+		},
+	}, map[string]*model.Task{
+		firstChannelFirst.GetUpstreamTaskID():   firstChannelFirst,
+		firstChannelSecond.GetUpstreamTaskID():  firstChannelSecond,
+		secondChannelFirst.GetUpstreamTaskID():  secondChannelFirst,
+		secondChannelSecond.GetUpstreamTaskID(): secondChannelSecond,
+	})
+
+	require.ErrorIs(t, err, context.DeadlineExceeded)
+	assert.ElementsMatch(t, []string{"upstream_a_1", "upstream_b_1"}, adaptor.fetchedTaskIDs())
+}
+
+func TestUpdateVideoTasksSlowChannelDoesNotBlockOtherChannels(t *testing.T) {
+	truncate(t)
+
+	const slowChannelID = 251
+	const fastChannelID = 252
+	seedTaskPollingChannel(t, slowChannelID, false)
+	seedTaskPollingChannel(t, fastChannelID, true)
+	slowTask := seedPollingTask(t, slowChannelID, "task_public_slow", "upstream_slow_1")
+	fastFirst := seedPollingTask(t, fastChannelID, "task_public_fast_1", "upstream_fast_parallel_1")
+	fastSecond := seedPollingTask(t, fastChannelID, "task_public_fast_2", "upstream_fast_parallel_2")
+
+	adaptor := &taskPollingFetchAdaptor{
+		fetched:      make(chan string, 4),
+		blockTaskID:  slowTask.GetUpstreamTaskID(),
+		blockStarted: make(chan struct{}),
+		releaseBlock: make(chan struct{}),
+	}
+	var releaseOnce sync.Once
+	releaseBlockedTask := func() {
+		releaseOnce.Do(func() {
+			close(adaptor.releaseBlock)
+		})
+	}
+	t.Cleanup(releaseBlockedTask)
+	previousFactory := GetTaskAdaptorFunc
+	GetTaskAdaptorFunc = func(constant.TaskPlatform) TaskPollingAdaptor { return adaptor }
+	t.Cleanup(func() { GetTaskAdaptorFunc = previousFactory })
+
+	errCh := make(chan error, 1)
+	gopool.Go(func() {
+		errCh <- UpdateVideoTasks(context.Background(), constant.TaskPlatform("kling"), map[int][]string{
+			slowChannelID: {
+				slowTask.GetUpstreamTaskID(),
+			},
+			fastChannelID: {
+				fastFirst.GetUpstreamTaskID(),
+				fastSecond.GetUpstreamTaskID(),
+			},
+		}, map[string]*model.Task{
+			slowTask.GetUpstreamTaskID():   slowTask,
+			fastFirst.GetUpstreamTaskID():  fastFirst,
+			fastSecond.GetUpstreamTaskID(): fastSecond,
+		})
+	})
+
+	select {
+	case <-adaptor.blockStarted:
+	case <-time.After(500 * time.Millisecond):
+		t.Fatal("slow channel did not start blocking")
+	}
+
+	require.Eventually(t, func() bool {
+		fetchedTaskIDs := adaptor.fetchedTaskIDs()
+		return len(fetchedTaskIDs) == 2 &&
+			fetchedTaskIDs[0] == fastFirst.GetUpstreamTaskID() &&
+			fetchedTaskIDs[1] == fastSecond.GetUpstreamTaskID()
+	}, 500*time.Millisecond, 10*time.Millisecond)
+
+	releaseBlockedTask()
+	require.NoError(t, <-errCh)
+	assert.ElementsMatch(t, []string{
+		slowTask.GetUpstreamTaskID(),
+		fastFirst.GetUpstreamTaskID(),
+		fastSecond.GetUpstreamTaskID(),
+	}, adaptor.fetchedTaskIDs())
+}
+
+func TestUpdateVideoTasksMixedChannelSleepSettings(t *testing.T) {
+	truncate(t)
+
+	const sleepyChannelID = 301
+	const fastChannelID = 302
+	seedTaskPollingChannel(t, sleepyChannelID, false)
+	seedTaskPollingChannel(t, fastChannelID, true)
+	sleepyFirst := seedPollingTask(t, sleepyChannelID, "task_public_9", "upstream_sleepy_1")
+	sleepySecond := seedPollingTask(t, sleepyChannelID, "task_public_10", "upstream_sleepy_2")
+	fastFirst := seedPollingTask(t, fastChannelID, "task_public_11", "upstream_fast_1")
+	fastSecond := seedPollingTask(t, fastChannelID, "task_public_12", "upstream_fast_2")
+
+	adaptor := &taskPollingFetchAdaptor{}
+	previousFactory := GetTaskAdaptorFunc
+	GetTaskAdaptorFunc = func(constant.TaskPlatform) TaskPollingAdaptor { return adaptor }
+	t.Cleanup(func() { GetTaskAdaptorFunc = previousFactory })
+
+	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+	defer cancel()
+
+	err := UpdateVideoTasks(ctx, constant.TaskPlatform("kling"), map[int][]string{
+		sleepyChannelID: {
+			sleepyFirst.GetUpstreamTaskID(),
+			sleepySecond.GetUpstreamTaskID(),
+		},
+		fastChannelID: {
+			fastFirst.GetUpstreamTaskID(),
+			fastSecond.GetUpstreamTaskID(),
+		},
+	}, map[string]*model.Task{
+		sleepyFirst.GetUpstreamTaskID():  sleepyFirst,
+		sleepySecond.GetUpstreamTaskID(): sleepySecond,
+		fastFirst.GetUpstreamTaskID():    fastFirst,
+		fastSecond.GetUpstreamTaskID():   fastSecond,
+	})
+
+	require.ErrorIs(t, err, context.DeadlineExceeded)
+	assert.ElementsMatch(t, []string{"upstream_sleepy_1", "upstream_fast_1", "upstream_fast_2"}, adaptor.fetchedTaskIDs())
+}