Files
NianAIGC/backend/internal/jobs/loop.go

81 lines
1.6 KiB
Go

package jobs
import (
"context"
"sync"
"sync/atomic"
"time"
)
type TickRunner interface {
Tick(context.Context, string) (TickResult, error)
}
type LoopConfig struct {
Interval time.Duration
WorkerID string
}
type WorkerLoop struct {
runner TickRunner
config LoopConfig
running atomic.Bool
mu sync.Mutex
cancel context.CancelFunc
done chan struct{}
}
func NewWorkerLoop(r TickRunner, c LoopConfig) *WorkerLoop {
if c.Interval <= 0 {
c.Interval = 5 * time.Second
}
if c.WorkerID == "" {
c.WorkerID = "embedded-worker"
}
return &WorkerLoop{runner: r, config: c}
}
func (l *WorkerLoop) Tick(ctx context.Context) (TickResult, bool) {
if l == nil || l.runner == nil || !l.running.CompareAndSwap(false, true) {
return TickResult{}, false
}
defer l.running.Store(false)
r, e := l.runner.Tick(ctx, l.config.WorkerID)
if e != nil {
return TickResult{WorkerID: l.config.WorkerID, Jobs: []TickJob{{Action: "failed", Error: e.Error()}}}, true
}
return r, true
}
func (l *WorkerLoop) Start(parent context.Context) {
l.mu.Lock()
defer l.mu.Unlock()
if l.cancel != nil {
return
}
ctx, cancel := context.WithCancel(parent)
l.cancel = cancel
l.done = make(chan struct{})
go func() {
defer close(l.done)
ticker := time.NewTicker(l.config.Interval)
defer ticker.Stop()
for {
_, _ = l.Tick(ctx)
select {
case <-ctx.Done():
return
case <-ticker.C:
}
}
}()
}
func (l *WorkerLoop) Stop() {
l.mu.Lock()
cancel, done := l.cancel, l.done
l.cancel = nil
l.done = nil
l.mu.Unlock()
if cancel == nil {
return
}
cancel()
<-done
}