am-downloader/downloader.go

package main

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

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.download(id)
			}
		}()
	}
}

func (a *App) download(id uint) {
	var t Task
	if a.db.First(&t, id).Error != nil || t.Status != "pending" {
		return
	}
	a.db.Model(&t).Updates(map[string]any{"status": "resolving", "error": "", "progress": 0})
	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, ""
	asset, reused, err := a.acquireAsset(&t)
	if err != nil {
		a.fail(&t, err)
		return
	}
	if reused {
		a.finishExisting(&t, asset.Size)
		return
	}
	a.db.Model(&t).Updates(map[string]any{"format": t.Format, "local_path": t.LocalPath})
	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"
	a.db.Save(t)
	// Existing local media is reused and therefore does not consume a credit.
	if err := a.refundCredit(t); err != nil {
		a.fail(t, err)
		return
	}
	a.upload(t)
}

func (a *App) downloadFile(t *Task, info playInfo) {
	a.db.Model(t).Updates(map[string]any{"status": "downloading", "format": t.Format, "local_path": t.LocalPath})
	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, t.TotalBytes, t.DownloadedBytes, t.Progress = "downloading", resp.ContentLength, 0, 0
	a.db.Save(t)

	tmp := t.LocalPath + ".part"
	f, err := os.Create(tmp)
	if err != nil {
		a.fail(t, err)
		return
	}
	n, err := a.copyWithProgress(f, resp.Body, t)
	closeErr := f.Close()
	if err != nil {
		a.fail(t, err)
		return
	}
	if closeErr != nil {
		a.fail(t, closeErr)
		return
	}
	if t.TotalBytes > 0 && n != t.TotalBytes {
		a.fail(t, fmt.Errorf("incomplete download: got %d of %d bytes", n, t.TotalBytes))
		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.Format, t.LocalPath, t.ActualQuality, t.DownloadedBytes = actual, finalPath, actualQuality(t.Quality, t.ResolvedQuality, actual), n
	if t.TotalBytes > 0 {
		t.Progress = 100
	}
	a.db.Save(t)
	a.markAssetDone(t)
	// The credit pays for the successful media download, not the later WebDAV sync.
	// Commit before upload so a storage outage cannot turn a successful download into a free retry.
	if err := a.commitCredit(t); err != nil {
		a.fail(t, err)
		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
				a.db.Model(t).Updates(map[string]any{"downloaded_bytes": n, "progress": progress})
				last = n
			}
		}
		if er == io.EOF {
			return n, nil
		}
		if er != nil {
			return n, er
		}
	}
}