Implemented downloader context

This commit is contained in:
bytedream 2022-02-13 15:03:53 +01:00
parent 0a4b9ec96e
commit d5fc8824cf

View file

@ -2,6 +2,7 @@ package crunchyroll
import ( import (
"bufio" "bufio"
"context"
"fmt" "fmt"
"github.com/grafov/m3u8" "github.com/grafov/m3u8"
"io/ioutil" "io/ioutil"
@ -52,20 +53,23 @@ func (f *Format) Download(downloader Downloader) error {
return err return err
} }
if err := download(f, downloader.TempDir, downloader.Goroutines, downloader.OnSegmentDownload); err != nil { if downloader.DeleteTempAfter {
defer os.RemoveAll(downloader.TempDir)
}
if err := download(downloader.Context, f, downloader.TempDir, downloader.Goroutines, downloader.OnSegmentDownload); err != nil {
return err return err
} }
if downloader.FFmpeg { if downloader.FFmpeg {
return mergeSegmentsFFmpeg(downloader.TempDir, downloader.Filename) return mergeSegmentsFFmpeg(downloader.Context, downloader.TempDir, downloader.Filename)
} else { } else {
return mergeSegments(downloader.TempDir, downloader.Filename) return mergeSegments(downloader.Context, downloader.TempDir, downloader.Filename)
} }
} }
// mergeSegments reads every file in tempDir and writes their content to the outputFile. // mergeSegments reads every file in tempDir and writes their content to the outputFile.
// The given output file gets created or overwritten if already existing // The given output file gets created or overwritten if already existing
func mergeSegments(tempDir string, outputFile string) error { func mergeSegments(context context.Context, tempDir string, outputFile string) error {
dir, err := os.ReadDir(tempDir) dir, err := os.ReadDir(tempDir)
if err != nil { if err != nil {
return err return err
@ -92,6 +96,12 @@ func mergeSegments(tempDir string, outputFile string) error {
}) })
for _, file := range dir { for _, file := range dir {
select {
case <-context.Done():
return context.Err()
default:
}
bodyAsBytes, err := ioutil.ReadFile(filepath.Join(tempDir, file.Name())) bodyAsBytes, err := ioutil.ReadFile(filepath.Join(tempDir, file.Name()))
if err != nil { if err != nil {
return err return err
@ -106,7 +116,7 @@ func mergeSegments(tempDir string, outputFile string) error {
// mergeSegmentsFFmpeg reads every file in tempDir and merges their content to the outputFile // mergeSegmentsFFmpeg reads every file in tempDir and merges their content to the outputFile
// with ffmpeg (https://ffmpeg.org/). // with ffmpeg (https://ffmpeg.org/).
// The given output file gets created or overwritten if already existing // The given output file gets created or overwritten if already existing
func mergeSegmentsFFmpeg(tempDir string, outputFile string) error { func mergeSegmentsFFmpeg(context context.Context, tempDir string, outputFile string) error {
dir, err := os.ReadDir(tempDir) dir, err := os.ReadDir(tempDir)
if err != nil { if err != nil {
return err return err
@ -125,5 +135,19 @@ func mergeSegmentsFFmpeg(tempDir string, outputFile string) error {
"-i", f.Name(), "-i", f.Name(),
"-c", "copy", "-c", "copy",
outputFile) outputFile)
return cmd.Run() if err := cmd.Start(); err != nil {
return err
}
cmdChan := make(chan error)
go func() {
cmdChan <- cmd.Wait()
}()
select {
case err := <-cmdChan:
return err
case <-context.Done():
return context.Err()
}
} }