Elijah
724d70e58b
Some checks failed
Automated Container Build / build-and-push (push) Failing after 12s
133 lines
3.1 KiB
Go
133 lines
3.1 KiB
Go
package workers
|
|
|
|
import (
|
|
"fmt"
|
|
"log"
|
|
"os"
|
|
"os/exec"
|
|
"path/filepath"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"git.elijahkuntz.com/Elijah/drive/config"
|
|
"git.elijahkuntz.com/Elijah/drive/database"
|
|
"github.com/disintegration/imaging"
|
|
)
|
|
|
|
type ThumbnailJob struct {
|
|
FilePath string
|
|
Checksum string
|
|
MimeType string
|
|
}
|
|
|
|
type ThumbnailManager struct {
|
|
Queue chan ThumbnailJob
|
|
Config *config.Config
|
|
DB *database.DB
|
|
wg sync.WaitGroup
|
|
}
|
|
|
|
var (
|
|
Instance *ThumbnailManager
|
|
once sync.Once
|
|
)
|
|
|
|
func InitThumbnailManager(cfg *config.Config, db *database.DB) {
|
|
once.Do(func() {
|
|
Instance = &ThumbnailManager{
|
|
Queue: make(chan ThumbnailJob, 1000),
|
|
Config: cfg,
|
|
DB: db,
|
|
}
|
|
Instance.Start(2) // Run 2 concurrent workers
|
|
})
|
|
}
|
|
|
|
func (m *ThumbnailManager) Start(workers int) {
|
|
for i := 0; i < workers; i++ {
|
|
m.wg.Add(1)
|
|
go m.worker(i)
|
|
}
|
|
}
|
|
|
|
func (m *ThumbnailManager) Stop() {
|
|
close(m.Queue)
|
|
m.wg.Wait()
|
|
}
|
|
|
|
func (m *ThumbnailManager) Enqueue(job ThumbnailJob) {
|
|
if job.Checksum == "" {
|
|
return // Cannot store thumbnail without checksum
|
|
}
|
|
select {
|
|
case m.Queue <- job:
|
|
default:
|
|
log.Println("Thumbnail queue is full, dropping job for", job.FilePath)
|
|
}
|
|
}
|
|
|
|
func (m *ThumbnailManager) worker(id int) {
|
|
defer m.wg.Done()
|
|
for job := range m.Queue {
|
|
// Check if thumbnail generation is enabled
|
|
if val, err := m.DB.GetSetting("thumbnail_images"); err == nil && val == "false" && strings.HasPrefix(job.MimeType, "image/") {
|
|
continue
|
|
}
|
|
if val, err := m.DB.GetSetting("thumbnail_videos"); err == nil && val == "false" && strings.HasPrefix(job.MimeType, "video/") {
|
|
continue
|
|
}
|
|
|
|
destPath := filepath.Join(m.Config.ThumbnailDir, job.Checksum+".jpg")
|
|
|
|
// Skip if thumbnail already exists
|
|
if _, err := os.Stat(destPath); err == nil {
|
|
continue
|
|
}
|
|
|
|
fullSourcePath := filepath.Join(m.Config.StorageDir, filepath.FromSlash(job.FilePath))
|
|
|
|
var err error
|
|
if strings.HasPrefix(job.MimeType, "image/") {
|
|
err = generateImageThumbnail(fullSourcePath, destPath)
|
|
} else if strings.HasPrefix(job.MimeType, "video/") {
|
|
err = generateVideoThumbnail(fullSourcePath, destPath)
|
|
}
|
|
|
|
if err != nil {
|
|
log.Printf("[Worker %d] Failed to generate thumbnail for %s: %v\n", id, job.FilePath, err)
|
|
} else {
|
|
log.Printf("[Worker %d] Generated thumbnail for %s\n", id, job.FilePath)
|
|
}
|
|
}
|
|
}
|
|
|
|
func generateImageThumbnail(src, dest string) error {
|
|
img, err := imaging.Open(src, imaging.AutoOrientation(true))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
// Resize to fit within 256x256 while preserving aspect ratio
|
|
thumb := imaging.Fit(img, 256, 256, imaging.Lanczos)
|
|
return imaging.Save(thumb, dest, imaging.JPEGQuality(80))
|
|
}
|
|
|
|
func generateVideoThumbnail(src, dest string) error {
|
|
// Extract a frame at 1 second mark, scale to max width/height 256
|
|
cmd := exec.Command("ffmpeg", "-y", "-i", src, "-ss", "00:00:01.000", "-vframes", "1", "-vf", "scale=256:256:force_original_aspect_ratio=decrease", dest)
|
|
|
|
// Set a timeout to prevent hanging on corrupted videos
|
|
timer := time.AfterFunc(15*time.Second, func() {
|
|
if cmd.Process != nil {
|
|
cmd.Process.Kill()
|
|
}
|
|
})
|
|
|
|
err := cmd.Run()
|
|
timer.Stop()
|
|
|
|
if err != nil {
|
|
return fmt.Errorf("ffmpeg failed: %w", err)
|
|
}
|
|
return nil
|
|
}
|