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 if err := p.downloadSvc.ExecuteDownload(d); err != nil {
68 log.Printf("Worker %d download %d failed: %v", id, d.ID, err)
69 } else {
70 log.Printf("Worker %d download %d completed", id, d.ID)
71 }
72 case <-p.ctx.Done():
73 log.Printf("Worker %d stopped", id)
74 return
75 }
76 }
77}
78
79func (p *Pool) queueChecker() {
80 ticker := time.NewTicker(2 * time.Second)
81 defer ticker.Stop()
82
83 for {
84 select {
85 case <-ticker.C:
86 p.checkQueue()
87 case <-p.ctx.Done():
88 return
89 }
90 }
91}
92
93func (p *Pool) checkQueue() {
94 // Check how many items are in queue vs active
95 downloads, err := p.downloadSvc.GetAll("queued", "date")
96 if err != nil {
97 log.Printf("Queue check error: %v", err)
98 return
99 }
100
101 for _, d := range downloads {
102 // Try to submit - if queue is full, it'll block briefly
103 select {
104 case p.queue <- d:
105 default:
106 // Queue is full, skip for now
107 return
108 }
109 }
110}
111
112func (p *Pool) GetQueueStatus() (queued, active, completed, failed int, err error) {
113 all, err := p.downloadSvc.GetAll("all", "date")
114 if err != nil {
115 return 0, 0, 0, 0, fmt.Errorf("get all downloads: %w", err)
116 }
117
118 for _, d := range all {
119 switch d.Status {
120 case "queued":
121 queued++
122 case "downloading":
123 active++
124 case "completed":
125 completed++
126 case "error":
127 failed++
128 }
129 }
130
131 return queued, active, completed, failed, nil
132}
133