internal/store/postgres/bootstrap_update_job_batch.go

Rama main · réplica de lectura del repositorio original.

package postgres

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"time"

	"telesrv/internal/domain"
	"telesrv/internal/store"
)

// BootstrapReadyBatchMetrics exposes only bounded aggregate signals. Selector
// identities are deliberately excluded from metrics.
type BootstrapReadyBatchMetrics interface {
	BootstrapReadyBatch(inputs int, matched int, d time.Duration, err error)
	BootstrapReadyPending(delta int)
}

type BootstrapReadyBatchConfig struct {
	MaxSize      int
	MaxWait      time.Duration
	QueueSize    int
	QueryTimeout time.Duration
	Metrics      BootstrapReadyBatchMetrics
}

type bootstrapReadyBatchKey struct {
	userID    int64
	authKeyID [8]byte
}

type bootstrapReadyBatchRequest struct {
	userID    int64
	authKeyID [8]byte
	sessionID int64
	result    chan bootstrapReadyBatchResult
}

type bootstrapReadyBatchResult struct {
	matched int
	err     error
}

type bootstrapReadyBatchBackend interface {
	store.BootstrapUpdateJobStore
	markReadyForSessions(context.Context, []bootstrapReadyBatchRequest) ([]int, error)
}

// BatchedBootstrapUpdateJobStore preserves the synchronous post-response
// delivery fence while combining independent readiness selectors into one
// PostgreSQL statement. Once accepted, a selector waits for a definitive
// commit/error; it is never converted into an unobserved background write.
type BatchedBootstrapUpdateJobStore struct {
	base   bootstrapReadyBatchBackend
	cfg    BootstrapReadyBatchConfig
	queue  chan bootstrapReadyBatchRequest
	stop   chan struct{}
	done   chan struct{}
	cancel context.CancelFunc
	once   sync.Once
	gate   sync.RWMutex
	closed bool
}

func NewBatchedBootstrapUpdateJobStore(
	base *BootstrapUpdateJobStore,
	cfg BootstrapReadyBatchConfig,
) (*BatchedBootstrapUpdateJobStore, error) {
	if base == nil || base.db == nil {
		return nil, errors.New("initialize bootstrap readiness batcher: nil store")
	}
	return newBatchedBootstrapUpdateJobStore(base, cfg)
}

func newBatchedBootstrapUpdateJobStore(
	base bootstrapReadyBatchBackend,
	cfg BootstrapReadyBatchConfig,
) (*BatchedBootstrapUpdateJobStore, error) {
	if base == nil {
		return nil, errors.New("initialize bootstrap readiness batcher: nil backend")
	}
	if cfg.MaxSize <= 0 || cfg.MaxSize > 4096 {
		return nil, fmt.Errorf("initialize bootstrap readiness batcher: max size %d outside [1,4096]", cfg.MaxSize)
	}
	if cfg.MaxWait <= 0 || cfg.MaxWait > time.Second {
		return nil, fmt.Errorf("initialize bootstrap readiness batcher: max wait %v outside (0,1s]", cfg.MaxWait)
	}
	if cfg.QueueSize < cfg.MaxSize || cfg.QueueSize > 1<<20 {
		return nil, fmt.Errorf("initialize bootstrap readiness batcher: queue size %d outside [%d,%d]", cfg.QueueSize, cfg.MaxSize, 1<<20)
	}
	if cfg.QueryTimeout <= 0 || cfg.QueryTimeout > 30*time.Second {
		return nil, fmt.Errorf("initialize bootstrap readiness batcher: query timeout %v outside (0,30s]", cfg.QueryTimeout)
	}
	workerCtx, cancel := context.WithCancel(context.Background())
	s := &BatchedBootstrapUpdateJobStore{
		base: base, cfg: cfg,
		queue: make(chan bootstrapReadyBatchRequest, cfg.QueueSize),
		stop:  make(chan struct{}), done: make(chan struct{}), cancel: cancel,
	}
	if cfg.Metrics != nil {
		cfg.Metrics.BootstrapReadyPending(0)
	}
	go s.run(workerCtx)
	return s, nil
}

func (s *BatchedBootstrapUpdateJobStore) EnqueueLoginMessage(
	ctx context.Context,
	job domain.BootstrapUpdateJob,
) (domain.BootstrapUpdateJob, error) {
	return s.base.EnqueueLoginMessage(ctx, job)
}

func (s *BatchedBootstrapUpdateJobStore) MarkReadyForSession(
	ctx context.Context,
	userID int64,
	authKeyID [8]byte,
	sessionID int64,
) (int, error) {
	if ctx == nil {
		ctx = context.Background()
	}
	request := bootstrapReadyBatchRequest{
		userID: userID, authKeyID: authKeyID, sessionID: sessionID,
		result: make(chan bootstrapReadyBatchResult, 1),
	}
	s.gate.RLock()
	if s.closed {
		s.gate.RUnlock()
		return 0, context.Canceled
	}
	select {
	case s.queue <- request:
		if s.cfg.Metrics != nil {
			s.cfg.Metrics.BootstrapReadyPending(1)
		}
	case <-ctx.Done():
		s.gate.RUnlock()
		return 0, ctx.Err()
	}
	s.gate.RUnlock()

	// Accepted work ignores later caller cancellation and waits for the worker's
	// definitive result. This prevents a physically delivered baseline from
	// leaving an unknown asynchronous readiness mutation behind.
	result := <-request.result
	return result.matched, result.err
}

func (s *BatchedBootstrapUpdateJobStore) ClaimReady(
	ctx context.Context,
	limit int,
	leaseTimeout time.Duration,
) ([]domain.BootstrapUpdateJob, error) {
	return s.base.ClaimReady(ctx, limit, leaseTimeout)
}

func (s *BatchedBootstrapUpdateJobStore) MarkPublished(ctx context.Context, id int64) error {
	return s.base.MarkPublished(ctx, id)
}

func (s *BatchedBootstrapUpdateJobStore) MarkFailed(ctx context.Context, id int64, lastError string) error {
	return s.base.MarkFailed(ctx, id, lastError)
}

func (s *BatchedBootstrapUpdateJobStore) Close() {
	s.once.Do(func() {
		s.gate.Lock()
		s.closed = true
		close(s.stop)
		s.cancel()
		s.gate.Unlock()
		<-s.done
	})
}

func (s *BatchedBootstrapUpdateJobStore) run(ctx context.Context) {
	defer close(s.done)
	pending := make([]bootstrapReadyBatchRequest, 0, s.cfg.MaxSize)
	for {
		if len(pending) == 0 {
			select {
			case request := <-s.queue:
				pending = append(pending, request)
			case <-s.stop:
				s.failQueued(context.Canceled, pending)
				return
			}
		}

		if len(pending) < s.cfg.MaxSize {
			timer := time.NewTimer(s.cfg.MaxWait)
		collect:
			for len(pending) < s.cfg.MaxSize {
				select {
				case request := <-s.queue:
					pending = append(pending, request)
				case <-timer.C:
					break collect
				case <-s.stop:
					stopAndDrainTimer(timer)
					s.failQueued(context.Canceled, pending)
					return
				}
			}
			stopAndDrainTimer(timer)
		}

		batch, remaining := selectDistinctBootstrapReadyBatch(pending, s.cfg.MaxSize)
		pending = remaining
		s.execute(ctx, batch)
	}
}

func stopAndDrainTimer(timer *time.Timer) {
	if timer != nil && !timer.Stop() {
		select {
		case <-timer.C:
		default:
		}
	}
}

func selectDistinctBootstrapReadyBatch(
	pending []bootstrapReadyBatchRequest,
	maxSize int,
) ([]bootstrapReadyBatchRequest, []bootstrapReadyBatchRequest) {
	batch := make([]bootstrapReadyBatchRequest, 0, min(maxSize, len(pending)))
	remaining := make([]bootstrapReadyBatchRequest, 0, len(pending))
	seen := make(map[bootstrapReadyBatchKey]struct{}, min(maxSize, len(pending)))
	for _, request := range pending {
		if len(batch) >= maxSize {
			remaining = append(remaining, request)
			continue
		}
		key := bootstrapReadyBatchKey{userID: request.userID, authKeyID: request.authKeyID}
		if _, exists := seen[key]; exists {
			remaining = append(remaining, request)
			continue
		}
		seen[key] = struct{}{}
		batch = append(batch, request)
	}
	return batch, remaining
}

func (s *BatchedBootstrapUpdateJobStore) execute(ctx context.Context, batch []bootstrapReadyBatchRequest) {
	if len(batch) == 0 {
		return
	}
	started := time.Now()
	queryCtx, cancel := context.WithTimeout(ctx, s.cfg.QueryTimeout)
	results, err := s.base.markReadyForSessions(queryCtx, batch)
	cancel()
	matched := 0
	if err == nil {
		if len(results) != len(batch) {
			err = fmt.Errorf("mark bootstrap readiness batch: result count %d, want %d", len(results), len(batch))
		} else {
			for _, count := range results {
				matched += count
			}
		}
	}
	if s.cfg.Metrics != nil {
		s.cfg.Metrics.BootstrapReadyBatch(len(batch), matched, time.Since(started), err)
	}
	for index, request := range batch {
		result := bootstrapReadyBatchResult{err: err}
		if err == nil {
			result.matched = results[index]
		}
		request.result <- result
		if s.cfg.Metrics != nil {
			s.cfg.Metrics.BootstrapReadyPending(-1)
		}
	}
}

func (s *BatchedBootstrapUpdateJobStore) failQueued(err error, pending []bootstrapReadyBatchRequest) {
	for _, request := range pending {
		s.failRequest(request, err)
	}
	for {
		select {
		case request := <-s.queue:
			s.failRequest(request, err)
		default:
			return
		}
	}
}

func (s *BatchedBootstrapUpdateJobStore) failRequest(request bootstrapReadyBatchRequest, err error) {
	request.result <- bootstrapReadyBatchResult{err: err}
	if s.cfg.Metrics != nil {
		s.cfg.Metrics.BootstrapReadyPending(-1)
	}
}

func (s *BootstrapUpdateJobStore) markReadyForSessions(
	ctx context.Context,
	requests []bootstrapReadyBatchRequest,
) ([]int, error) {
	results := make([]int, len(requests))
	if len(requests) == 0 {
		return results, nil
	}
	userIDs := make([]int64, len(requests))
	authKeyIDs := make([]int64, len(requests))
	sessionIDs := make([]int64, len(requests))
	seen := make(map[bootstrapReadyBatchKey]struct{}, len(requests))
	for index, request := range requests {
		key := bootstrapReadyBatchKey{userID: request.userID, authKeyID: request.authKeyID}
		if _, duplicate := seen[key]; duplicate {
			return nil, fmt.Errorf("mark bootstrap readiness batch: duplicate fence at index %d", index)
		}
		seen[key] = struct{}{}
		userIDs[index] = request.userID
		authKeyIDs[index] = authKeyIDToInt64(request.authKeyID)
		sessionIDs[index] = request.sessionID
	}
	rows, err := s.db.Query(ctx, `
WITH input AS (
  SELECT *
  FROM unnest(
    $1::bigint[],
    $2::bigint[],
    $3::bigint[]
  ) WITH ORDINALITY AS value(user_id, auth_key_id, session_id, ordinal)
), candidates AS MATERIALIZED (
  SELECT input.ordinal,
         input.session_id,
         jobs.id
  FROM input
  JOIN bootstrap_update_jobs AS jobs
    ON jobs.user_id = input.user_id
   AND jobs.auth_key_id = input.auth_key_id
   AND jobs.status = 'pending'
  ORDER BY jobs.id, input.ordinal
  FOR UPDATE OF jobs
), updated AS (
  UPDATE bootstrap_update_jobs AS jobs
  SET status = 'ready',
      session_id = candidates.session_id,
      ready_at = now(),
      updated_at = now()
  FROM candidates
  WHERE jobs.id = candidates.id
  RETURNING candidates.ordinal
)
SELECT ordinal, count(*)::bigint
FROM updated
GROUP BY ordinal
ORDER BY ordinal`, userIDs, authKeyIDs, sessionIDs)
	if err != nil {
		return nil, fmt.Errorf("mark bootstrap readiness batch: %w", err)
	}
	defer rows.Close()
	for rows.Next() {
		var ordinal, count int64
		if err := rows.Scan(&ordinal, &count); err != nil {
			return nil, fmt.Errorf("scan bootstrap readiness batch: %w", err)
		}
		index := int(ordinal - 1)
		if index < 0 || index >= len(results) || results[index] != 0 || count <= 0 {
			return nil, fmt.Errorf("mark bootstrap readiness batch: invalid ordinal/count %d/%d", ordinal, count)
		}
		results[index] = int(count)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("mark bootstrap readiness batch rows: %w", err)
	}
	return results, nil
}

var _ store.BootstrapUpdateJobStore = (*BatchedBootstrapUpdateJobStore)(nil)