back/internal/v2/lib/downloader/v2.go

package downloader

import (
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"recook/tools"
	"sync"
)

type worker2 interface {
	GetLocalName(raw Item) string
	Download(file string, raw Item) error
}

type Item struct {
	Name string
	Raw  string
}

type downloader2 struct {
	worker    worker2
	workerNum int
	Queue     chan Item
	group     sync.WaitGroup
	file      string
}

func (o *downloader2) AddTask(raw Item) {
	o.Queue <- raw
}

func (o *downloader2) Start() {
	go func() {
		for i := 0; i < o.workerNum; i++ {
			o.group.Add(1)
			go o.forever()
		}
	}()
}

func (o *downloader2) GetPath(raw Item) string {
	return filepath.Join(o.file, o.worker.GetLocalName(raw))
}

func (o *downloader2) forever() {
	for {
		if value, ok := <-o.Queue; ok {
			file := filepath.Join(o.file)
			o.worker.Download(file, value)
		} else {
			break
		}
	}
	o.group.Done()
}

func (o *downloader2) Wait() {
	close(o.Queue)
	o.group.Wait()
}

func New2(worker worker2, file string, workerNum int) *downloader2 {
	if worker == nil {
		worker = defaultWorker2{}
	}
	return &downloader2{
		worker:    worker,
		workerNum: workerNum,
		Queue:     make(chan Item, 10000),
		file:      file,
	}
}

type defaultWorker2 struct {
}

func (o defaultWorker2) GetLocalName(raw Item) string {
	if raw.Name != "" {
		return raw.Name
	} else {
		res := filepath.Ext(raw.Raw)
		return filepath.Join(tools.MD5(raw.Raw) + res)
	}
}

func (o defaultWorker2) Download(file string, raw Item) error {
	allPath := filepath.Join(file, o.GetLocalName(raw))
	if !Exists(file) {
		_ = os.MkdirAll(file, os.ModePerm)
	}
	if Exists(allPath) {
		return nil
	}
	resp, err := http.Get(raw.Raw)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	data, err := io.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	err = os.WriteFile(allPath, data, os.ModePerm)
	if err != nil {
		fmt.Println(err.Error(), "===")
	}
	return err
}