Skip to content

Commit

Permalink
We always want thumbs index to start from zero as users are expecting…
Browse files Browse the repository at this point in the history
… this (#1262)

* We always want thumbs index to start from zero as users are expecting this

* fix test when run separately

* extract func
  • Loading branch information
mjh1 committed May 20, 2024
1 parent 78607b4 commit ac54bfd
Show file tree
Hide file tree
Showing 3 changed files with 104 additions and 25 deletions.
2 changes: 1 addition & 1 deletion handlers/ffmpeg/ffmpeg.go
Original file line number Diff line number Diff line change
Expand Up @@ -89,7 +89,7 @@ func (h *HandlersCollection) NewFile() httprouter.Handle {
if job.ThumbnailsTargetURL == nil {
return
}
if err := thumbnails.GenerateThumb(filename, content, job.ThumbnailsTargetURL); err != nil {
if err := thumbnails.GenerateThumb(filename, content, job.ThumbnailsTargetURL, 0); err != nil {
log.LogError(job.RequestID, "generate thumb failed", err, "in", path.Join(targetURLBase, filename), "out", job.ThumbnailsTargetURL)
}
}()
Expand Down
44 changes: 36 additions & 8 deletions thumbnails/thumbnails.go
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,18 @@ func getMediaManifest(requestID string, input string) (*m3u8.MediaPlaylist, erro
return mediaPlaylist, nil
}

func getSegmentOffset(mediaPlaylist *m3u8.MediaPlaylist) (int64, error) {
segments := mediaPlaylist.GetAllSegments()
if len(segments) < 1 {
return 0, fmt.Errorf("no segments found for")
}
segmentOffset, err := segmentIndex(path.Base(segments[0].URI))
if err != nil {
return 0, fmt.Errorf("failed to get segment index: %w", err)
}
return segmentOffset, nil
}

func GenerateThumbsVTT(requestID string, input string, output *url.URL) error {
// download and parse the manifest
mediaPlaylist, err := getMediaManifest(requestID, input)
Expand All @@ -69,11 +81,15 @@ func GenerateThumbsVTT(requestID string, input string, output *url.URL) error {
if err != nil {
return err
}
segmentOffset, err := getSegmentOffset(mediaPlaylist)
if err != nil {
return err
}

var currentTime time.Time
// loop through each segment, generate a vtt entry for it
for _, segment := range mediaPlaylist.GetAllSegments() {
filename, err := thumbFilename(path.Base(segment.URI))
filename, err := thumbFilename(path.Base(segment.URI), segmentOffset)
if err != nil {
return err
}
Expand Down Expand Up @@ -106,7 +122,7 @@ func GenerateThumbsVTT(requestID string, input string, output *url.URL) error {
return nil
}

func GenerateThumb(segmentURI string, input []byte, output *url.URL) error {
func GenerateThumb(segmentURI string, input []byte, output *url.URL, segmentOffset int64) error {
tempDir, err := os.MkdirTemp(os.TempDir(), "thumbs-*")
if err != nil {
return fmt.Errorf("failed to make temp dir: %w", err)
Expand All @@ -119,7 +135,7 @@ func GenerateThumb(segmentURI string, input []byte, output *url.URL) error {
return err
}

filename, err := thumbFilename(segmentURI)
filename, err := thumbFilename(segmentURI, segmentOffset)
if err != nil {
return err
}
Expand Down Expand Up @@ -159,6 +175,10 @@ func GenerateThumbsFromManifest(requestID, input string, output *url.URL) error
if err != nil {
return err
}
segmentOffset, err := getSegmentOffset(mediaPlaylist)
if err != nil {
return err
}

// parallelise the thumb uploads
uploadGroup, _ := errgroup.WithContext(context.Background())
Expand All @@ -185,7 +205,7 @@ func GenerateThumbsFromManifest(requestID, input string, output *url.URL) error
}

// generate thumbnail for the segment
return GenerateThumb(path.Base(segment.URI), bs, output)
return GenerateThumb(path.Base(segment.URI), bs, output, segmentOffset)
})
}
return uploadGroup.Wait()
Expand Down Expand Up @@ -218,15 +238,23 @@ func processSegment(input string, thumbOut string) error {

var segmentPrefix = []string{"index", "clip_"}

func thumbFilename(segmentURI string) (string, error) {
func segmentIndex(segmentURI string) (int64, error) {
// segmentURI will be indexX.ts or clip_X.ts
for _, prefix := range segmentPrefix {
segmentURI = strings.TrimPrefix(segmentURI, prefix)
}
index := strings.TrimSuffix(segmentURI, ".ts")
i, err := strconv.ParseInt(index, 10, 32)
i, err := strconv.ParseInt(index, 10, 64)
if err != nil {
return 0, fmt.Errorf("thumbFilename failed for %s: %w", segmentURI, err)
}
return i, nil
}

func thumbFilename(segmentURI string, segmentOffset int64) (string, error) {
i, err := segmentIndex(segmentURI)
if err != nil {
return "", fmt.Errorf("thumbFilename failed for %s: %w", segmentURI, err)
return "", err
}
return fmt.Sprintf("keyframes_%d.jpg", i), nil
return fmt.Sprintf("keyframes_%d.jpg", i-segmentOffset), nil
}
83 changes: 67 additions & 16 deletions thumbnails/thumbnails_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package thumbnails

import (
"context"
"fmt"
"net/url"
"os"
"path"
Expand All @@ -15,7 +16,7 @@ import (
func generateThumb(t *testing.T, filename string, out *url.URL) {
bs, err := os.ReadFile(filename)
require.NoError(t, err)
err = GenerateThumb(path.Base(filename), bs, out)
err = GenerateThumb(path.Base(filename), bs, out, 0)
require.NoError(t, err)
}

Expand All @@ -35,7 +36,7 @@ func TestGenerateThumbs(t *testing.T) {
generateThumb(t, path.Join(wd, "..", "test/fixtures/seg-1.ts"), out)
generateThumb(t, path.Join(wd, "..", "test/fixtures/seg-2.ts"), out)

testGenerateThumbsRun(t, outDir)
testGenerateThumbsRun(t, outDir, path.Join(wd, "..", "test/fixtures/tiny.m3u8"))

// Test the recording flow
outDir, err = os.MkdirTemp(os.TempDir(), "thumbs*")
Expand All @@ -47,17 +48,58 @@ func TestGenerateThumbs(t *testing.T) {
err = GenerateThumbsFromManifest("req ID", path.Join(wd, "..", "test/fixtures/tiny.m3u8"), out)
require.NoError(t, err)

testGenerateThumbsRun(t, outDir)
testGenerateThumbsRun(t, outDir, path.Join(wd, "..", "test/fixtures/tiny.m3u8"))
}

func testGenerateThumbsRun(t *testing.T, outDir string) {
func TestGenerateThumbsOffset(t *testing.T) {
// Test manifest with an segment offset i.e. the first segment has index >0

wd, err := os.Getwd()
require.NoError(t, err)
outDir, err := os.MkdirTemp(os.TempDir(), "thumbs*")
require.NoError(t, err)
defer os.RemoveAll(outDir)
out, err := url.Parse(outDir)
require.NoError(t, err)
err = os.Mkdir(path.Join(outDir, "in"), 0755)
require.NoError(t, err)

wd, err := os.Getwd()
// Generate a manifest with an offset
inputFile := path.Join(outDir, "in", "index.m3u8")
manifest := `
#EXTM3U
#EXT-X-VERSION:3
#EXT-X-MEDIA-SEQUENCE:0
#EXT-X-TARGETDURATION:10
#EXTINF:10.000000,
clip_100.ts
#EXTINF:10.000000,
clip_101.ts
#EXTINF:10.000000,
clip_102.ts
#EXT-X-ENDLIST
`
err = os.WriteFile(inputFile, []byte(manifest), 0644)
require.NoError(t, err)
// copy segments
for i := 0; i < 3; i++ {
b, err := os.ReadFile(path.Join(wd, "..", fmt.Sprintf("test/fixtures/seg-%d.ts", i)))
require.NoError(t, err)
err = os.WriteFile(path.Join(outDir, "in", fmt.Sprintf("clip_%d.ts", i+100)), b, 0644)
require.NoError(t, err)
}

err = GenerateThumbsFromManifest("req ID", inputFile, out)
require.NoError(t, err)

testGenerateThumbsRun(t, outDir, inputFile)
}

err = GenerateThumbsVTT("req ID", path.Join(wd, "..", "test/fixtures/tiny.m3u8"), out)
func testGenerateThumbsRun(t *testing.T, outDir, input string) {
out, err := url.Parse(outDir)
require.NoError(t, err)

err = GenerateThumbsVTT("req ID", input, out)
require.NoError(t, err)

expectedVtt := `WEBVTT
Expand Down Expand Up @@ -92,24 +134,33 @@ keyframes_2.jpg

func Test_thumbFilename(t *testing.T) {
tests := []struct {
name string
segmentURI string
want string
name string
segmentURI string
segmentOffset int64
want string
}{
{
name: "index",
segmentURI: "index0.ts",
want: "keyframes_0.jpg",
name: "index",
segmentURI: "index0.ts",
segmentOffset: 0,
want: "keyframes_0.jpg",
},
{
name: "clip",
segmentURI: "clip_1.ts",
segmentOffset: 0,
want: "keyframes_1.jpg",
},
{
name: "clip",
segmentURI: "clip_1.ts",
want: "keyframes_1.jpg",
name: "clip",
segmentURI: "clip_101.ts",
segmentOffset: 100,
want: "keyframes_1.jpg",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got, err := thumbFilename(tt.segmentURI)
got, err := thumbFilename(tt.segmentURI, tt.segmentOffset)
require.NoError(t, err)
require.Equal(t, tt.want, got)
})
Expand Down

0 comments on commit ac54bfd

Please sign in to comment.