104 lines
2.5 KiB
Go
104 lines
2.5 KiB
Go
package download
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
)
|
|
|
|
// Result pairs a finished download with its outcome and a final stat snapshot.
|
|
type Result struct {
|
|
Name string
|
|
Index int // position in the slice passed to Run, so callers can restore input order
|
|
Err error
|
|
Stat Stat
|
|
}
|
|
|
|
// Engine runs a set of downloads, at most maxConcurrent at a time, and exposes
|
|
// a live snapshot of the running ones for the progress renderer. A buffered
|
|
// channel of slots enforces the concurrency limit, and one goroutine per
|
|
// download replaces an explicit scheduler tick.
|
|
type Engine struct {
|
|
maxConcurrent int
|
|
|
|
mu sync.Mutex
|
|
active []Download
|
|
}
|
|
|
|
// NewEngine returns an engine that runs at most maxConcurrent downloads at once.
|
|
func NewEngine(maxConcurrent int) *Engine {
|
|
if maxConcurrent < 1 {
|
|
maxConcurrent = 1
|
|
}
|
|
return &Engine{maxConcurrent: maxConcurrent}
|
|
}
|
|
|
|
// Run starts every download, honouring the concurrency limit, and returns once
|
|
// all of them have finished or ctx is cancelled. There is one Result per
|
|
// download; results arrive in completion order, but each Result carries its
|
|
// input Index so callers can sort back to the original slice order.
|
|
func (e *Engine) Run(ctx context.Context, downloads []Download) []Result {
|
|
slots := make(chan struct{}, e.maxConcurrent)
|
|
results := make(chan Result, len(downloads))
|
|
|
|
var wg sync.WaitGroup
|
|
for i, d := range downloads {
|
|
wg.Add(1)
|
|
go func(i int, d Download) {
|
|
defer wg.Done()
|
|
|
|
// Acquire a concurrency slot, or bail if we're shutting down
|
|
// before this download ever started.
|
|
select {
|
|
case slots <- struct{}{}:
|
|
case <-ctx.Done():
|
|
results <- Result{Name: d.Name(), Index: i, Err: ctx.Err(), Stat: d.Stat()}
|
|
return
|
|
}
|
|
defer func() { <-slots }()
|
|
|
|
e.track(d)
|
|
defer e.untrack(d)
|
|
|
|
err := d.Run(ctx)
|
|
results <- Result{Name: d.Name(), Index: i, Err: err, Stat: d.Stat()}
|
|
}(i, d)
|
|
}
|
|
|
|
wg.Wait()
|
|
close(results)
|
|
|
|
out := make([]Result, 0, len(downloads))
|
|
for r := range results {
|
|
out = append(out, r)
|
|
}
|
|
return out
|
|
}
|
|
|
|
// Snapshot returns a Stat for every currently running download.
|
|
func (e *Engine) Snapshot() []Stat {
|
|
e.mu.Lock()
|
|
defer e.mu.Unlock()
|
|
out := make([]Stat, 0, len(e.active))
|
|
for _, d := range e.active {
|
|
out = append(out, d.Stat())
|
|
}
|
|
return out
|
|
}
|
|
|
|
func (e *Engine) track(d Download) {
|
|
e.mu.Lock()
|
|
e.active = append(e.active, d)
|
|
e.mu.Unlock()
|
|
}
|
|
|
|
func (e *Engine) untrack(d Download) {
|
|
e.mu.Lock()
|
|
defer e.mu.Unlock()
|
|
for i, x := range e.active {
|
|
if x == d {
|
|
e.active = append(e.active[:i], e.active[i+1:]...)
|
|
return
|
|
}
|
|
}
|
|
}
|