node build fixed
This commit is contained in:
ra_ma
@@ -0,0 +1,389 @@
|
||||
package image_downloader
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"image"
|
||||
_ "image/gif"
|
||||
_ "image/jpeg"
|
||||
_ "image/png"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"seanime/internal/util"
|
||||
"seanime/internal/util/limiter"
|
||||
"slices"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/goccy/go-json"
|
||||
"github.com/google/uuid"
|
||||
"github.com/rs/zerolog"
|
||||
_ "golang.org/x/image/bmp"
|
||||
_ "golang.org/x/image/tiff"
|
||||
_ "golang.org/x/image/webp"
|
||||
)
|
||||
|
||||
const (
|
||||
RegistryFilename = "registry.json"
|
||||
)
|
||||
|
||||
type (
|
||||
ImageDownloader struct {
|
||||
downloadDir string
|
||||
registry Registry
|
||||
cancelChannel chan struct{}
|
||||
logger *zerolog.Logger
|
||||
actionMu sync.Mutex
|
||||
registryMu sync.Mutex
|
||||
}
|
||||
|
||||
Registry struct {
|
||||
content *RegistryContent
|
||||
logger *zerolog.Logger
|
||||
downloadDir string
|
||||
registryPath string
|
||||
mu sync.Mutex
|
||||
}
|
||||
RegistryContent struct {
|
||||
UrlToId map[string]string `json:"url_to_id"`
|
||||
IdToUrl map[string]string `json:"id_to_url"`
|
||||
IdToExt map[string]string `json:"id_to_ext"`
|
||||
}
|
||||
)
|
||||
|
||||
func NewImageDownloader(downloadDir string, logger *zerolog.Logger) *ImageDownloader {
|
||||
_ = os.MkdirAll(downloadDir, os.ModePerm)
|
||||
|
||||
return &ImageDownloader{
|
||||
downloadDir: downloadDir,
|
||||
logger: logger,
|
||||
registry: Registry{
|
||||
logger: logger,
|
||||
registryPath: filepath.Join(downloadDir, RegistryFilename),
|
||||
downloadDir: downloadDir,
|
||||
content: &RegistryContent{},
|
||||
},
|
||||
cancelChannel: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
// DownloadImages downloads multiple images concurrently.
|
||||
func (id *ImageDownloader) DownloadImages(urls []string) (err error) {
|
||||
id.cancelChannel = make(chan struct{})
|
||||
|
||||
if err = id.registry.setup(); err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
rateLimiter := limiter.NewLimiter(1*time.Second, 10)
|
||||
var wg sync.WaitGroup
|
||||
for _, url := range urls {
|
||||
wg.Add(1)
|
||||
go func(url string) {
|
||||
defer wg.Done()
|
||||
select {
|
||||
case <-id.cancelChannel:
|
||||
id.logger.Warn().Msg("image downloader: Download process canceled")
|
||||
return
|
||||
default:
|
||||
rateLimiter.Wait()
|
||||
id.downloadImage(url)
|
||||
}
|
||||
}(url)
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
if err = id.registry.save(urls); err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
func (id *ImageDownloader) DeleteDownloads() {
|
||||
id.actionMu.Lock()
|
||||
defer id.actionMu.Unlock()
|
||||
|
||||
id.registryMu.Lock()
|
||||
defer id.registryMu.Unlock()
|
||||
|
||||
_ = os.RemoveAll(id.downloadDir)
|
||||
id.registry.content = &RegistryContent{}
|
||||
}
|
||||
|
||||
// CancelDownload cancels the download process.
|
||||
func (id *ImageDownloader) CancelDownload() {
|
||||
close(id.cancelChannel)
|
||||
}
|
||||
|
||||
func (id *ImageDownloader) GetImageFilenameByUrl(url string) (filename string, ok bool) {
|
||||
id.actionMu.Lock()
|
||||
defer id.actionMu.Unlock()
|
||||
|
||||
id.registryMu.Lock()
|
||||
defer id.registryMu.Unlock()
|
||||
|
||||
if err := id.registry.setup(); err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
var imgID string
|
||||
imgID, ok = id.registry.content.UrlToId[url]
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
filename = imgID + "." + id.registry.content.IdToExt[imgID]
|
||||
return
|
||||
}
|
||||
|
||||
// GetImageFilenamesByUrls returns a map of URLs to image filenames.
|
||||
//
|
||||
// e.g., {"url1": "filename1.png", "url2": "filename2.jpg"}
|
||||
func (id *ImageDownloader) GetImageFilenamesByUrls(urls []string) (ret map[string]string, err error) {
|
||||
id.actionMu.Lock()
|
||||
defer id.actionMu.Unlock()
|
||||
|
||||
id.registryMu.Lock()
|
||||
defer id.registryMu.Unlock()
|
||||
|
||||
ret = make(map[string]string)
|
||||
|
||||
if err = id.registry.setup(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
for _, url := range urls {
|
||||
imgID, ok := id.registry.content.UrlToId[url]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
|
||||
ret[url] = imgID + "." + id.registry.content.IdToExt[imgID]
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func (id *ImageDownloader) DeleteImagesByUrls(urls []string) (err error) {
|
||||
id.actionMu.Lock()
|
||||
defer id.actionMu.Unlock()
|
||||
|
||||
id.registryMu.Lock()
|
||||
defer id.registryMu.Unlock()
|
||||
|
||||
if err = id.registry.setup(); err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
for _, url := range urls {
|
||||
imgID, ok := id.registry.content.UrlToId[url]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
|
||||
err = os.Remove(filepath.Join(id.downloadDir, imgID+"."+id.registry.content.IdToExt[imgID]))
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
delete(id.registry.content.UrlToId, url)
|
||||
delete(id.registry.content.IdToUrl, imgID)
|
||||
delete(id.registry.content.IdToExt, imgID)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// downloadImage downloads an image from a URL.
|
||||
func (id *ImageDownloader) downloadImage(url string) {
|
||||
|
||||
defer util.HandlePanicInModuleThen("util/image_downloader/downloadImage", func() {
|
||||
})
|
||||
|
||||
if url == "" {
|
||||
id.logger.Warn().Msg("image downloader: Empty URL provided, skipping download")
|
||||
return
|
||||
}
|
||||
|
||||
// Check if the image has already been downloaded
|
||||
id.registryMu.Lock()
|
||||
if _, ok := id.registry.content.UrlToId[url]; ok {
|
||||
id.registryMu.Unlock()
|
||||
id.logger.Debug().Msgf("image downloader: Image from URL %s has already been downloaded", url)
|
||||
return
|
||||
}
|
||||
id.registryMu.Unlock()
|
||||
|
||||
// Download image from URL
|
||||
id.logger.Info().Msgf("image downloader: Downloading image from URL: %s", url)
|
||||
|
||||
imgID := uuid.NewString()
|
||||
|
||||
// Download the image
|
||||
resp, err := http.Get(url)
|
||||
if err != nil {
|
||||
id.logger.Error().Err(err).Msgf("image downloader: Failed to download image from URL %s", url)
|
||||
return
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
buf, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
id.logger.Error().Err(err).Msgf("image downloader: Failed to read image data from URL %s", url)
|
||||
return
|
||||
}
|
||||
|
||||
// Get the image format
|
||||
_, format, err := image.DecodeConfig(bytes.NewReader(buf))
|
||||
if err != nil {
|
||||
id.logger.Error().Err(err).Msgf("image downloader: Failed to decode image format from URL %s", url)
|
||||
return
|
||||
}
|
||||
|
||||
// Create the file
|
||||
filePath := filepath.Join(id.downloadDir, imgID+"."+format)
|
||||
file, err := os.Create(filePath)
|
||||
if err != nil {
|
||||
id.logger.Error().Err(err).Msgf("image downloader: Failed to create file for image %s", imgID)
|
||||
return
|
||||
}
|
||||
defer file.Close()
|
||||
|
||||
// Copy the image data to the file
|
||||
_, err = io.Copy(file, bytes.NewReader(buf))
|
||||
if err != nil {
|
||||
id.logger.Error().Err(err).Msgf("image downloader: Failed to write image data to file for image from %s", url)
|
||||
return
|
||||
}
|
||||
|
||||
// Update registry
|
||||
id.registryMu.Lock()
|
||||
id.registry.addUrl(imgID, url, format)
|
||||
id.registryMu.Unlock()
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
func (r *Registry) setup() (err error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
|
||||
defer util.HandlePanicInModuleThen("util/image_downloader/setup", func() {
|
||||
err = fmt.Errorf("image downloader: Failed to setup registry")
|
||||
})
|
||||
|
||||
if r.content.IdToUrl != nil && r.content.UrlToId != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
r.content.UrlToId = make(map[string]string)
|
||||
r.content.IdToUrl = make(map[string]string)
|
||||
r.content.IdToExt = make(map[string]string)
|
||||
|
||||
// Check if the registry exists
|
||||
_ = os.MkdirAll(filepath.Dir(r.registryPath), os.ModePerm)
|
||||
_, err = os.Stat(r.registryPath)
|
||||
if os.IsNotExist(err) {
|
||||
// Create the registry file
|
||||
err = os.WriteFile(r.registryPath, []byte("{}"), os.ModePerm)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
// Read the registry file
|
||||
file, err := os.Open(r.registryPath)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer file.Close()
|
||||
|
||||
// Decode the registry file if there is content
|
||||
if file != nil {
|
||||
r.logger.Debug().Msg("image downloader: Reading registry content")
|
||||
err = json.NewDecoder(file).Decode(&r.content)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
if r.content == nil {
|
||||
r.content = &RegistryContent{
|
||||
UrlToId: make(map[string]string),
|
||||
IdToUrl: make(map[string]string),
|
||||
IdToExt: make(map[string]string),
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// save verifies and saves the registry content.
|
||||
func (r *Registry) save(urls []string) (err error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
|
||||
defer util.HandlePanicInModuleThen("util/image_downloader/save", func() {
|
||||
err = fmt.Errorf("image downloader: Failed to save registry content")
|
||||
})
|
||||
|
||||
// Verify all images have been downloaded
|
||||
allDownloaded := true
|
||||
for _, url := range urls {
|
||||
if url == "" {
|
||||
continue
|
||||
}
|
||||
if _, ok := r.content.UrlToId[url]; !ok {
|
||||
allDownloaded = false
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if !allDownloaded {
|
||||
// Clean up downloaded images
|
||||
go func() {
|
||||
r.logger.Error().Msg("image downloader: Not all images have been downloaded, aborting")
|
||||
// Read the directory
|
||||
files, err := os.ReadDir(r.downloadDir)
|
||||
if err != nil {
|
||||
r.logger.Error().Err(err).Msg("image downloader: Failed to abort")
|
||||
return
|
||||
}
|
||||
// Delete all files that have been downloaded (are in the registry)
|
||||
for _, file := range files {
|
||||
fileNameWithoutExt := file.Name()[:len(file.Name())-len(filepath.Ext(file.Name()))]
|
||||
if url, ok := r.content.IdToUrl[fileNameWithoutExt]; ok && slices.Contains(urls, url) {
|
||||
err = os.Remove(filepath.Join(r.downloadDir, file.Name()))
|
||||
if err != nil {
|
||||
r.logger.Error().Err(err).Msgf("image downloader: Failed to delete file %s", file.Name())
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
return fmt.Errorf("image downloader: Not all images have been downloaded, operation aborted")
|
||||
}
|
||||
|
||||
data, err := json.Marshal(r.content)
|
||||
if err != nil {
|
||||
r.logger.Error().Err(err).Msg("image downloader: Failed to marshal registry content")
|
||||
}
|
||||
// Overwrite the registry file
|
||||
err = os.WriteFile(r.registryPath, data, 0644)
|
||||
if err != nil {
|
||||
r.logger.Error().Err(err).Msg("image downloader: Failed to write registry content")
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *Registry) addUrl(imgID, url, format string) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.content.UrlToId[url] = imgID
|
||||
r.content.IdToUrl[imgID] = url
|
||||
r.content.IdToExt[imgID] = format
|
||||
}
|
||||
Reference in New Issue
Block a user