ddc/internal/jobs/manager.go

package jobs

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"errors"
	"fmt"
	"os"
	"path"
	"path/filepath"
	"sort"
	"sync"
	"time"

	"ddc/internal/config"
	"ddc/internal/diskspace"
	"ddc/internal/fetch"
	"ddc/internal/files"
	"ddc/internal/guard"
	"ddc/internal/libraries"
	"ddc/internal/preferences"
)

var (
	ErrJobNotFound    = errors.New("download not found")
	ErrJobBusy        = errors.New("download is busy, try again when it is ready")
	ErrInvalidRequest = errors.New("invalid request")
)

type Manager struct {
	settings    config.Settings
	preferences *preferences.Store
	policy      *guard.Policy
	downloader  *fetch.Downloader
	rootContext context.Context
	cancelAll   context.CancelFunc
	slots       slotLimiter
	mutex       sync.RWMutex
	jobs        map[string]*Job
}

func NewManager(settings config.Settings, store *preferences.Store) *Manager {
	policy := guard.NewPolicy(store.Get().AllowedPorts)
	rootContext, cancelAll := context.WithCancel(context.Background())
	return &Manager{
		settings:    settings,
		preferences: store,
		policy:      policy,
		downloader:  fetch.NewDownloader(policy),
		rootContext: rootContext,
		cancelAll:   cancelAll,
		jobs:        make(map[string]*Job),
	}
}

func (manager *Manager) Destinations() []libraries.Destination {
	return libraries.Discover(manager.settings.LibraryRoots)
}

func (manager *Manager) Preferences() preferences.Preferences {
	return manager.preferences.Get()
}

func (manager *Manager) UpdatePreferences(next preferences.Preferences) error {
	if err := manager.preferences.Update(next); err != nil {
		return err
	}
	manager.policy.SetAllowedPorts(next.AllowedPorts)
	return nil
}

func (manager *Manager) concurrencyLimit() int {
	return manager.preferences.Get().MaxConcurrentDownloads
}

func (manager *Manager) reservedBytes() int64 {
	return manager.preferences.Get().MinFreeBytes()
}

func (manager *Manager) AvailableBytes() (int64, error) {
	return diskspace.Available(manager.settings.DownloadDir)
}

func (manager *Manager) Create(rawURL string, autoExtract bool) (Snapshot, error) {
	target, err := manager.policy.ValidateURL(rawURL)
	if err != nil {
		return Snapshot{}, invalid(err)
	}
	if err := diskspace.EnsureFree(manager.settings.DownloadDir, 0, manager.reservedBytes()); err != nil {
		return Snapshot{}, err
	}

	id, err := newJobID()
	if err != nil {
		return Snapshot{}, err
	}
	directory := filepath.Join(manager.settings.DownloadDir, id)
	if err := os.Mkdir(directory, 0o755); err != nil {
		return Snapshot{}, err
	}

	job := manager.newJob(id, directory, time.Now())
	job.sourceURL = target.String()
	job.displayURL = target.Redacted()
	job.autoExtract = autoExtract
	job.status = StatusQueued
	job.fileName = initialFileName(target.Path, target.Host)
	job.workers.Add(1)

	manager.store(job)
	go manager.download(job)
	return job.snapshot(), nil
}

func (manager *Manager) List() []Snapshot {
	manager.mutex.RLock()
	snapshots := make([]Snapshot, 0, len(manager.jobs))
	for _, job := range manager.jobs {
		snapshots = append(snapshots, job.snapshot())
	}
	manager.mutex.RUnlock()

	sort.Slice(snapshots, func(first, second int) bool {
		return snapshots[first].CreatedAt.After(snapshots[second].CreatedAt)
	})
	return snapshots
}

func (manager *Manager) Files(id string) ([]files.Node, error) {
	job, err := manager.lookup(id)
	if err != nil {
		return nil, err
	}
	if job.currentStatus() != StatusReady {
		return nil, ErrJobBusy
	}
	return files.BuildTree(job.directory)
}

func (manager *Manager) Delete(id string) error {
	job, err := manager.lookup(id)
	if err != nil {
		return err
	}
	if !job.markDeleted() {
		return ErrJobBusy
	}

	job.cancel()
	job.workers.Wait()

	manager.mutex.Lock()
	delete(manager.jobs, id)
	manager.mutex.Unlock()

	return os.RemoveAll(job.directory)
}

func (manager *Manager) CancelAll() {
	manager.cancelAll()
}

func (manager *Manager) newJob(id, directory string, createdAt time.Time) *Job {
	jobContext, cancel := context.WithCancel(manager.rootContext)
	return &Job{
		id:         id,
		directory:  directory,
		createdAt:  createdAt,
		jobContext: jobContext,
		cancel:     cancel,
	}
}

func (manager *Manager) store(job *Job) {
	manager.mutex.Lock()
	defer manager.mutex.Unlock()
	manager.jobs[job.id] = job
}

func (manager *Manager) lookup(id string) (*Job, error) {
	manager.mutex.RLock()
	defer manager.mutex.RUnlock()
	job, found := manager.jobs[id]
	if !found {
		return nil, ErrJobNotFound
	}
	return job, nil
}

func newJobID() (string, error) {
	buffer := make([]byte, 16)
	if _, err := rand.Read(buffer); err != nil {
		return "", err
	}
	return hex.EncodeToString(buffer), nil
}

func initialFileName(urlPath, host string) string {
	if name := fetch.SanitizeFileName(path.Base(urlPath)); name != "" {
		return name
	}
	return host
}

func invalid(err error) error {
	return fmt.Errorf("%w: %w", ErrInvalidRequest, err)
}