From ee3d8916ed81cae1ad73a9886c4f9f66b01fffd3 Mon Sep 17 00:00:00 2001
From: Zoe Roux <zoe.roux@zoriya.dev>
Date: Mon, 19 Feb 2024 14:27:51 +0100
Subject: [PATCH] Use cmap for thumbnails

---
 transcoder/main.go           | 17 ++++----
 transcoder/src/thumbnails.go | 84 +++++++++++++-----------------------
 2 files changed, 39 insertions(+), 62 deletions(-)

diff --git a/transcoder/main.go b/transcoder/main.go
index 1a548dd3..af4983a0 100644
--- a/transcoder/main.go
+++ b/transcoder/main.go
@@ -187,9 +187,10 @@ func (h *Handler) GetInfo(c echo.Context) error {
 	}
 	// Run extractors to have them in cache
 	src.Extract(ret.Path, sha, route)
-	go h.thumbnails.ExtractThumbnail(
+	go src.ExtractThumbnail(
 		ret.Path,
-		fmt.Sprintf("%s/thumbnails.png", route),
+		route,
+		&sha,
 	)
 	return c.JSON(http.StatusOK, ret)
 }
@@ -265,9 +266,10 @@ func (h *Handler) GetThumbnails(c echo.Context) error {
 		return err
 	}
 
-	out, err := h.thumbnails.ExtractThumbnail(
+	out, err := src.ExtractThumbnail(
 		path,
-		fmt.Sprintf("%s/thumbnails.png", GetRoute(c)),
+		GetRoute(c),
+		nil,
 	)
 	if err != nil {
 		return err
@@ -288,9 +290,10 @@ func (h *Handler) GetThumbnailsVtt(c echo.Context) error {
 		return err
 	}
 
-	out, err := h.thumbnails.ExtractThumbnail(
+	out, err := src.ExtractThumbnail(
 		path,
-		fmt.Sprintf("%s/thumbnails.png", GetRoute(c)),
+		GetRoute(c),
+		nil,
 	)
 	if err != nil {
 		return err
@@ -301,7 +304,6 @@ func (h *Handler) GetThumbnailsVtt(c echo.Context) error {
 
 type Handler struct {
 	transcoder *src.Transcoder
-	thumbnails *src.ThumbnailsCreator
 }
 
 func main() {
@@ -316,7 +318,6 @@ func main() {
 	}
 	h := Handler{
 		transcoder: transcoder,
-		thumbnails: src.NewThumbnailsCreator(),
 	}
 
 	e.GET("/direct", DirectStream)
diff --git a/transcoder/src/thumbnails.go b/transcoder/src/thumbnails.go
index aed62a85..993e8cc4 100644
--- a/transcoder/src/thumbnails.go
+++ b/transcoder/src/thumbnails.go
@@ -20,71 +20,47 @@ var default_interval = 10
 // Setting this too high allows really long processing times.
 var max_numcaps = 150
 
-type ThumbnailsCreator struct {
-	created map[string]string
-	running map[string]chan struct{}
-	lock    sync.Mutex
+type Thumbnail struct {
+	ready sync.WaitGroup
+	path  string
 }
 
-func NewThumbnailsCreator() *ThumbnailsCreator {
-	return &ThumbnailsCreator{
-		created: make(map[string]string),
-		running: make(map[string]chan struct{}),
-	}
+var thumbnails = NewCMap[string, *Thumbnail]()
+
+func ExtractThumbnail(path string, route string, sha *string) (string, error) {
+	ret, _ := thumbnails.GetOrCreate(path, func() *Thumbnail {
+		if sha == nil {
+			nsha, err := GetHash(path)
+			if err != nil {
+				return nil
+			}
+			sha = &nsha
+		}
+		ret := &Thumbnail{
+			path: fmt.Sprintf("%s/%s", Settings.Metadata, sha),
+		}
+		ret.ready.Add(1)
+		go extractThumbnail(path, ret.path, fmt.Sprintf("%s/thumbnails.png", route))
+		return ret
+	})
+	ret.ready.Wait()
+	return ret.path, nil
 }
 
-func (t *ThumbnailsCreator) ExtractThumbnail(path string, name string) (string, error) {
-	t.lock.Lock()
-	defer t.lock.Unlock()
-
-	out, ok := t.created[path]
-	if ok {
-		return out, nil
-	}
-
-	wait, ok := t.running[path]
-	if ok {
-		t.lock.Unlock()
-		<-wait
-		t.lock.Lock()
-		out = t.created[path]
-		return out, nil
-	}
-	wait = make(chan struct{})
-	t.running[path] = wait
-
-	t.lock.Unlock()
-
-	out, err := extractThumbnail(path, name)
-	t.lock.Lock()
-	// this will be unlocked by the defer at the top of the function.
-
-	delete(t.running, path)
-	if err == nil {
-		t.created[path] = out
-	}
-	return out, err
-}
-
-func extractThumbnail(path string, name string) (string, error) {
+func extractThumbnail(path string, out string, name string) error {
 	defer printExecTime("extracting thumbnails for %s", path)()
-	sha, err := GetHash(path)
-	if err != nil {
-		return "", err
-	}
-	out := fmt.Sprintf("%s/%s", Settings.Metadata, sha)
 	os.MkdirAll(out, 0o755)
 	sprite_path := fmt.Sprintf("%s/sprite.png", out)
 	vtt_path := fmt.Sprintf("%s/sprite.vtt", out)
 
 	if _, err := os.Stat(sprite_path); err == nil {
-		return out, nil
+		return nil
 	}
 
 	gen, err := screengen.NewGenerator(path)
 	if err != nil {
 		log.Printf("Error reading video file: %v", err)
-		return "", err
+		return err
 	}
 	defer gen.Close()
 
@@ -115,7 +91,7 @@ func extractThumbnail(path string, name string) (string, error) {
 		img, err := gen.ImageWxH(int64(ts*1000), width, height)
 		if err != nil {
 			log.Printf("Could not generate screenshot %s", err)
-			return "", err
+			return err
 		}
 
 		x := (i % columns) * width
@@ -138,13 +114,13 @@ func extractThumbnail(path string, name string) (string, error) {
 
 	err = os.WriteFile(vtt_path, []byte(vtt), 0o644)
 	if err != nil {
-		return "", err
+		return err
 	}
 	err = imaging.Save(sprite, sprite_path)
 	if err != nil {
-		return "", err
+		return err
 	}
-	return out, nil
+	return nil
 }
 
 func tsToVttTime(ts int) string {