149 lines
3.7 KiB
Go
149 lines
3.7 KiB
Go
package worker
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"log/slog"
|
|
"sync"
|
|
|
|
"galaxy/mail/internal/service/executeattempt"
|
|
)
|
|
|
|
// AttemptExecutionService executes one claimed in-progress attempt.
|
|
type AttemptExecutionService interface {
|
|
// Execute runs one claimed attempt through provider execution and durable
|
|
// state mutation.
|
|
Execute(context.Context, executeattempt.WorkItem) error
|
|
}
|
|
|
|
// AttemptWorkerPoolConfig stores the dependencies used by AttemptWorkerPool.
|
|
type AttemptWorkerPoolConfig struct {
|
|
// Concurrency stores how many workers run concurrently.
|
|
Concurrency int
|
|
|
|
// WorkQueue stores the claimed attempt handoff channel produced by the
|
|
// scheduler.
|
|
WorkQueue <-chan executeattempt.WorkItem
|
|
|
|
// Service executes one claimed attempt.
|
|
Service AttemptExecutionService
|
|
}
|
|
|
|
// AttemptWorkerPool executes claimed attempts concurrently.
|
|
type AttemptWorkerPool struct {
|
|
concurrency int
|
|
workQueue <-chan executeattempt.WorkItem
|
|
service AttemptExecutionService
|
|
logger *slog.Logger
|
|
}
|
|
|
|
// NewAttemptWorkerPool constructs one attempt worker pool.
|
|
func NewAttemptWorkerPool(cfg AttemptWorkerPoolConfig, logger *slog.Logger) (*AttemptWorkerPool, error) {
|
|
switch {
|
|
case cfg.Concurrency <= 0:
|
|
return nil, errors.New("new attempt worker pool: concurrency must be positive")
|
|
case cfg.WorkQueue == nil:
|
|
return nil, errors.New("new attempt worker pool: nil work queue")
|
|
case cfg.Service == nil:
|
|
return nil, errors.New("new attempt worker pool: nil attempt execution service")
|
|
}
|
|
if logger == nil {
|
|
logger = slog.Default()
|
|
}
|
|
|
|
return &AttemptWorkerPool{
|
|
concurrency: cfg.Concurrency,
|
|
workQueue: cfg.WorkQueue,
|
|
service: cfg.Service,
|
|
logger: logger.With("component", "attempt_worker_pool", "concurrency", cfg.Concurrency),
|
|
}, nil
|
|
}
|
|
|
|
// Run starts the attempt worker pool and blocks until ctx is canceled or one
|
|
// worker returns an execution error.
|
|
func (pool *AttemptWorkerPool) Run(ctx context.Context) error {
|
|
if ctx == nil {
|
|
return errors.New("run attempt worker pool: nil context")
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return err
|
|
}
|
|
if pool == nil {
|
|
return errors.New("run attempt worker pool: nil pool")
|
|
}
|
|
|
|
pool.logger.Info("attempt worker pool started")
|
|
defer pool.logger.Info("attempt worker pool stopped")
|
|
|
|
runCtx, cancel := context.WithCancel(ctx)
|
|
defer cancel()
|
|
|
|
errs := make(chan error, pool.concurrency)
|
|
var waitGroup sync.WaitGroup
|
|
|
|
for index := 0; index < pool.concurrency; index++ {
|
|
waitGroup.Add(1)
|
|
go func(workerIndex int) {
|
|
defer waitGroup.Done()
|
|
if err := pool.runWorker(runCtx, workerIndex); err != nil {
|
|
errs <- err
|
|
}
|
|
}(index)
|
|
}
|
|
|
|
done := make(chan struct{})
|
|
go func() {
|
|
waitGroup.Wait()
|
|
close(done)
|
|
}()
|
|
|
|
select {
|
|
case <-ctx.Done():
|
|
cancel()
|
|
<-done
|
|
return ctx.Err()
|
|
case err := <-errs:
|
|
cancel()
|
|
<-done
|
|
return err
|
|
case <-done:
|
|
if ctx.Err() != nil {
|
|
return ctx.Err()
|
|
}
|
|
return errors.New("run attempt worker pool: workers exited without shutdown")
|
|
}
|
|
}
|
|
|
|
func (pool *AttemptWorkerPool) runWorker(ctx context.Context, workerIndex int) error {
|
|
pool.logger.Debug("attempt worker started", "worker_index", workerIndex)
|
|
defer pool.logger.Debug("attempt worker stopped", "worker_index", workerIndex)
|
|
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
case item, ok := <-pool.workQueue:
|
|
if !ok {
|
|
return nil
|
|
}
|
|
if err := pool.service.Execute(ctx, item); err != nil {
|
|
return fmt.Errorf("attempt worker %d: %w", workerIndex, err)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Shutdown stops the attempt worker pool within ctx. The pool does not own
|
|
// additional resources beyond its run loop.
|
|
func (pool *AttemptWorkerPool) Shutdown(ctx context.Context) error {
|
|
if ctx == nil {
|
|
return errors.New("shutdown attempt worker pool: nil context")
|
|
}
|
|
if pool == nil {
|
|
return nil
|
|
}
|
|
|
|
return nil
|
|
}
|