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 }