vidarchive

vidarchive/internal/worker/pool.go
pool.go
⎇
Raw
1package worker
2
3import (
4 "context"
5 "fmt"
6 "log"
7 "time"
8
9 "vidarchive/internal/models"
10 "vidarchive/internal/service"
11)
12
13type Pool struct {
14 downloadSvc *service.DownloadService
15 workers int
16 queue chan *models.Download
17 ctx context.Context
18 cancel context.CancelFunc
19}
20
21func New(downloadSvc *service.DownloadService, workers int) *Pool {
22 ctx, cancel := context.WithCancel(context.Background())
23 return &Pool{
24 downloadSvc: downloadSvc,
25 workers: workers,
26 queue: make(chan *models.Download, 100),
27 ctx: ctx,
28 cancel: cancel,
29 }
30}
31
32func (p *Pool) Start() {
33 // Reset any downloads that were in progress during a previous run
34 if err := p.downloadSvc.ResetStalledDownloads(); err != nil {
35 log.Printf("Warning: failed to reset stalled downloads: %v", err)
36 }
37
38 for i := 0; i < p.workers; i++ {
39 go p.worker(i)
40 }
41
42 // Queue checker - polls DB for queued downloads
43 go p.queueChecker()
44}
45
46func (p *Pool) Stop() {
47 p.cancel()
48}
49
50func (p *Pool) Submit(d *models.Download) {
51 select {
52 case p.queue <- d:
53 log.Printf("Download %d queued", d.ID)
54 case <-p.ctx.Done():
55 }
56}
57
58func (p *Pool) worker(id int) {
59 log.Printf("Worker %d started", id)
60 for {
61 select {
62 case d := <-p.queue:
63 if d == nil {
64 continue
65 }
66 log.Printf("Worker %d processing download %d", id, d.ID)
67 processed, err := p.downloadSvc.ExecuteDownload(d)
68 switch {
69 case err != nil:
70 log.Printf("Worker %d download %d failed: %v", id, d.ID, err)
71 case processed:
72 log.Printf("Worker %d download %d completed", id, d.ID)
73 default:
74 log.Printf("Worker %d download %d already claimed by another worker, skipping", id, d.ID)
75 }
76 case <-p.ctx.Done():
77 log.Printf("Worker %d stopped", id)
78 return
79 }
80 }
81}
82
83func (p *Pool) queueChecker() {
84 ticker := time.NewTicker(2 * time.Second)
85 defer ticker.Stop()
86
87 for {
88 select {
89 case <-ticker.C:
90 p.checkQueue()
91 case <-p.ctx.Done():
92 return
93 }
94 }
95}
96
97func (p *Pool) checkQueue() {
98 // Check how many items are in queue vs active
99 downloads, err := p.downloadSvc.GetAll("queued", "date")
100 if err != nil {
101 log.Printf("Queue check error: %v", err)
102 return
103 }
104
105 for _, d := range downloads {
106 // Try to submit - if queue is full, it'll block briefly
107 select {
108 case p.queue <- d:
109 default:
110 // Queue is full, skip for now
111 return
112 }
113 }
114}
115
116func (p *Pool) GetQueueStatus() (queued, active, completed, failed int, err error) {
117 all, err := p.downloadSvc.GetAll("all", "date")
118 if err != nil {
119 return 0, 0, 0, 0, fmt.Errorf("get all downloads: %w", err)
120 }
121
122 for _, d := range all {
123 switch d.Status {
124 case "queued":
125 queued++
126 case "downloading":
127 active++
128 case "completed":
129 completed++
130 case "error":
131 failed++
132 }
133 }
134
135 return queued, active, completed, failed, nil
136}
137