vidarchive

vidarchive/internal/worker/pool.go
pool.go
⎇
Raw
1package worker
2
3import (
4 "context"
5 "errors"
6 "log/slog"
7 "sync"
8 "time"
9
10 "vidarchive/internal/models"
11 "vidarchive/internal/service"
12)
13
14type Pool struct {
15 downloadSvc *service.DownloadService
16 workers int
17 queue chan *models.Download
18 ctx context.Context
19 cancel context.CancelFunc
20 wg sync.WaitGroup
21}
22
23func New(downloadSvc *service.DownloadService, workers int) *Pool {
24 ctx, cancel := context.WithCancel(context.Background())
25 return &Pool{
26 downloadSvc: downloadSvc,
27 workers: workers,
28 queue: make(chan *models.Download, 100),
29 ctx: ctx,
30 cancel: cancel,
31 }
32}
33
34func (p *Pool) Start() {
35 // Reset any downloads that were in progress during a previous run, and clear
36 // the temp dirs they left behind so a re-run doesn't import a duplicate.
37 if err := p.downloadSvc.ResetStalledDownloads(); err != nil {
38 slog.Error("failed to reset stalled downloads", "err", err)
39 }
40
41 for i := 0; i < p.workers; i++ {
42 p.wg.Add(1)
43 go p.worker(i)
44 }
45
46 p.wg.Add(1)
47 go p.queueChecker()
48}
49
50// Stop cancels running downloads and waits for the workers to return, so the
51// process doesn't exit while a yt-dlp child is still writing to the temp dir.
52func (p *Pool) Stop() {
53 p.cancel()
54 p.downloadSvc.CancelAll()
55 p.wg.Wait()
56}
57
58// Submit enqueues a download without blocking. When the buffer is full the row
59// stays "queued" in the database and the queue checker picks it up on a later
60// tick, so a burst of submissions can't stall the HTTP handler that made them.
61func (p *Pool) Submit(d *models.Download) {
62 if p.ctx.Err() != nil {
63 return
64 }
65
66 select {
67 case p.queue <- d:
68 slog.Debug("download handed to a worker", "download_id", d.ID)
69 default:
70 slog.Info("download left in database queue, worker buffer full", "download_id", d.ID)
71 }
72}
73
74func (p *Pool) worker(id int) {
75 defer p.wg.Done()
76 slog.Debug("worker started", "worker", id)
77 for {
78 select {
79 case d := <-p.queue:
80 slog.Info("processing download", "worker", id, "download_id", d.ID)
81 processed, err := p.downloadSvc.ExecuteDownload(p.ctx, d)
82 switch {
83 case errors.Is(err, service.ErrCancelled):
84 slog.Info("download cancelled", "worker", id, "download_id", d.ID)
85 case err != nil:
86 slog.Error("download failed", "worker", id, "download_id", d.ID, "err", err)
87 case processed:
88 slog.Info("download completed", "worker", id, "download_id", d.ID)
89 default:
90 slog.Debug("download already claimed by another worker", "worker", id, "download_id", d.ID)
91 }
92 case <-p.ctx.Done():
93 slog.Debug("worker stopped", "worker", id)
94 return
95 }
96 }
97}
98
99func (p *Pool) queueChecker() {
100 defer p.wg.Done()
101 ticker := time.NewTicker(2 * time.Second)
102 defer ticker.Stop()
103
104 for {
105 select {
106 case <-ticker.C:
107 p.checkQueue()
108 case <-p.ctx.Done():
109 return
110 }
111 }
112}
113
114func (p *Pool) checkQueue() {
115 // Pull at most a channel's worth of the oldest queued downloads rather than
116 // loading every queued row each tick. Anything beyond the buffer is picked up
117 // on a later tick; duplicates are harmless (MarkStarted claims atomically).
118 downloads, err := p.downloadSvc.GetQueued(cap(p.queue))
119 if err != nil {
120 slog.Error("queue check failed", "err", err)
121 return
122 }
123
124 for _, d := range downloads {
125 select {
126 case p.queue <- d:
127 default:
128 return
129 }
130 }
131}
132