new-api/service/channel_quota_reset_task.go

package service

import (
	"context"
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/logger"
	"github.com/QuantumNous/new-api/model"

	"github.com/bytedance/gopkg/util/gopool"
)

const (
	channelQuotaResetTickInterval = 1 * time.Minute
	channelQuotaResetBatchSize    = 300
)

var (
	channelQuotaResetOnce    sync.Once
	channelQuotaResetRunning atomic.Bool
)

func StartChannelQuotaResetTask() {
	channelQuotaResetOnce.Do(func() {
		if !common.IsMasterNode {
			return
		}
		gopool.Go(func() {
			logger.LogInfo(context.Background(), fmt.Sprintf("channel quota reset task started: tick=%s", channelQuotaResetTickInterval))
			ticker := time.NewTicker(channelQuotaResetTickInterval)
			defer ticker.Stop()

			runChannelQuotaResetOnce()
			for range ticker.C {
				runChannelQuotaResetOnce()
			}
		})
	})
}

func runChannelQuotaResetOnce() {
	if !channelQuotaResetRunning.CompareAndSwap(false, true) {
		return
	}
	defer channelQuotaResetRunning.Store(false)

	ctx := context.Background()
	totalReset := 0
	for {
		n, err := model.ResetDueChannels(channelQuotaResetBatchSize)
		if err != nil {
			logger.LogWarn(ctx, fmt.Sprintf("channel quota reset task failed: %v", err))
			return
		}
		if n == 0 {
			break
		}
		totalReset += n
		if n < channelQuotaResetBatchSize {
			break
		}
	}
	if common.DebugEnabled && totalReset > 0 {
		logger.LogDebug(ctx, "channel quota maintenance: reset_count=%d", totalReset)
	}
}