am-downloader/downloader.go

package main

import (
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"
)

func (a *App) startWorkers() {
	for i := 0; i < a.cfg.WorkerCount; i++ {
		a.wg.Add(1)
		go func() {
			defer a.wg.Done()
			for id := range a.queue {
				a.processTask(id)
			}
		}()
	}
}

func (a *App) processTask(id uint) {
	var t Task
	if a.db.First(&t, id).Error != nil {
		return
	}
	if t.Status == "upload_pending" {
		a.upload(&t)
		return
	}
	if t.Status != "pending" {
		return
	}
	a.download(id)
}

func (a *App) download(id uint) {
	t, err := a.claimTask(id)
	if err != nil {
		return
	}

	info, err := a.am.Resolve(t.Source, t.RID, t.Quality)
	if err != nil {
		a.fail(t, err)
		return
	}

	ext := normalizedFormat(info.Format)
	if ext == "" {
		ext = normalizedFormat(filepath.Ext(info.URL))
	}
	if ext == "" {
		ext = "mp3"
	}
	t.Format, t.ResolvedQuality, t.ActualQuality = ext, info.Level, ""
	if err := a.db.Model(t).Updates(map[string]any{
		"format": t.Format, "resolved_quality": t.ResolvedQuality,
	}).Error; err != nil {
		a.fail(t, err)
		return
	}

	asset, reused, err := a.acquireAsset(t)
	if err != nil {
		a.fail(t, err)
		return
	}
	if reused {
		a.finishExisting(t, asset.Size)
		return
	}
	if err := a.db.Model(t).Updates(map[string]any{
		"local_path": t.LocalPath, "status": "downloading",
	}).Error; err != nil {
		a.invalidateAsset(t)
		a.fail(t, err)
		return
	}
	a.downloadFile(t, info)
}

func (a *App) finishExisting(t *Task, size int64) {
	actual, err := detectAudioFormat(t.LocalPath)
	if err == nil {
		actual = normalizedFormat(actual)
		if actual != t.Format {
			if renamed, renameErr := renameToActualFormat(t.LocalPath, actual); renameErr == nil {
				t.LocalPath = renamed
			} else {
				a.fail(t, renameErr)
				return
			}
		}
		t.Format = actual
		t.ActualQuality = actualQuality(t.Quality, t.ResolvedQuality, actual)
	}
	t.TotalBytes, t.DownloadedBytes, t.Progress, t.Status = size, size, 100, "done"
	if err := a.db.Save(t).Error; err != nil {
		return
	}
	if err := a.refundCredit(t); err != nil {
		a.fail(t, err)
		return
	}
	a.upload(t)
}

func (a *App) downloadFile(t *Task, info playInfo) {
	req, err := http.NewRequest("GET", info.URL, nil)
	if err != nil {
		a.fail(t, err)
		return
	}
	for k, v := range info.Headers {
		req.Header.Set(k, v)
	}
	resp, err := a.httpClient.Do(req)
	if err != nil {
		a.fail(t, err)
		return
	}
	defer resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		a.fail(t, fmt.Errorf("media http %d", resp.StatusCode))
		return
	}

	t.Status = "downloading"
	t.TotalBytes = resp.ContentLength
	t.DownloadedBytes = 0
	t.Progress = 0
	_ = a.db.Model(t).Updates(map[string]any{
		"status": t.Status, "total_bytes": t.TotalBytes,
		"downloaded_bytes": 0, "progress": 0,
	}).Error

	if err := os.MkdirAll(filepath.Dir(t.LocalPath), 0755); err != nil {
		a.fail(t, err)
		return
	}
	tmp := fmt.Sprintf("%s.task-%d.part", t.LocalPath, t.ID)
	_ = os.Remove(tmp)
	f, err := os.Create(tmp)
	if err != nil {
		a.fail(t, err)
		return
	}

	stop := make(chan struct{})
	ticker := time.NewTicker(30 * time.Second)
	go func() {
		for {
			select {
			case <-ticker.C:
				a.markAssetHeartbeat(t.AssetID, t.ID)
			case <-stop:
				return
			}
		}
	}()
	n, copyErr := a.copyWithProgress(f, resp.Body, t)
	closeErr := f.Close()
	ticker.Stop()
	close(stop)

	if copyErr != nil {
		a.fail(t, copyErr)
		_ = os.Remove(tmp)
		return
	}
	if closeErr != nil {
		a.fail(t, closeErr)
		_ = os.Remove(tmp)
		return
	}
	if t.TotalBytes > 0 && n != t.TotalBytes {
		a.fail(t, fmt.Errorf("incomplete download: got %d of %d bytes", n, t.TotalBytes))
		_ = os.Remove(tmp)
		return
	}
	if !a.assertAssetOwner(t) {
		_ = os.Remove(tmp)
		return
	}

	actual, err := detectAudioFormat(tmp)
	if err != nil {
		a.fail(t, fmt.Errorf("detect audio format: %w", err))
		_ = os.Remove(tmp)
		return
	}
	actual = normalizedFormat(actual)
	finalPath := strings.TrimSuffix(t.LocalPath, filepath.Ext(t.LocalPath)) + "." + actual
	if err := os.Rename(tmp, finalPath); err != nil {
		a.fail(t, err)
		_ = os.Remove(tmp)
		return
	}
	t.LocalPath = finalPath
	t.Format = actual
	t.ActualQuality = actualQuality(t.Quality, t.ResolvedQuality, actual)
	t.DownloadedBytes = n
	t.Progress = 100
	if err := a.db.Model(t).Updates(map[string]any{
		"format": t.Format, "local_path": t.LocalPath,
		"actual_quality":   t.ActualQuality,
		"downloaded_bytes": n, "progress": 100,
	}).Error; err != nil {
		a.invalidateAsset(t)
		return
	}
	if !a.finishAsset(t) {
		a.fail(t, fmt.Errorf("shared asset ownership lost before completion"))
		return
	}
	if err := a.commitCredit(t); err != nil {
		a.fail(t, err)
		return
	}
	t.Status, t.Error = "upload_pending", ""
	if err := a.db.Model(t).Updates(map[string]any{
		"status": "upload_pending", "error": "",
	}).Error; err != nil {
		return
	}
	a.upload(t)
}
func (a *App) copyWithProgress(dst *os.File, src io.Reader, t *Task) (int64, error) {
	var n, last int64
	buf := make([]byte, 256*1024)
	for {
		nr, er := src.Read(buf)
		if nr > 0 {
			wn, ew := dst.Write(buf[:nr])
			n += int64(wn)
			if ew != nil {
				return n, ew
			}
			if n-last >= 1024*1024 {
				progress := 0.0
				if t.TotalBytes > 0 {
					progress = float64(n) * 100 / float64(t.TotalBytes)
				}
				t.DownloadedBytes, t.Progress = n, progress
				if err := a.db.Model(t).Updates(map[string]any{
					"downloaded_bytes": n, "progress": progress,
				}).Error; err != nil {
					return n, err
				}
				a.markAssetHeartbeat(t.AssetID, t.ID)
				last = n
			}
		}
		if er == io.EOF {
			return n, nil
		}
		if er != nil {
			return n, er
		}
	}
}