iotlcm/worker-svr/internal/producer/dailyusage/producer.go

// Package dailyusage produces the previous-day usage full-sync run.
//
// It is the daily-usage twin of producer/monthlyusage and producer/simstatus:
// it scans the eligible cards once, groups them by channel, and enqueues one
// Redis ZSet task per card. The zset consumer owns pacing, retries, failure
// pooling, archiving and replay.
//
// The only structural difference from the monthly producer is that the run
// period is a day (yyyyMMdd) rather than a month (yyyyMM). That value travels
// in zset.Task.QueryMonth, which is a period identifier rather than a month
// name; see that field's documentation.
package dailyusage

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"strings"
	"sync/atomic"
	"time"

	"github.com/gogf/gf/v2/database/gdb"
	"github.com/gogf/gf/v2/frame/g"

	"git.5iot.com/iotlcm-public/iotlcm-worker-svr/internal/dao"
	"git.5iot.com/iotlcm-public/iotlcm-worker-svr/internal/model"
	"git.5iot.com/iotlcm-public/iotlcm-worker-svr/internal/model/entity"
	"git.5iot.com/iotlcm-public/iotlcm-worker-svr/internal/queue/zset"
)

// FullSyncOperation identifies a daily usage run in run metadata, archives and
// replay. contract.SyncKindForOperation maps it to traffic_usage_daily.
const FullSyncOperation = model.OperationPreviousDayUsageFullSync

// scopeSampleLimit bounds how many channel names are copied into run metadata.
// The run hash and its archive JSON must stay bounded regardless of how many
// channels a scope selects.
const scopeSampleLimit = 10

// scopeDigest is a stable short fingerprint of the selected channels. It lets
// operators confirm two runs covered the same scope without storing the list.
func scopeDigest(scope *model.SyncScope) string {
	if scope == nil || len(scope.Channels) == 0 {
		return ""
	}
	digest := sha256.Sum256([]byte(strings.Join(scope.Channels, "\n")))
	return hex.EncodeToString(digest[:8])
}

type Config struct {
	ScanMode             string
	ScanConcurrency      int
	PageSize             int
	GlobalHighWaterMark  int64
	ChannelHighWaterMark int64
	PollInterval         time.Duration
	RunLockKey           string
	RunLockTTL           time.Duration
	RoundRetryEnabled    bool
	MaxRounds            int
	FailureRetention     time.Duration
}

func (c Config) withDefaults() Config {
	if c.ScanMode == "" {
		c.ScanMode = "serial"
	}
	if c.ScanConcurrency == 0 {
		c.ScanConcurrency = 4
	}
	if c.PageSize <= 0 {
		c.PageSize = 100
	}
	if c.PollInterval <= 0 {
		c.PollInterval = 2 * time.Second
	}
	if c.RunLockKey == "" {
		c.RunLockKey = "daily:producer:active"
	}
	if c.RunLockTTL <= 0 {
		c.RunLockTTL = 5 * time.Minute
	}
	if c.MaxRounds <= 0 {
		c.MaxRounds = 3
	}
	if c.FailureRetention <= 0 {
		c.FailureRetention = 30 * 24 * time.Hour
	}
	return c
}

type CardScanner interface {
	SnapshotMaxID(context.Context) (int64, error)
	Scan(context.Context, int64, int64, int) ([]Card, error)
}

type Card struct {
	ID          int64
	ICCID       string
	MSISDN      string
	IMSI        string
	ChanKey     string
	TempKey     string
	PackageCode string
	SyncEnabled bool
	Operator    int
}

type channelMetadata struct {
	ChanKey     string `orm:"chan_key"`
	TempKey     string `orm:"temp_key"`
	SyncEnabled bool   `orm:"sync_enabled"`
	Operator    int    `orm:"operator"`
}

type DatabaseScanner struct {
	channelMetadata map[string]channelMetadata
}

func (s *DatabaseScanner) SnapshotMaxID(ctx context.Context) (int64, error) {
	var card entity.TSimCardInfo
	cardTable := dao.TSimCardInfo.Table()
	columns := dao.TSimCardInfo.Columns()
	if err := dailyUsageEligibleCardModel(ctx).
		Fields(cardTable + "." + columns.Id + " AS " + columns.Id).
		OrderDesc(cardTable + "." + columns.Id).
		Limit(1).
		Scan(&card); err != nil {
		return 0, fmt.Errorf("read card snapshot max id: %w", err)
	}
	return card.Id, nil
}

func (s *DatabaseScanner) Scan(ctx context.Context, lastID, maxID int64, pageSize int) ([]Card, error) {
	if maxID <= lastID {
		return nil, nil
	}
	if lastID == 0 || s.channelMetadata == nil {
		var metadata []channelMetadata
		if err := g.DB().Model("t_sim_channels").Fields("chan_key,temp_key,sync_enabled,operator").Scan(&metadata); err != nil {
			return nil, fmt.Errorf("scan channel metadata: %w", err)
		}
		s.channelMetadata = make(map[string]channelMetadata, len(metadata))
		for _, item := range metadata {
			s.channelMetadata[item.ChanKey] = item
		}
	}
	var rows []entity.TSimCardInfo
	columns := dao.TSimCardInfo.Columns()
	cardTable := dao.TSimCardInfo.Table()
	query := dailyUsageEligibleCardModel(ctx).
		Fields(
			cardTable+"."+columns.Id+" AS "+columns.Id,
			cardTable+"."+columns.Iccid+" AS "+columns.Iccid,
			cardTable+"."+columns.Msisdn+" AS "+columns.Msisdn,
			cardTable+"."+columns.Imsi+" AS "+columns.Imsi,
			cardTable+"."+columns.ChanKey+" AS "+columns.ChanKey,
			// The daily usage child table stores the package code and the
			// diagnosis page reads it back, so it must be carried on the task.
			cardTable+"."+columns.CurrentPackageCode+" AS "+columns.CurrentPackageCode,
		).
		WhereGT(cardTable+"."+columns.Id, lastID).
		WhereLTE(cardTable+"."+columns.Id, maxID).
		OrderAsc(cardTable + "." + columns.Id).
		Limit(pageSize)
	if err := query.Scan(&rows); err != nil {
		return nil, fmt.Errorf("scan card information: %w", err)
	}
	cards := make([]Card, 0, len(rows))
	for _, row := range rows {
		channel := s.channelMetadata[row.ChanKey]
		cards = append(cards, Card{
			ID:          row.Id,
			ICCID:       row.Iccid,
			MSISDN:      row.Msisdn,
			IMSI:        row.Imsi,
			ChanKey:     row.ChanKey,
			TempKey:     channel.TempKey,
			PackageCode: row.CurrentPackageCode,
			SyncEnabled: channel.SyncEnabled,
			Operator:    channel.Operator,
		})
	}
	return cards, nil
}

// dailyUsageEligibleCardModel mirrors the monthly usage eligibility rule and the
// legacy admin SyncDaily scan: the card must belong to a live department, carry
// an ICCID and channel, and have a current package code.
func dailyUsageEligibleCardModel(ctx context.Context) *gdb.Model {
	cardTable := dao.TSimCardInfo.Table()
	deptTable := dao.SysDept.Table()
	cardColumns := dao.TSimCardInfo.Columns()
	deptColumns := dao.SysDept.Columns()

	return dao.TSimCardInfo.Ctx(ctx).
		InnerJoin(deptTable, cardTable+"."+cardColumns.DeptId+" = "+deptTable+"."+deptColumns.DeptId+" AND "+cardTable+"."+cardColumns.TenantId+" = "+deptTable+"."+deptColumns.TenantId).
		Where(cardTable+"."+cardColumns.Iccid+" <> ?", "").
		Where(cardTable+"."+cardColumns.ChanKey+" <> ?", "").
		Where(cardTable+"."+cardColumns.CurrentPackageCode+" <> ?", "").
		WhereNull(deptTable + "." + deptColumns.DeletedAt)
}

type Producer struct {
	scanner     CardScanner
	store       *zset.Store
	config      Config
	running     atomic.Int32
	owner       string
	observation *supplyObservation
}

type productionStats struct {
	pages               int64
	scanned             int64
	eligible            int64
	enqueued            int64
	skippedEmptyChan    int64
	skippedNoTemplate   int64
	skippedSyncDisabled int64
	skippedOutOfScope   int64
	stoppedTasks        int64
	channels            map[string]struct{}
	scopeChannels       int
	skippedChannels     int
}

func (s *productionStats) channelCount() int {
	return len(s.channels)
}

func (s *productionStats) log(ctx context.Context, status, runID, queryDate, stage, reason string, startedAt time.Time, maxID int64, err error) {
	message := fmt.Sprintf("event=daily_usage_producer status=%s runId=%s queryDate=%s snapshotMaxId=%d pages=%d scanned=%d eligible=%d enqueued=%d skippedEmptyChan=%d skippedNoTemplate=%d skippedSyncDisabled=%d skippedOutOfScope=%d stoppedTasks=%d channelCount=%d scopeChannels=%d skippedChannels=%d durationMs=%d", status, runID, queryDate, maxID, s.pages, s.scanned, s.eligible, s.enqueued, s.skippedEmptyChan, s.skippedNoTemplate, s.skippedSyncDisabled, s.skippedOutOfScope, s.stoppedTasks, s.channelCount(), s.scopeChannels, s.skippedChannels, time.Since(startedAt).Milliseconds())
	if s.scanned > 0 {
		message += fmt.Sprintf(" cardsPerSec=%.2f", float64(s.scanned)/time.Since(startedAt).Seconds())
	}
	if stage != "" {
		message += " stage=" + stage
	}
	if reason != "" {
		message += " reason=" + reason
	}
	if err != nil {
		message += fmt.Sprintf(" error=%q", err.Error())
	}
	if status == "failed" {
		g.Log().Error(ctx, message)
		return
	}
	if status == "skipped" {
		g.Log().Warning(ctx, message)
		return
	}
	g.Log().Info(ctx, message)
}

func New(scanner CardScanner, store *zset.Store, config Config) (*Producer, error) {
	if scanner == nil || store == nil {
		return nil, fmt.Errorf("daily usage producer dependencies are incomplete")
	}
	config = config.withDefaults()
	if err := ValidateScanSettings(config.ScanMode, config.ScanConcurrency); err != nil {
		return nil, err
	}
	if config.ScanMode == "channel" {
		if _, ok := scanner.(channelScanner); !ok {
			return nil, fmt.Errorf("channel scan requires ListChannels and ScanChannel")
		}
	}
	return &Producer{scanner: scanner, store: store, config: config, owner: newOwner()}, nil
}

var ErrAlreadyRunning = fmt.Errorf("daily usage producer already running")

// ErrRunFinished is returned when a redelivered start message names a run that
// already reached a terminal state. The run must not be resurrected; callers
// should log and ack instead of retrying.
var ErrRunFinished = fmt.Errorf("daily usage run already finished")

// runAlreadyFinished reports whether runID exists and has already completed.
// A missing run is not an error: the message may be the first delivery.
func (p *Producer) runAlreadyFinished(ctx context.Context, runID string) (bool, error) {
	meta, err := p.store.GetRunMeta(ctx, runID)
	if err != nil {
		if errors.Is(err, zset.ErrRunNotFound) {
			return false, nil
		}
		return false, err
	}
	return meta.Status == "completed" || meta.Status == "completed_with_failures", nil
}

// Produce scans the eligible cards once and enqueues one task per card.
//
// message.QueryDate (yyyyMMdd) is stored in zset.Task.QueryMonth, the queue's
// period-identifier field, so the zset store, run metadata, archive, replay and
// Lua scripts need no change for a daily run.
func (p *Producer) Produce(ctx context.Context, message model.TrafficUsageDailyStartMessage) (err error) {
	startedAt := time.Now()
	stats := productionStats{channels: make(map[string]struct{})}
	if !p.running.CompareAndSwap(0, 1) {
		stats.log(ctx, "skipped", message.RunID, message.QueryDate, "precheck", "already_running", startedAt, 0, nil)
		return ErrAlreadyRunning
	}
	defer p.running.Store(0)

	if message.RunID == "" {
		err = fmt.Errorf("runId is empty")
		stats.log(ctx, "failed", message.RunID, message.QueryDate, "validate", "", startedAt, 0, err)
		return err
	}
	if message.QueryDate == "" {
		err = fmt.Errorf("queryDate is empty")
		stats.log(ctx, "failed", message.RunID, message.QueryDate, "validate", "", startedAt, 0, err)
		return err
	}
	scope, err := model.NormalizeScope(message.Scope)
	if err != nil {
		stats.log(ctx, "failed", message.RunID, message.QueryDate, "validate", "invalid_scope", startedAt, 0, err)
		return err
	}
	stats.scopeChannels = scope.ChannelCount()
	// A redelivered start message must not resurrect a run that already
	// finished; that would create a second producer pass over the same RunID.
	if p.config.RoundRetryEnabled {
		if finished, checkErr := p.runAlreadyFinished(ctx, message.RunID); checkErr != nil {
			stats.log(ctx, "failed", message.RunID, message.QueryDate, "run_state", "", startedAt, 0, checkErr)
			return checkErr
		} else if finished {
			stats.log(ctx, "skipped", message.RunID, message.QueryDate, "run_state", "run_finished", startedAt, 0, nil)
			return ErrRunFinished
		}
	}
	acquired, err := p.store.AcquireLockOrTakeover(ctx, p.config.RunLockKey, p.owner, message.RunID, p.config.RunLockTTL)
	if err != nil {
		stats.log(ctx, "failed", message.RunID, message.QueryDate, "lock", "", startedAt, 0, err)
		return err
	}
	if !acquired {
		stats.log(ctx, "skipped", message.RunID, message.QueryDate, "lock", "already_running", startedAt, 0, nil)
		return ErrAlreadyRunning
	}
	runCtx, cancel := context.WithCancel(ctx)
	defer cancel()
	defer func() {
		releaseCtx, releaseCancel := context.WithTimeout(context.Background(), time.Second)
		defer releaseCancel()
		_, _ = p.store.ReleaseLock(releaseCtx, p.config.RunLockKey, p.owner)
	}()
	pending, err := p.hasPendingTasks(ctx, message.RunID)
	if err != nil {
		stats.log(ctx, "failed", message.RunID, message.QueryDate, "pending_check", "", startedAt, 0, err)
		return err
	}
	if pending {
		stats.log(ctx, "skipped", message.RunID, message.QueryDate, "pending_check", "pending_tasks", startedAt, 0, nil)
		return ErrAlreadyRunning
	}
	renewErr := make(chan error, 1)
	go p.renewLock(runCtx, cancel, renewErr)
	ctx = runCtx
	if message.Operation == "" {
		message.Operation = FullSyncOperation
	}
	if p.config.RoundRetryEnabled {
		if err = p.store.CreateRun(ctx, zset.RunMeta{
			RunID:               message.RunID,
			Operation:           message.Operation,
			QueryMonth:          message.QueryDate,
			Status:              "producing",
			CurrentRound:        1,
			ScopeChannels:       scope.ChannelCount(),
			ScopeDigest:         scopeDigest(scope),
			ScopeChannelsSample: strings.Join(scope.ChannelSample(scopeSampleLimit), ","),
		}); err != nil {
			stats.log(ctx, "failed", message.RunID, message.QueryDate, "run_init", "", startedAt, 0, err)
			return err
		}
	}
	maxID, err := p.scanner.SnapshotMaxID(ctx)
	if err != nil {
		stats.log(ctx, "failed", message.RunID, message.QueryDate, "snapshot", "", startedAt, 0, err)
		return err
	}
	g.Log().Infof(ctx, "event=daily_usage_producer status=started runId=%s queryDate=%s snapshotMaxId=%d scopeChannels=%d scanMode=%s scanConcurrency=%d pageSize=%d", message.RunID, message.QueryDate, maxID, scope.ChannelCount(), p.config.ScanMode, p.config.ScanConcurrency, p.config.PageSize)
	stage := "scan"
	defer func() {
		if err != nil {
			stats.log(ctx, "failed", message.RunID, message.QueryDate, stage, "", startedAt, maxID, err)
			return
		}
		stats.log(ctx, "completed", message.RunID, message.QueryDate, "", "", startedAt, maxID, nil)
	}()
	if p.config.RoundRetryEnabled && p.config.ScanMode == "channel" {
		stage = "channel_scan"
		if err = p.produceChannels(ctx, message, scope, maxID, startedAt, &stats); err != nil {
			if isProductionClosed(err) {
				// A same-Run redelivery after producerDone leaves consumption alone.
				err = nil
				return nil
			}
			return err
		}
		// Verify ownership immediately before closing production. The periodic
		// renewer still cancels scanners and capacity waits on lease loss.
		var owned bool
		owned, err = p.store.RenewLock(ctx, p.config.RunLockKey, p.owner, p.config.RunLockTTL)
		if err != nil {
			return err
		}
		if !owned {
			return fmt.Errorf("producer lock token is no longer owned")
		}
		if err = ctx.Err(); err != nil {
			return err
		}
		stage = "producer_done"
		err = p.store.MarkProducerDone(ctx, message.RunID)
		return err
	}
	lastID := int64(0)
	for lastID < maxID {
		stage = "capacity"
		if err = p.waitForCapacity(ctx, message.RunID); err != nil {
			return err
		}
		stage = "scan"
		cards, scanErr := p.scanner.Scan(ctx, lastID, maxID, p.config.PageSize)
		if scanErr != nil {
			err = scanErr
			return err
		}
		if len(cards) == 0 {
			break
		}
		stats.pages++
		stats.scanned += int64(len(cards))
		groupedTasks := make(map[string][]zset.Task)
		for _, card := range cards {
			lastID = card.ID
			if card.ChanKey == "" {
				stats.skippedEmptyChan++
				g.Log().Warningf(ctx, "event=daily_usage_card status=skipped stage=validate cardId=%d reason=empty_chan_key", card.ID)
				continue
			}
			// Scope is applied before the sync_enabled and template checks so a
			// partial run only ever reports on the channels it was asked for.
			if !scope.Allows(card.ChanKey) {
				stats.skippedOutOfScope++
				continue
			}
			if card.SyncEnabled {
				stats.skippedSyncDisabled++
				g.Log().Warningf(ctx, "event=daily_usage_card status=skipped stage=validate cardId=%d chanKey=%s reason=sync_disabled", card.ID, card.ChanKey)
				continue
			}
			if card.TempKey == "" {
				stats.skippedNoTemplate++
				g.Log().Warningf(ctx, "event=daily_usage_card status=skipped stage=validate cardId=%d chanKey=%s reason=missing_channel_template", card.ID, card.ChanKey)
				continue
			}
			stats.eligible++
			stats.channels[card.ChanKey] = struct{}{}
			task := zset.Task{
				TaskID:      zset.StableTaskID(message.RunID, message.Operation, card.ID),
				RunID:       message.RunID,
				Operation:   message.Operation,
				QueryMonth:  message.QueryDate,
				ChanKey:     card.ChanKey,
				TempKey:     card.TempKey,
				Operator:    card.Operator,
				CardID:      card.ID,
				ICCID:       card.ICCID,
				MSISDN:      card.MSISDN,
				IMSI:        card.IMSI,
				PackageCode: card.PackageCode,
			}
			groupedTasks[card.ChanKey] = append(groupedTasks[card.ChanKey], task)
		}
		now := time.Now()
		for chanKey, tasks := range groupedTasks {
			stage = "enqueue"
			if err = p.waitForChannelCapacity(ctx, message.RunID, chanKey); err != nil {
				return err
			}
			var added int
			var enqueueErr error
			if p.config.RoundRetryEnabled {
				added, enqueueErr = p.store.EnqueueRunBatch(ctx, message.RunID, chanKey, tasks, now)
			} else {
				added, enqueueErr = p.store.EnqueueBatch(ctx, chanKey, tasks, now)
			}
			if enqueueErr != nil {
				err = enqueueErr
				return err
			}
			stats.enqueued += int64(added)
			if added < len(tasks) {
				stats.stoppedTasks += int64(len(tasks) - added)
			}
			if added == 0 {
				g.Log().Warningf(ctx, "event=daily_usage_channel status=skipped runId=%s queryDate=%s chanKey=%s count=%d reason=stopped", message.RunID, message.QueryDate, chanKey, len(tasks))
			}
		}
	}
	if p.config.RoundRetryEnabled {
		if err = p.store.MarkProducerDone(ctx, message.RunID); err != nil {
			return err
		}
	}
	return nil
}

func (p *Producer) hasPendingTasks(ctx context.Context, currentRunID ...string) (bool, error) {
	currentID := ""
	if len(currentRunID) > 0 {
		currentID = currentRunID[0]
	}
	if p.config.RoundRetryEnabled {
		runs, err := p.store.ActiveRuns(ctx)
		if err != nil {
			return false, err
		}
		for _, runID := range runs {
			if runID == currentID {
				// A redelivered message for the same run may resume a producer
				// that stopped before producerDone. Stable task IDs make the
				// rescan idempotent; other runs still provide backpressure.
				continue
			}
			meta, metaErr := p.store.GetRunMeta(ctx, runID)
			if metaErr != nil {
				// The active index is a global key written outside the run-scoped
				// Lua transaction, so it can list a run whose hash is already gone
				// (expired, evicted, or cleaned by hand). Returning an error here
				// would surface as a non-retryable failure and make the consumer
				// redeliver this request forever. Treat it as a stale index entry
				// and drop it.
				if errors.Is(metaErr, zset.ErrRunNotFound) {
					if _, removeErr := p.store.RedisClient().Do(ctx, "SREM", p.store.ActiveRunsKey(), runID); removeErr != nil {
						return false, removeErr
					}
					g.Log().Warningf(ctx, "event=daily_usage_producer status=skipped stage=pending_check reason=stale_index runId=%s", runID)
					continue
				}
				return false, metaErr
			}
			if meta.Status != "producing" && meta.Status != "running" {
				continue
			}
			// Channel blocking can set running before production ends. Only
			// producerDone identifies consumer-owned work; replay has its own lock.
			if meta.ProducerDone || meta.SyncType == "replay" {
				return true, nil
			}
			// Recheck the unique producer lease immediately before reclaiming
			// another run. A lost lease or an unspecified current run is not proof
			// that the old producer has stopped.
			if currentID == "" {
				return true, nil
			}
			owned, lockErr := p.store.RenewLock(ctx, p.config.RunLockKey, p.owner, p.config.RunLockTTL)
			if lockErr != nil {
				return false, lockErr
			}
			if !owned {
				return true, nil
			}
			// Close production, not the run: consumers retain queued, leased and
			// next-round work and apply retention/index cleanup on normal completion.
			reclaimed, reclaimErr := p.store.MarkProducerAbandoned(ctx, runID)
			if reclaimErr != nil {
				return false, reclaimErr
			}
			g.Log().Warningf(ctx, "event=daily_usage_producer status=skipped stage=pending_check reason=zombie_run runId=%s reclaimed=%t activeCount=%d inflightCount=%d nextRoundCount=%d", runID, reclaimed, meta.ActiveCount, meta.InflightCount, meta.NextRoundCount)
			if !reclaimed || meta.ActiveCount > 0 || meta.InflightCount > 0 || meta.NextRoundCount > 0 {
				return true, nil
			}
		}
	}
	channels, err := p.store.ActiveChannels(ctx)
	if err != nil {
		return false, err
	}
	for _, chanKey := range channels {
		size, sizeErr := p.store.Size(ctx, chanKey)
		if sizeErr != nil {
			return false, sizeErr
		}
		if size > 0 {
			return true, nil
		}
	}
	return false, nil
}

func (p *Producer) renewLock(ctx context.Context, cancel context.CancelFunc, errors chan<- error) {
	interval := p.config.RunLockTTL / 3
	if interval <= 0 {
		interval = time.Second
	}
	ticker := time.NewTicker(interval)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			ok, err := p.store.RenewLock(ctx, p.config.RunLockKey, p.owner, p.config.RunLockTTL)
			if err != nil || !ok {
				if err == nil {
					err = fmt.Errorf("producer lock token is no longer owned")
				}
				select {
				case errors <- fmt.Errorf("renew producer lock: %w", err):
				default:
				}
				cancel()
				return
			}
		}
	}
}

func newOwner() string {
	var bytes [12]byte
	if _, err := rand.Read(bytes[:]); err == nil {
		return "producer-" + hex.EncodeToString(bytes[:])
	}
	return fmt.Sprintf("producer-%d", time.Now().UnixNano())
}

func (p *Producer) waitForCapacity(ctx context.Context, runID ...string) error {
	if p.config.GlobalHighWaterMark <= 0 {
		return nil
	}
	for {
		if p.config.RoundRetryEnabled && len(runID) > 0 && runID[0] != "" {
			meta, err := p.store.GetRunMeta(ctx, runID[0])
			if err != nil {
				return err
			}
			channels, err := p.store.RunChannels(ctx, runID[0])
			if err != nil {
				return err
			}
			var total int64
			for _, chanKey := range channels {
				key, keyErr := p.store.RoundActiveKey(runID[0], meta.CurrentRound, chanKey)
				if keyErr != nil {
					return keyErr
				}
				size, sizeErr := p.store.SizeKey(ctx, key)
				if sizeErr != nil {
					return sizeErr
				}
				total += size
			}
			if total < p.config.GlobalHighWaterMark {
				return nil
			}
			if err = sleepContext(ctx, p.config.PollInterval); err != nil {
				return err
			}
			continue
		}
		channels, err := p.store.ActiveChannels(ctx)
		if err != nil {
			return err
		}
		var total int64
		for _, chanKey := range channels {
			size, sizeErr := p.store.Size(ctx, chanKey)
			if sizeErr != nil {
				return sizeErr
			}
			total += size
		}
		if total < p.config.GlobalHighWaterMark {
			return nil
		}
		if err = sleepContext(ctx, p.config.PollInterval); err != nil {
			return err
		}
	}
}

func (p *Producer) waitForChannelCapacity(ctx context.Context, runIDAndChan ...string) error {
	if p.config.ChannelHighWaterMark <= 0 {
		return nil
	}
	chanKey := ""
	runID := ""
	if len(runIDAndChan) == 1 {
		chanKey = runIDAndChan[0]
	} else if len(runIDAndChan) >= 2 {
		runID, chanKey = runIDAndChan[0], runIDAndChan[1]
	}
	for {
		var size int64
		var err error
		if p.config.RoundRetryEnabled && runID != "" {
			meta, metaErr := p.store.GetRunMeta(ctx, runID)
			if metaErr != nil {
				return metaErr
			}
			key, keyErr := p.store.RoundActiveKey(runID, meta.CurrentRound, chanKey)
			if keyErr != nil {
				return keyErr
			}
			size, err = p.store.SizeKey(ctx, key)
		} else {
			size, err = p.store.Size(ctx, chanKey)
		}
		if err != nil {
			return err
		}
		if size < p.config.ChannelHighWaterMark {
			return nil
		}
		if err = sleepContext(ctx, p.config.PollInterval); err != nil {
			return err
		}
	}
}

func sleepContext(ctx context.Context, duration time.Duration) error {
	timer := time.NewTimer(duration)
	defer timer.Stop()
	select {
	case <-timer.C:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}