diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 30172e8..f0e65a0 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -19,9 +19,9 @@ jobs: permissions: contents: read steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 - - uses: actions/setup-go@v6 + - uses: actions/setup-go@v7 with: go-version-file: go.mod cache: false @@ -64,7 +64,7 @@ jobs: permissions: contents: read steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 - name: Set up QEMU uses: docker/setup-qemu-action@v4 diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index e10d825..4e0297f 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -19,9 +19,9 @@ jobs: permissions: contents: read steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 - - uses: actions/setup-go@v6 + - uses: actions/setup-go@v7 with: go-version-file: go.mod cache: false diff --git a/.github/workflows/docker.yml b/.github/workflows/docker.yml index 9c9216b..5aa0818 100644 --- a/.github/workflows/docker.yml +++ b/.github/workflows/docker.yml @@ -11,6 +11,7 @@ on: env: REGISTRY: ghcr.io + UPSTREAM_REPOSITORY: TypeType-Video/TypeType FORCE_JAVASCRIPT_ACTIONS_TO_NODE24: true concurrency: @@ -30,18 +31,37 @@ jobs: steps: - name: Checkout - uses: actions/checkout@v6 + uses: actions/checkout@v7 - name: Isolate Docker credentials run: echo "DOCKER_CONFIG=$RUNNER_TEMP/docker-config" >> "$GITHUB_ENV" - name: Resolve build metadata id: build-info + env: + GH_TOKEN: ${{ github.token }} run: | - base_version="0.1.0" if [[ "$GITHUB_REF" == refs/tags/v* ]]; then - version="${GITHUB_REF_NAME#v}" + base_version="${GITHUB_REF_NAME#v}" + if [[ ! "$base_version" =~ ^[0-9]+\.[0-9]+\.[0-9]+$ ]]; then + echo "Invalid component release tag: $GITHUB_REF_NAME" + exit 1 + fi + gh api "repos/${UPSTREAM_REPOSITORY}/git/ref/tags/v${base_version}" >/dev/null + version="$base_version" else + upstream_tag="$(gh api "repos/${UPSTREAM_REPOSITORY}/releases/latest" --jq '.tag_name')" + if [[ ! "$upstream_tag" =~ ^v([0-9]+\.[0-9]+\.[0-9]+)$ ]]; then + echo "Invalid upstream TypeType release tag: $upstream_tag" + exit 1 + fi + base_version="${BASH_REMATCH[1]}" + fi + if [[ "$GITHUB_REF_NAME" == "main" ]]; then + version="$base_version" + elif [[ "$GITHUB_REF_NAME" == "dev" ]]; then + version="$base_version-dev.$GITHUB_RUN_NUMBER" + elif [[ "$GITHUB_REF" != refs/tags/v* ]]; then channel="${GITHUB_REF_NAME//[^0-9A-Za-z-]/-}" version="$base_version-$channel.$GITHUB_RUN_NUMBER" fi @@ -69,6 +89,7 @@ jobs: with: images: ${{ steps.build-info.outputs.image }} tags: | + type=raw,value=${{ steps.build-info.outputs.version }} type=sha,prefix=sha-,format=short type=ref,event=branch type=ref,event=tag diff --git a/Dockerfile b/Dockerfile index b97f2bd..154d938 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,12 +1,10 @@ -FROM golang:1.26-trixie AS build +FROM golang:1.26-alpine3.23 AS build -ARG BUILD_VERSION=0.1.0 +ARG BUILD_VERSION=1.2.4-dev ARG BUILD_REVISION=development ARG BUILD_TIME=unknown WORKDIR /src -RUN apt-get update \ - && apt-get install -y --no-install-recommends pkg-config libavformat-dev libavcodec-dev libavutil-dev \ - && rm -rf /var/lib/apt/lists/* +RUN apk add --no-cache build-base pkgconf ffmpeg-dev COPY go.mod go.sum ./ RUN go mod download @@ -14,13 +12,11 @@ RUN go mod download COPY . . RUN CGO_ENABLED=1 go build -trimpath -ldflags="-s -w -X typetype-downloader-go/internal/buildinfo.Version=$BUILD_VERSION -X typetype-downloader-go/internal/buildinfo.Revision=$BUILD_REVISION -X typetype-downloader-go/internal/buildinfo.BuildTime=$BUILD_TIME" -o /out/typetype-downloader-go ./cmd/server -FROM debian:trixie-slim +FROM alpine:3.23 -RUN apt-get update \ - && apt-get install -y --no-install-recommends ca-certificates ffmpeg libavformat61 libavcodec61 libavutil59 \ - && rm -rf /var/lib/apt/lists/* +RUN apk add --no-cache ca-certificates ffmpeg -RUN useradd --system --create-home --home-dir /app typetype +RUN addgroup -S typetype && adduser -S -G typetype -h /app typetype WORKDIR /app COPY --from=build /out/typetype-downloader-go /usr/local/bin/typetype-downloader-go RUN mkdir -p /app/data && chown -R typetype:typetype /app diff --git a/Dockerfile.wolfi b/Dockerfile.wolfi index d0ee6ed..bd3540b 100644 --- a/Dockerfile.wolfi +++ b/Dockerfile.wolfi @@ -1,6 +1,6 @@ FROM cgr.dev/chainguard/go:latest-dev AS build -ARG BUILD_VERSION=0.1.0 +ARG BUILD_VERSION=1.2.4-dev ARG BUILD_REVISION=development ARG BUILD_TIME=unknown USER root diff --git a/cmd/server/main.go b/cmd/server/main.go index cf40cb6..2865276 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -39,12 +39,12 @@ func main() { store := job.NewStore(cfg.PublicBaseURL, sinks...) store.Restore(restored) pendingIDs := store.RestorePending(pending) - runner := pipeline.NewRunner(cfg, store, files) disk, err := storage.NewMonitor(cfg.DataDir, cfg.MinFreeBytes, cfg.MinFreePercent) if err != nil { slog.Error("storage monitor failed", "error", err) os.Exit(1) } + runner := pipeline.NewRunner(cfg, store, files, disk) runner.Start(ctx) cleanup.Start(ctx, cfg.DataDir, cfg.StorageBackend) for _, id := range pendingIDs { diff --git a/docker-compose.yml b/docker-compose.yml index 0356cd0..e492318 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -14,7 +14,7 @@ services: retries: 30 dragonfly: - image: docker.dragonflydb.io/dragonflydb/dragonfly:v1.38.0 + image: docker.dragonflydb.io/dragonflydb/dragonfly:v1.39.0 ports: - "56379:6379" diff --git a/internal/api/version_test.go b/internal/api/version_test.go index 3af8fb9..a7fc10f 100644 --- a/internal/api/version_test.go +++ b/internal/api/version_test.go @@ -21,7 +21,9 @@ func TestVersionReturnsBuildMetadata(t *testing.T) { if err := json.Unmarshal(response.Body.Bytes(), &body); err != nil { t.Fatal(err) } - if body["revision"] != buildinfo.Revision || body["service"] != "downloader" { + if body["service"] != "downloader" || + body["version"] != buildinfo.Version || + body["revision"] != buildinfo.Revision { t.Fatalf("body = %#v", body) } } diff --git a/internal/buildinfo/info.go b/internal/buildinfo/info.go index bcb0787..1d50edf 100644 --- a/internal/buildinfo/info.go +++ b/internal/buildinfo/info.go @@ -1,6 +1,6 @@ package buildinfo -var Version = "0.1.0" +var Version = "1.2.4-dev" var Revision = "development" var BuildTime = "unknown" diff --git a/internal/downloader/range_downloader.go b/internal/downloader/range_downloader.go index 4c5b009..950dcc4 100644 --- a/internal/downloader/range_downloader.go +++ b/internal/downloader/range_downloader.go @@ -3,13 +3,9 @@ package downloader import ( "context" "fmt" - "io" "net/http" - "net/url" "os" - "strings" "sync" - "sync/atomic" "time" ) @@ -37,31 +33,21 @@ type Progress struct { type ProgressFunc func(Progress) -func DownloadFile(ctx context.Context, client *http.Client, source Source, output string, options Options, progress ProgressFunc) error { +func DownloadFile( + ctx context.Context, + client *http.Client, + source Source, + output string, + options Options, + progress ProgressFunc, +) error { if source.URL == "" { return fmt.Errorf("invalid source %s", source.Name) } if client == nil { client = http.DefaultClient } - if options.ChunkSize <= 0 { - options.ChunkSize = 10 << 20 - } - if options.Workers <= 0 { - options.Workers = 8 - } - if options.Retries <= 0 { - options.Retries = 4 - } - if options.BufferSize <= 0 { - options.BufferSize = 256 * 1024 - } - if options.ProgressBytes <= 0 { - options.ProgressBytes = 4 << 20 - } - if options.RangeMode == "" { - options.RangeMode = "header" - } + options = normalizedOptions(options) if source.Size <= 0 { size, err := probeSourceSize(ctx, client, source.URL, options.RangeMode) if err != nil { @@ -76,26 +62,78 @@ func DownloadFile(ctx context.Context, client *http.Client, source Source, outpu if err != nil { return err } - defer file.Close() + committed := false + defer func() { + _ = file.Close() + if !committed { + _ = os.Remove(tmp) + } + }() if err := file.Truncate(source.Size); err != nil { return err } + if err := downloadChunks(ctx, client, file, source, options, progress); err != nil { + return err + } + if err := file.Close(); err != nil { + return err + } + if err := os.Rename(tmp, output); err != nil { + return err + } + committed = true + return nil +} + +func normalizedOptions(options Options) Options { + if options.ChunkSize <= 0 { + options.ChunkSize = 10 << 20 + } + if options.Workers <= 0 { + options.Workers = 8 + } + if options.Retries <= 0 { + options.Retries = 4 + } + if options.BufferSize <= 0 { + options.BufferSize = defaultCopyBufferSize + } + if options.ProgressBytes <= 0 { + options.ProgressBytes = 4 << 20 + } + if options.RangeMode == "" { + options.RangeMode = "header" + } + return options +} + +func downloadChunks( + ctx context.Context, + client *http.Client, + file *os.File, + source Source, + options Options, + progress ProgressFunc, +) error { downloadCtx, cancel := context.WithCancel(ctx) defer cancel() - jobs := make(chan chunk) + chunkCount := int((source.Size + options.ChunkSize - 1) / options.ChunkSize) + workers := min(options.Workers, chunkCount) + jobs := make(chan chunk, workers) errs := make(chan error, 1) - var downloaded atomic.Int64 - started := time.Now() + tracker := newProgressTracker(source, options.ProgressBytes, progress) var wg sync.WaitGroup - for range options.Workers { + for range workers { wg.Add(1) go func() { defer wg.Done() - for job := range jobs { - if err := downloadChunk(downloadCtx, client, file, source, job, options, &downloaded, started, progress); err != nil { + buffer := borrowCopyBuffer(options.BufferSize) + defer releaseCopyBuffer(buffer) + for part := range jobs { + if err := downloadChunk(downloadCtx, client, file, source, part, options, buffer, tracker); err != nil { select { case errs <- err: default: @@ -107,21 +145,9 @@ func DownloadFile(ctx context.Context, client *http.Client, source Source, outpu }() } -sendLoop: - for start := int64(0); start < source.Size; start += options.ChunkSize { - end := start + options.ChunkSize - 1 - if end >= source.Size { - end = source.Size - 1 - } - select { - case <-downloadCtx.Done(): - break sendLoop - case jobs <- chunk{start: start, end: end}: - } - } + queueChunks(downloadCtx, jobs, source.Size, options.ChunkSize) close(jobs) wg.Wait() - select { case err := <-errs: return err @@ -130,128 +156,20 @@ sendLoop: if err := ctx.Err(); err != nil { return err } - if got := downloaded.Load(); got != source.Size { + if got := tracker.downloaded.Load(); got != source.Size { return fmt.Errorf("downloaded %d bytes, expected %d", got, source.Size) } - if err := file.Close(); err != nil { - return err - } - return os.Rename(tmp, output) -} - -type chunk struct { - start int64 - end int64 -} - -func downloadChunk(ctx context.Context, client *http.Client, file *os.File, source Source, part chunk, options Options, downloaded *atomic.Int64, started time.Time, progress ProgressFunc) error { - var lastErr error - for attempt := 0; attempt < options.Retries; attempt++ { - if attempt > 0 { - select { - case <-ctx.Done(): - return ctx.Err() - case <-time.After(time.Duration(attempt) * 300 * time.Millisecond): - } - } - err := fetchChunk(ctx, client, file, source, part, options, downloaded, started, progress) - if err == nil { - return nil - } - lastErr = err - } - return fmt.Errorf("%s bytes %d-%d failed: %w", source.Name, part.start, part.end, lastErr) -} - -func fetchChunk(ctx context.Context, client *http.Client, file *os.File, source Source, part chunk, options Options, downloaded *atomic.Int64, started time.Time, progress ProgressFunc) error { - rangeMode := effectiveRangeMode(source.URL, options.RangeMode) - requestURL, err := rangedURL(source.URL, part, rangeMode) - if err != nil { - return err - } - requestURL, err = stripFragment(requestURL) - if err != nil { - return err - } - req, err := http.NewRequestWithContext(ctx, http.MethodGet, requestURL, nil) - if err != nil { - return err - } - req.Header.Set("Accept-Encoding", "identity") - if rangeMode == "header" { - req.Header.Set("Range", fmt.Sprintf("bytes=%d-%d", part.start, part.end)) - } - applyMediaHeaders(req, source.URL) - - res, err := client.Do(req) - if err != nil { - return err - } - defer res.Body.Close() - if res.StatusCode != http.StatusPartialContent && !(rangeMode == "query" && res.StatusCode == http.StatusOK) { - return fmt.Errorf("unexpected HTTP status %d", res.StatusCode) - } - if contentRange := res.Header.Get("Content-Range"); rangeMode == "header" && !strings.HasPrefix(contentRange, fmt.Sprintf("bytes %d-%d/", part.start, part.end)) { - return fmt.Errorf("unexpected Content-Range %q", contentRange) - } - - buf := make([]byte, options.BufferSize) - position := part.start - written := int64(0) - for { - n, readErr := res.Body.Read(buf) - if n > 0 { - if position+int64(n)-1 > part.end { - return fmt.Errorf("chunk overflow") - } - if _, err := file.WriteAt(buf[:n], position); err != nil { - return err - } - position += int64(n) - written += int64(n) - } - if readErr == io.EOF { - break - } - if readErr != nil { - return readErr - } - } - if position != part.end+1 { - return fmt.Errorf("short chunk: got %d expected %d", position-part.start, part.end-part.start+1) - } - current := downloaded.Add(written) - if progress != nil { - elapsed := time.Since(started).Seconds() - if elapsed > 0 { - progress(Progress{Name: source.Name, Downloaded: current, Total: source.Size, Speed: float64(current) / elapsed}) - } - } + tracker.finish(time.Now()) return nil } -func stripFragment(rawURL string) (string, error) { - parsed, err := url.Parse(rawURL) - if err != nil { - return "", err - } - parsed.Fragment = "" - return parsed.String(), nil -} - -func rangedURL(rawURL string, part chunk, mode string) (string, error) { - if mode == "header" { - return rawURL, nil - } - if mode != "query" { - return "", fmt.Errorf("unsupported range mode %q", mode) - } - parsed, err := url.Parse(rawURL) - if err != nil { - return "", err +func queueChunks(ctx context.Context, jobs chan<- chunk, size int64, chunkSize int64) { + for start := int64(0); start < size; start += chunkSize { + end := min(start+chunkSize, size) - 1 + select { + case <-ctx.Done(): + return + case jobs <- chunk{start: start, end: end}: + } } - query := parsed.Query() - query.Set("range", fmt.Sprintf("%d-%d", part.start, part.end)) - parsed.RawQuery = query.Encode() - return parsed.String(), nil } diff --git a/internal/downloader/range_downloader_test.go b/internal/downloader/range_downloader_test.go new file mode 100644 index 0000000..f48449d --- /dev/null +++ b/internal/downloader/range_downloader_test.go @@ -0,0 +1,253 @@ +package downloader + +import ( + "bytes" + "fmt" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strconv" + "strings" + "sync/atomic" + "testing" +) + +func TestDownloadFileWritesConcurrentRangesExactly(t *testing.T) { + data := repeatedPayload(3*64*1024 + 713) + for _, mode := range []string{"header", "query"} { + t.Run(mode, func(t *testing.T) { + var requests atomic.Int64 + server := httptest.NewServer(rangeServer(data, &requests, nil)) + defer server.Close() + + output := filepath.Join(t.TempDir(), "media.bin") + var updates []Progress + err := DownloadFile(t.Context(), server.Client(), Source{ + Name: "video", + URL: server.URL + "/media#cookie=ignored", + Size: int64(len(data)), + }, output, Options{ + ChunkSize: 64 * 1024, + Workers: 32, + Retries: 1, + BufferSize: 16 * 1024, + RangeMode: mode, + ProgressBytes: 64 * 1024, + }, func(progress Progress) { + updates = append(updates, progress) + }) + if err != nil { + t.Fatalf("DownloadFile() error = %v", err) + } + + got, err := os.ReadFile(output) + if err != nil { + t.Fatal(err) + } + if !bytes.Equal(got, data) { + t.Fatal("downloaded data differs from source") + } + if got := requests.Load(); got != 4 { + t.Fatalf("requests = %d, want 4", got) + } + assertProgressMonotonic(t, updates, int64(len(data))) + }) + } +} + +func TestDownloadFileRetriesTruncatedRange(t *testing.T) { + data := repeatedPayload(80 * 1024) + var attempts atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + start, end, err := parseRequestedRange(request) + if err != nil { + http.Error(response, err.Error(), http.StatusBadRequest) + return + } + part := data[start : end+1] + response.Header().Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, end, len(data))) + response.Header().Set("Content-Length", strconv.Itoa(len(part))) + response.WriteHeader(http.StatusPartialContent) + if attempts.Add(1) == 1 { + _, _ = response.Write(part[:len(part)/2]) + return + } + _, _ = response.Write(part) + })) + defer server.Close() + + output := filepath.Join(t.TempDir(), "media.bin") + err := DownloadFile(t.Context(), server.Client(), Source{ + Name: "audio", + URL: server.URL, + Size: int64(len(data)), + }, output, Options{ + ChunkSize: int64(len(data)), + Workers: 1, + Retries: 2, + BufferSize: 32 * 1024, + RangeMode: "header", + }, nil) + if err != nil { + t.Fatalf("DownloadFile() error = %v", err) + } + if got := attempts.Load(); got != 2 { + t.Fatalf("attempts = %d, want 2", got) + } + got, err := os.ReadFile(output) + if err != nil { + t.Fatal(err) + } + if !bytes.Equal(got, data) { + t.Fatal("downloaded data differs after retry") + } +} + +func TestDownloadFilePreservesOutputAndRemovesPartialFileOnFailure(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, _ *http.Request) { + http.Error(response, "unavailable", http.StatusServiceUnavailable) + })) + defer server.Close() + + dir := t.TempDir() + output := filepath.Join(dir, "media.bin") + if err := os.WriteFile(output, []byte("existing"), 0o644); err != nil { + t.Fatal(err) + } + err := DownloadFile(t.Context(), server.Client(), Source{ + Name: "video", + URL: server.URL, + Size: 1024, + }, output, Options{Retries: 1}, nil) + if err == nil { + t.Fatal("DownloadFile() error = nil, want failure") + } + got, err := os.ReadFile(output) + if err != nil { + t.Fatal(err) + } + if string(got) != "existing" { + t.Fatalf("output = %q, want existing data", got) + } + if _, err := os.Stat(output + ".part"); !os.IsNotExist(err) { + t.Fatalf("partial file still exists: %v", err) + } +} + +func TestDownloadFileRejectsChunkOverflow(t *testing.T) { + data := repeatedPayload(32 * 1024) + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, _ *http.Request) { + flusher := response.(http.Flusher) + response.WriteHeader(http.StatusOK) + _, _ = response.Write(data) + flusher.Flush() + _, _ = response.Write([]byte{1}) + })) + defer server.Close() + + output := filepath.Join(t.TempDir(), "media.bin") + err := DownloadFile(t.Context(), server.Client(), Source{ + Name: "audio", + URL: server.URL, + Size: int64(len(data)), + }, output, Options{ + ChunkSize: int64(len(data)), + Workers: 1, + Retries: 1, + RangeMode: "query", + }, nil) + if err == nil || !strings.Contains(err.Error(), "overflow") { + t.Fatalf("DownloadFile() error = %v, want overflow", err) + } +} + +func BenchmarkDownloadFile(b *testing.B) { + data := repeatedPayload(16 << 20) + server := httptest.NewServer(rangeServer(data, nil, nil)) + defer server.Close() + output := filepath.Join(b.TempDir(), "media.bin") + options := Options{ + ChunkSize: 1 << 20, + Workers: 8, + Retries: 1, + BufferSize: defaultCopyBufferSize, + RangeMode: "header", + } + + b.SetBytes(int64(len(data))) + b.ReportAllocs() + b.ResetTimer() + for range b.N { + if err := DownloadFile(b.Context(), server.Client(), Source{ + Name: "video", + URL: server.URL, + Size: int64(len(data)), + }, output, options, nil); err != nil { + b.Fatal(err) + } + } +} + +func rangeServer(data []byte, requests *atomic.Int64, truncated *atomic.Bool) http.HandlerFunc { + return func(response http.ResponseWriter, request *http.Request) { + if requests != nil { + requests.Add(1) + } + start, end, err := parseRequestedRange(request) + if err != nil { + http.Error(response, err.Error(), http.StatusBadRequest) + return + } + part := data[start : end+1] + response.Header().Set("Content-Length", strconv.Itoa(len(part))) + if request.Header.Get("Range") != "" { + response.Header().Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, end, len(data))) + response.WriteHeader(http.StatusPartialContent) + } + if truncated != nil && truncated.CompareAndSwap(false, true) { + part = part[:len(part)/2] + } + for len(part) > 0 { + size := min(3*1024, len(part)) + _, _ = response.Write(part[:size]) + part = part[size:] + } + } +} + +func parseRequestedRange(request *http.Request) (int, int, error) { + value := request.URL.Query().Get("range") + if value == "" { + value = strings.TrimPrefix(request.Header.Get("Range"), "bytes=") + } + startText, endText, ok := strings.Cut(value, "-") + if !ok { + return 0, 0, fmt.Errorf("invalid range %q", value) + } + start, startErr := strconv.Atoi(startText) + end, endErr := strconv.Atoi(endText) + if startErr != nil || endErr != nil || start < 0 || end < start { + return 0, 0, fmt.Errorf("invalid range %q", value) + } + return start, end, nil +} + +func repeatedPayload(size int) []byte { + pattern := []byte("typetype-downloader-performance") + return bytes.Repeat(pattern, (size+len(pattern)-1)/len(pattern))[:size] +} + +func assertProgressMonotonic(t *testing.T, updates []Progress, total int64) { + t.Helper() + var previous int64 + for _, update := range updates { + if update.Downloaded <= previous { + t.Fatalf("progress moved from %d to %d", previous, update.Downloaded) + } + previous = update.Downloaded + } + if previous != total { + t.Fatalf("final progress = %d, want %d", previous, total) + } +} diff --git a/internal/downloader/range_request.go b/internal/downloader/range_request.go new file mode 100644 index 0000000..ad94f6d --- /dev/null +++ b/internal/downloader/range_request.go @@ -0,0 +1,110 @@ +package downloader + +import ( + "context" + "fmt" + "io" + "net/http" + "net/url" + "os" + "strings" +) + +func fetchChunk( + ctx context.Context, + client *http.Client, + file *os.File, + source Source, + part chunk, + configuredMode string, + buffer []byte, +) error { + rangeMode := effectiveRangeMode(source.URL, configuredMode) + request, err := rangedRequest(ctx, source.URL, part, rangeMode) + if err != nil { + return err + } + response, err := client.Do(request) + if err != nil { + return err + } + defer response.Body.Close() + if err := validateRangeResponse(response, part, rangeMode); err != nil { + return err + } + if err := copyChunk(response.Body, file, part, buffer); err != nil { + return err + } + if response.ContentLength < 0 { + var trailing [1]byte + if n, readErr := response.Body.Read(trailing[:]); n > 0 || readErr != io.EOF { + return fmt.Errorf("chunk overflow") + } + } + return nil +} + +func rangedRequest(ctx context.Context, rawURL string, part chunk, rangeMode string) (*http.Request, error) { + requestURL, err := rangedURL(rawURL, part, rangeMode) + if err != nil { + return nil, err + } + requestURL, err = stripFragment(requestURL) + if err != nil { + return nil, err + } + request, err := http.NewRequestWithContext(ctx, http.MethodGet, requestURL, nil) + if err != nil { + return nil, err + } + request.Header.Set("Accept-Encoding", "identity") + if rangeMode == "header" { + request.Header.Set("Range", fmt.Sprintf("bytes=%d-%d", part.start, part.end)) + } + applyMediaHeaders(request, rawURL) + return request, nil +} + +func validateRangeResponse(response *http.Response, part chunk, configuredMode string) error { + expected := part.end - part.start + 1 + if response.StatusCode != http.StatusPartialContent && + !(configuredMode == "query" && response.StatusCode == http.StatusOK) { + return fmt.Errorf("unexpected HTTP status %d", response.StatusCode) + } + if configuredMode == "header" { + want := fmt.Sprintf("bytes %d-%d/", part.start, part.end) + if contentRange := response.Header.Get("Content-Range"); !strings.HasPrefix(contentRange, want) { + return fmt.Errorf("unexpected Content-Range %q", contentRange) + } + } + if response.ContentLength >= 0 && response.ContentLength != expected { + return fmt.Errorf("unexpected content length %d, expected %d", response.ContentLength, expected) + } + return nil +} + +func stripFragment(rawURL string) (string, error) { + parsed, err := url.Parse(rawURL) + if err != nil { + return "", err + } + parsed.Fragment = "" + return parsed.String(), nil +} + +func rangedURL(rawURL string, part chunk, mode string) (string, error) { + if mode == "header" { + return rawURL, nil + } + if mode != "query" { + return "", fmt.Errorf("unsupported range mode %q", mode) + } + parsed, err := url.Parse(rawURL) + if err != nil { + return "", err + } + query := parsed.Query() + query.Set("range", fmt.Sprintf("%d-%d", part.start, part.end)) + parsed.RawQuery = query.Encode() + return parsed.String(), nil +} diff --git a/internal/downloader/range_worker.go b/internal/downloader/range_worker.go new file mode 100644 index 0000000..5573127 --- /dev/null +++ b/internal/downloader/range_worker.go @@ -0,0 +1,159 @@ +package downloader + +import ( + "context" + "fmt" + "io" + "net/http" + "os" + "sync" + "sync/atomic" + "time" +) + +const defaultCopyBufferSize = 256 * 1024 + +var copyBufferPool = sync.Pool{ + New: func() any { + return make([]byte, defaultCopyBufferSize) + }, +} + +type chunk struct { + start int64 + end int64 +} + +type progressTracker struct { + source Source + interval int64 + progress ProgressFunc + started time.Time + downloaded atomic.Int64 + reportMu sync.Mutex + reported int64 +} + +func newProgressTracker(source Source, interval int64, progress ProgressFunc) *progressTracker { + return &progressTracker{ + source: source, + interval: interval, + progress: progress, + started: time.Now(), + } +} + +func (tracker *progressTracker) add(bytes int64, now time.Time) { + current := tracker.downloaded.Add(bytes) + if tracker.progress == nil || current >= tracker.source.Size { + return + } + tracker.reportMu.Lock() + defer tracker.reportMu.Unlock() + if current <= tracker.reported || current-tracker.reported < tracker.interval { + return + } + tracker.reported = current + tracker.report(current, now) +} + +func (tracker *progressTracker) finish(now time.Time) { + if tracker.progress == nil { + return + } + tracker.reportMu.Lock() + defer tracker.reportMu.Unlock() + current := tracker.downloaded.Load() + if current <= tracker.reported { + return + } + tracker.reported = current + tracker.report(current, now) +} + +func (tracker *progressTracker) report(downloaded int64, now time.Time) { + elapsed := now.Sub(tracker.started).Seconds() + if elapsed <= 0 { + return + } + tracker.progress(Progress{ + Name: tracker.source.Name, + Downloaded: downloaded, + Total: tracker.source.Size, + Speed: float64(downloaded) / elapsed, + }) +} + +func borrowCopyBuffer(size int) []byte { + if size == defaultCopyBufferSize { + return copyBufferPool.Get().([]byte) + } + return make([]byte, size) +} + +func releaseCopyBuffer(buffer []byte) { + if cap(buffer) == defaultCopyBufferSize { + copyBufferPool.Put(buffer[:defaultCopyBufferSize]) + } +} + +func downloadChunk( + ctx context.Context, + client *http.Client, + file *os.File, + source Source, + part chunk, + options Options, + buffer []byte, + tracker *progressTracker, +) error { + var lastErr error + for attempt := 0; attempt < options.Retries; attempt++ { + if err := waitForRetry(ctx, attempt); err != nil { + return err + } + if err := fetchChunk(ctx, client, file, source, part, options.RangeMode, buffer); err != nil { + lastErr = err + continue + } + tracker.add(part.end-part.start+1, time.Now()) + return nil + } + return fmt.Errorf("%s bytes %d-%d failed: %w", source.Name, part.start, part.end, lastErr) +} + +func waitForRetry(ctx context.Context, attempt int) error { + if attempt == 0 { + return nil + } + timer := time.NewTimer(time.Duration(attempt) * 300 * time.Millisecond) + defer timer.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil + } +} + +func copyChunk(body io.Reader, file *os.File, part chunk, buffer []byte) error { + position := part.start + for position <= part.end { + size := min(int64(len(buffer)), part.end-position+1) + n, err := io.ReadFull(body, buffer[:size]) + if n > 0 { + written, writeErr := file.WriteAt(buffer[:n], position) + if writeErr != nil { + return writeErr + } + if written != n { + return io.ErrShortWrite + } + position += int64(n) + } + if err != nil { + return fmt.Errorf("short chunk: got %d expected %d: %w", position-part.start, part.end-part.start+1, err) + } + } + return nil +} diff --git a/internal/mux/real_remux_test.go b/internal/mux/real_remux_test.go new file mode 100644 index 0000000..c943c6d --- /dev/null +++ b/internal/mux/real_remux_test.go @@ -0,0 +1,30 @@ +package mux + +import ( + "context" + "os" + "path/filepath" + "testing" + "time" +) + +func TestRealRemux(t *testing.T) { + videoPath := os.Getenv("TYPETYPE_REMUX_VIDEO") + audioPath := os.Getenv("TYPETYPE_REMUX_AUDIO") + if videoPath == "" || audioPath == "" { + t.Skip("set TYPETYPE_REMUX_VIDEO and TYPETYPE_REMUX_AUDIO to enable the real remux test") + } + outputPath := os.Getenv("TYPETYPE_REMUX_OUTPUT") + if outputPath == "" { + outputPath = filepath.Join(t.TempDir(), "output.mp4") + } + started := time.Now() + if err := RemuxAVFormat(context.Background(), videoPath, audioPath, outputPath); err != nil { + t.Fatal(err) + } + info, err := os.Stat(outputPath) + if err != nil { + t.Fatal(err) + } + t.Logf("duration=%s bytes=%d", time.Since(started), info.Size()) +} diff --git a/internal/pipeline/capacity.go b/internal/pipeline/capacity.go new file mode 100644 index 0000000..cbce44d --- /dev/null +++ b/internal/pipeline/capacity.go @@ -0,0 +1,85 @@ +package pipeline + +import ( + "math" + + "typetype-downloader-go/internal/selector" +) + +const ( + outputHeadroomTenths = uint64(11) + localPeakTenths = uint64(21) +) + +func (r *Runner) reserveVideo(id string, selection *selector.Selection, duration int64) (func(), error) { + bytes := videoReservationBytes(selection, duration) + if bytes == 0 { + return noRelease, nil + } + return r.disk.Reserve(id, bytes) +} + +func (r *Runner) reserveAudio(id string, selection *selector.AudioSelection, duration int64) (func(), error) { + bytes := audioReservationBytes(selection, duration) + if bytes == 0 { + return noRelease, nil + } + return r.disk.Reserve(id, bytes) +} + +func videoReservationBytes(selection *selector.Selection, duration int64) uint64 { + videoBytes := mediaBytes(selection.Video.ContentLength, selection.Video.Bitrate, duration) + audioBytes := mediaBytes(selection.Audio.ContentLength, selection.Audio.Bitrate, duration) + total := knownSum(videoBytes, audioBytes) + if total == 0 { + return 0 + } + multiplier := outputHeadroomTenths + if selection.Video.DeliveryMethod == "sabr" || + selection.Video.ContentLength > 0 && selection.Audio.ContentLength > 0 && !usesRemoteMux(selection) { + multiplier = localPeakTenths + } + return scaledBytes(total, multiplier) +} + +func audioReservationBytes(selection *selector.AudioSelection, duration int64) uint64 { + bytes := mediaBytes(selection.Audio.ContentLength, selection.Audio.Bitrate, duration) + if bytes == 0 { + return 0 + } + multiplier := outputHeadroomTenths + if selection.Audio.DeliveryMethod == "sabr" { + multiplier = localPeakTenths + } + return scaledBytes(bytes, multiplier) +} + +func mediaBytes(contentLength int64, bitrate *int, duration int64) uint64 { + if contentLength > 0 { + return uint64(contentLength) + } + if bitrate == nil || *bitrate <= 0 || duration <= 0 { + return 0 + } + rate := uint64(*bitrate) + if rate > math.MaxUint64/uint64(duration) { + return math.MaxUint64 + } + return rate * uint64(duration) / 8 +} + +func knownSum(left uint64, right uint64) uint64 { + if left > math.MaxUint64-right { + return math.MaxUint64 + } + return left + right +} + +func scaledBytes(bytes uint64, tenths uint64) uint64 { + if bytes > math.MaxUint64/tenths { + return math.MaxUint64 + } + return bytes * tenths / 10 +} + +func noRelease() {} diff --git a/internal/pipeline/capacity_test.go b/internal/pipeline/capacity_test.go new file mode 100644 index 0000000..fe1eef7 --- /dev/null +++ b/internal/pipeline/capacity_test.go @@ -0,0 +1,68 @@ +package pipeline + +import ( + "math" + "testing" + + "typetype-downloader-go/internal/selector" + "typetype-downloader-go/internal/typetype" +) + +func TestKnownSumUsesAvailableInputs(t *testing.T) { + if got := knownSum(100, 50); got != 150 { + t.Fatalf("sum = %d", got) + } + if got := knownSum(100, 0); got != 100 { + t.Fatalf("partial sum = %d", got) + } +} + +func TestMediaBytesPrefersContentLengthAndEstimatesFallback(t *testing.T) { + bitrate := 128_000 + if got := mediaBytes(42, &bitrate, 36_000); got != 42 { + t.Fatalf("content bytes = %d", got) + } + if got := mediaBytes(0, &bitrate, 36_000); got != 576_000_000 { + t.Fatalf("estimated bytes = %d", got) + } +} + +func TestCapacityEstimateSaturates(t *testing.T) { + if got := scaledBytes(100, 21); got != 210 { + t.Fatalf("scaled bytes = %d", got) + } + if got := scaledBytes(math.MaxUint64, 21); got != math.MaxUint64 { + t.Fatalf("saturated bytes = %d", got) + } +} + +func TestVideoReservationCoversLocalAssemblyPeak(t *testing.T) { + selection := &selector.Selection{ + Video: typetype.VideoStreamItem{URL: "video", ContentLength: 100}, + Audio: typetype.AudioStreamItem{URL: "audio", ContentLength: 50}, + } + if got := videoReservationBytes(selection, 0); got != 315 { + t.Fatalf("local reservation = %d", got) + } + selection.Video.URL = "https://example.test/video.m3u8" + if got := videoReservationBytes(selection, 0); got != 165 { + t.Fatalf("remote reservation = %d", got) + } + selection.Video.DeliveryMethod = "sabr" + if got := videoReservationBytes(selection, 0); got != 315 { + t.Fatalf("SABR reservation = %d", got) + } +} + +func TestAudioReservationCoversSABRAssemblyPeak(t *testing.T) { + selection := &selector.AudioSelection{ + Audio: typetype.AudioStreamItem{ContentLength: 100}, + } + if got := audioReservationBytes(selection, 0); got != 110 { + t.Fatalf("direct reservation = %d", got) + } + selection.Audio.DeliveryMethod = "sabr" + if got := audioReservationBytes(selection, 0); got != 210 { + t.Fatalf("SABR reservation = %d", got) + } +} diff --git a/internal/pipeline/failure.go b/internal/pipeline/failure.go new file mode 100644 index 0000000..3789a09 --- /dev/null +++ b/internal/pipeline/failure.go @@ -0,0 +1,18 @@ +package pipeline + +import ( + "context" + "errors" + + "typetype-downloader-go/internal/storage" +) + +func failureCode(ctx context.Context, err error) string { + if ctx.Err() != nil || errors.Is(err, context.Canceled) { + return "cancelled" + } + if errors.Is(err, storage.ErrInsufficientStorage) { + return "insufficient_storage" + } + return "download_failed" +} diff --git a/internal/pipeline/failure_test.go b/internal/pipeline/failure_test.go new file mode 100644 index 0000000..7d8bc70 --- /dev/null +++ b/internal/pipeline/failure_test.go @@ -0,0 +1,16 @@ +package pipeline + +import ( + "context" + "fmt" + "testing" + + "typetype-downloader-go/internal/storage" +) + +func TestFailureCodePreservesStorageFailure(t *testing.T) { + err := fmt.Errorf("reserve: %w", storage.ErrInsufficientStorage) + if got := failureCode(context.Background(), err); got != "insufficient_storage" { + t.Fatalf("failure code = %q", got) + } +} diff --git a/internal/pipeline/remote.go b/internal/pipeline/remote.go index 671245f..fc9111b 100644 --- a/internal/pipeline/remote.go +++ b/internal/pipeline/remote.go @@ -13,7 +13,7 @@ import ( func (r *Runner) runRemote(ctx context.Context, id string, title string, selection *selector.Selection) error { paths := artifact.Build(r.cfg.DataDir, id, title, selection.Container) r.store.Resolve(id, title, resolvedOutput(selection, paths.Name)) - return r.runSABRArtifact(ctx, id, paths, func() (int64, error) { + return r.runSABRArtifact(ctx, id, paths, func() (artifactTimings, error) { started := time.Now() r.store.Progress(id, job.Progress{Stage: "download"}) err := ffmpeg.DownloadRemote( @@ -22,17 +22,17 @@ func (r *Runner) runRemote(ctx context.Context, id string, title string, selecti r.streams.ProxyMediaURL(selection.Audio.URL), paths.Output, ) - return time.Since(started).Milliseconds(), err + return artifactTimings{downloadMs: time.Since(started).Milliseconds()}, err }) } func (r *Runner) runRemoteAudio(ctx context.Context, id string, title string, selection *selector.AudioSelection) error { paths := artifact.Build(r.cfg.DataDir, id, title, selection.Container) r.store.Resolve(id, title, audioResolvedOutput(selection, paths.Name)) - return r.runSABRArtifact(ctx, id, paths, func() (int64, error) { + return r.runSABRArtifact(ctx, id, paths, func() (artifactTimings, error) { started := time.Now() r.store.Progress(id, job.Progress{Stage: "download"}) err := ffmpeg.DownloadRemoteAudio(ctx, r.streams.ProxyMediaURL(selection.Audio.URL), paths.Output) - return time.Since(started).Milliseconds(), err + return artifactTimings{downloadMs: time.Since(started).Milliseconds()}, err }) } diff --git a/internal/pipeline/runner.go b/internal/pipeline/runner.go index c9393f2..544559b 100644 --- a/internal/pipeline/runner.go +++ b/internal/pipeline/runner.go @@ -2,7 +2,6 @@ package pipeline import ( "context" - "errors" "fmt" "log/slog" "net/http" @@ -14,6 +13,7 @@ import ( "typetype-downloader-go/internal/config" "typetype-downloader-go/internal/job" "typetype-downloader-go/internal/selector" + "typetype-downloader-go/internal/storage" "typetype-downloader-go/internal/typetype" ) @@ -22,16 +22,18 @@ type Runner struct { store *job.Store streams *typetype.Client storage artifact.Store + disk *storage.Monitor http *http.Client queue chan string } -func NewRunner(cfg config.Config, store *job.Store, storage artifact.Store) *Runner { +func NewRunner(cfg config.Config, store *job.Store, files artifact.Store, disk *storage.Monitor) *Runner { return &Runner{ cfg: cfg, store: store, streams: typetype.NewClient(cfg.TypeTypeAPIBase), - storage: storage, + storage: files, + disk: disk, http: newHTTPClient(cfg.DownloadWorkers, cfg.HTTP2), queue: make(chan string, cfg.MaxQueueSize), } @@ -87,13 +89,6 @@ func (r *Runner) process(parent context.Context, id string) { } } -func failureCode(ctx context.Context, err error) string { - if ctx.Err() != nil || errors.Is(err, context.Canceled) { - return "cancelled" - } - return "download_failed" -} - func (r *Runner) run(ctx context.Context, id string, record *job.Record) error { started := time.Now() var stream *typetype.StreamResponse @@ -111,6 +106,11 @@ func (r *Runner) run(ctx context.Context, id string, record *job.Record) error { if err != nil { return err } + release, err := r.reserveVideo(id, selection, stream.Duration) + if err != nil { + return err + } + defer release() if selection.Video.DeliveryMethod == "sabr" { return r.runSABR(ctx, id, record, stream.Title, selection) } @@ -174,6 +174,11 @@ func (r *Runner) runAudioOnly(ctx context.Context, id string, record *job.Record if err != nil { return err } + release, err := r.reserveAudio(id, selection, stream.Duration) + if err != nil { + return err + } + defer release() if selection.Audio.DeliveryMethod == "sabr" { return r.runSABRAudio(ctx, id, record, stream.Title, selection) } diff --git a/internal/pipeline/sabr.go b/internal/pipeline/sabr.go index 0f4cb68..0080bd1 100644 --- a/internal/pipeline/sabr.go +++ b/internal/pipeline/sabr.go @@ -24,10 +24,9 @@ func (r *Runner) downloadSABR(ctx context.Context, id string, record *job.Record VideoItag: selection.Video.Itag, AudioItag: selection.Audio.Itag, AudioTrackID: trackID, - Workers: r.cfg.DownloadWorkers, - WorkDir: paths.WorkDir, VideoPath: paths.Video, AudioPath: paths.Audio, + ExpectedBytes: totalBytes, }, r.sabrProgress(id, started, totalBytes)) return time.Since(started).Milliseconds(), err } @@ -46,9 +45,8 @@ func (r *Runner) downloadSABRAudio(ctx context.Context, id string, record *job.R AudioItag: selection.Audio.Itag, AudioTrackID: trackID, AudioOnly: true, - Workers: r.cfg.DownloadWorkers, - WorkDir: paths.WorkDir, AudioPath: paths.Output, + ExpectedBytes: totalBytes, }, r.sabrProgress(id, started, totalBytes)) return time.Since(started).Milliseconds(), err } diff --git a/internal/pipeline/sabr_runner.go b/internal/pipeline/sabr_runner.go index 08672cf..d36a22c 100644 --- a/internal/pipeline/sabr_runner.go +++ b/internal/pipeline/sabr_runner.go @@ -14,30 +14,40 @@ import ( func (r *Runner) runSABR(ctx context.Context, id string, record *job.Record, title string, selection *selector.Selection) error { paths := artifact.Build(r.cfg.DataDir, id, title, selection.Container) r.store.Resolve(id, title, resolvedOutput(selection, paths.Name)) - return r.runSABRArtifact(ctx, id, paths, func() (int64, error) { + return r.runSABRArtifact(ctx, id, paths, func() (artifactTimings, error) { downloadMs, err := r.downloadSABR(ctx, id, record, selection, paths) if err != nil { - return downloadMs, err + return artifactTimings{downloadMs: downloadMs}, err } totalBytes := selection.Video.ContentLength + selection.Audio.ContentLength r.store.Progress(id, job.Progress{Stage: "mux", DownloadedBytes: totalBytes, TotalBytes: totalBytes}) + muxStarted := time.Now() err = retry(ctx, 2, "mux", func() error { _ = os.Remove(paths.Output) return merge(ctx, r.cfg.Muxer, paths.Video, paths.Audio, paths.Output) }) - return downloadMs, err + return artifactTimings{ + downloadMs: downloadMs, + muxMs: time.Since(muxStarted).Milliseconds(), + }, err }) } func (r *Runner) runSABRAudio(ctx context.Context, id string, record *job.Record, title string, selection *selector.AudioSelection) error { paths := artifact.Build(r.cfg.DataDir, id, title, selection.Container) r.store.Resolve(id, title, audioResolvedOutput(selection, paths.Name)) - return r.runSABRArtifact(ctx, id, paths, func() (int64, error) { - return r.downloadSABRAudio(ctx, id, record, selection, paths) + return r.runSABRArtifact(ctx, id, paths, func() (artifactTimings, error) { + downloadMs, err := r.downloadSABRAudio(ctx, id, record, selection, paths) + return artifactTimings{downloadMs: downloadMs}, err }) } -func (r *Runner) runSABRArtifact(ctx context.Context, id string, paths artifact.Paths, download func() (int64, error)) error { +func (r *Runner) runSABRArtifact( + ctx context.Context, + id string, + paths artifact.Paths, + process func() (artifactTimings, error), +) error { preserveOutput := false defer func() { cleanupWork(paths, preserveOutput) }() if err := os.MkdirAll(paths.WorkDir, 0o755); err != nil { @@ -46,7 +56,7 @@ func (r *Runner) runSABRArtifact(ctx context.Context, id string, paths artifact. if err := os.MkdirAll(filepath.Dir(paths.Output), 0o755); err != nil { return err } - downloadMs, err := download() + timings, err := process() if err != nil { return err } @@ -62,7 +72,12 @@ func (r *Runner) runSABRArtifact(ctx context.Context, id string, paths artifact. if !saved.Expires.IsZero() { expires = &saved.Expires } - r.store.Done(id, saved.Location, saved.Backend, expires, downloadMs, 0) + r.store.Done(id, saved.Location, saved.Backend, expires, timings.downloadMs, timings.muxMs) preserveOutput = saved.Backend == "local" return nil } + +type artifactTimings struct { + downloadMs int64 + muxMs int64 +} diff --git a/internal/pipeline/sabr_runner_test.go b/internal/pipeline/sabr_runner_test.go new file mode 100644 index 0000000..e370d50 --- /dev/null +++ b/internal/pipeline/sabr_runner_test.go @@ -0,0 +1,57 @@ +package pipeline + +import ( + "context" + "net/http" + "path/filepath" + "testing" + + "typetype-downloader-go/internal/artifact" + "typetype-downloader-go/internal/job" +) + +func TestRunSABRArtifactRecordsProcessingTimings(t *testing.T) { + store := job.NewStore("http://localhost") + store.Restore([]*job.Record{{ID: "job", Status: job.StatusRunning}}) + runner := &Runner{store: store, storage: timingArtifactStore{}} + root := t.TempDir() + paths := artifact.Paths{ + WorkDir: filepath.Join(root, "work"), + Output: filepath.Join(root, "output.mp4"), + Key: "artifact", + } + + err := runner.runSABRArtifact(t.Context(), "job", paths, func() (artifactTimings, error) { + return artifactTimings{downloadMs: 123, muxMs: 45}, nil + }) + if err != nil { + t.Fatal(err) + } + record, ok := store.Get("job") + if !ok || record.DownloadMs == nil || *record.DownloadMs != 123 || + record.MuxMs == nil || *record.MuxMs != 45 { + t.Fatalf("record = %#v", record) + } +} + +type timingArtifactStore struct{} + +func (timingArtifactStore) Name() string { + return "test" +} + +func (timingArtifactStore) Health(context.Context) error { + return nil +} + +func (timingArtifactStore) Save(context.Context, string, string) (artifact.Saved, error) { + return artifact.Saved{Backend: "test", Location: "artifact"}, nil +} + +func (timingArtifactStore) ServeHTTP(http.ResponseWriter, *http.Request, artifact.Saved, string) error { + return nil +} + +func (timingArtifactStore) Delete(context.Context, artifact.Saved) error { + return nil +} diff --git a/internal/sabr/download.go b/internal/sabr/download.go index 6312a38..08d6f50 100644 --- a/internal/sabr/download.go +++ b/internal/sabr/download.go @@ -5,84 +5,78 @@ import ( "fmt" "net/http" "net/url" - "os" - "path/filepath" "strconv" "strings" + "time" ) func Download(ctx context.Context, client *http.Client, options Options, progress ProgressFunc) error { - manifestURL, err := buildManifestURL(options) - if err != nil { - return err - } - tracks, err := fetchManifest(ctx, client, manifestURL, options.Authorization, options.AudioOnly) - if err != nil { - return err - } - tempDir, err := os.MkdirTemp(options.WorkDir, "sabr-") - if err != nil { - return err - } - defer os.RemoveAll(tempDir) - files, plans := planFiles(tracks, tempDir, options) + parts := downloadPartCount(options) reporter := newReporter(progress) - if err := downloadFiles(ctx, client, files, options.Authorization, options.Workers, reporter); err != nil { - return err - } - for _, plan := range plans { - if err := assemble(ctx, plan.Target, plan.Parts); err != nil { - return err - } - } - reporter.finish() - return nil -} - -func fetchManifest(ctx context.Context, client *http.Client, rawURL string, authorization string, audioOnly bool) ([]Track, error) { var last error - for attempt := 1; attempt <= 4; attempt++ { - response, err := request(ctx, client, rawURL, authorization) - if err == nil { - tracks, parseErr := parseManifest(response.Body, response.Request.URL, audioOnly) - response.Body.Close() - if parseErr != nil { - return nil, parseErr - } - return tracks, nil + for attempt := 1; attempt <= downloadAttempts; attempt++ { + reporter.beginAttempt() + if err := downloadAttempt(ctx, client, options, parts, reporter); err == nil { + reporter.finish() + return nil + } else { + last = err } - last = err - if attempt < 4 { + if ctx.Err() != nil { + return ctx.Err() + } + if attempt < downloadAttempts { if err := retryDelay(ctx, attempt); err != nil { - return nil, err + return err } } } - return nil, fmt.Errorf("fetch SABR manifest failed after 4 attempts: %w", last) + return fmt.Errorf("SABR download failed after %d attempts: %w", downloadAttempts, last) +} + +func downloadAttempt( + ctx context.Context, + client *http.Client, + options Options, + parts int, + progress *reporter, +) error { + defer cleanupDownloadFiles(options, parts) + if err := downloadParts(ctx, client, options, parts, progress); err != nil { + return err + } + return assembleDownload(options, parts) } -func buildManifestURL(options Options) (string, error) { +func buildDownloadURL(options Options, part int, parts int) (string, error) { parsed, err := url.Parse(options.ManifestURL) if err != nil { return "", err } query := parsed.Query() - query.Set("workload", "download") query.Set("audioItag", strconv.Itoa(options.AudioItag)) if options.VideoItag > 0 { query.Set("videoItag", strconv.Itoa(options.VideoItag)) } + parsed.Path = strings.Replace(parsed.Path, "/sabr/manifest/", "/sabr/download/", 1) if options.AudioTrackID != "" { query.Set("audioTrackId", options.AudioTrackID) } if options.AudioOnly { query.Set("audioOnly", "true") } + query.Set("part", strconv.Itoa(part)) + query.Set("parts", strconv.Itoa(parts)) parsed.RawQuery = query.Encode() return parsed.String(), nil } -func request(ctx context.Context, client *http.Client, rawURL string, authorization string) (*http.Response, error) { +func requestDownload( + ctx context.Context, + client *http.Client, + rawURL string, + authorization string, +) (*http.Response, error) { req, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil) if err != nil { return nil, err @@ -98,36 +92,54 @@ func request(ctx context.Context, client *http.Client, rawURL string, authorizat response.Body.Close() return nil, fmt.Errorf("GET %s returned %s", req.URL.Path, response.Status) } + if mediaType := strings.TrimSpace(strings.Split(response.Header.Get("Content-Type"), ";")[0]); mediaType != downloadMediaType { + response.Body.Close() + return nil, fmt.Errorf("GET %s returned unexpected content type %q", req.URL.Path, mediaType) + } return response, nil } -func planFiles(tracks []Track, tempDir string, options Options) ([]filePlan, []trackPlan) { - plans := make([]trackPlan, 0, len(tracks)) - for trackIndex, track := range tracks { - target := options.AudioPath - if track.Kind == "video" { - target = options.VideoPath - } - parts := make([]string, 0, len(track.URLs)) - for partIndex := range track.URLs { - path := filepath.Join(tempDir, fmt.Sprintf("%d-%06d.part", trackIndex, partIndex)) - parts = append(parts, path) +const downloadAttempts = 2 + +func downloadPartCount(options Options) int { + if options.Parts > 0 { + if options.Parts > maxDownloadParts { + return maxDownloadParts } - plans = append(plans, trackPlan{Parts: parts, Target: target}) + return options.Parts } - files := make([]filePlan, 0) - for partIndex := 0; ; partIndex++ { - added := false - for trackIndex, track := range tracks { - if partIndex >= len(track.URLs) { - continue - } - files = append(files, filePlan{URL: track.URLs[partIndex], Path: plans[trackIndex].Parts[partIndex]}) - added = true + if options.AudioOnly { + if options.ExpectedBytes >= 16<<20 { + return 4 } - if !added { - break + if options.ExpectedBytes >= 4<<20 { + return 2 } + return 1 + } + switch { + case options.ExpectedBytes >= 256<<20: + return 12 + case options.ExpectedBytes >= 128<<20: + return 6 + case options.ExpectedBytes >= 16<<20: + return 4 + case options.ExpectedBytes >= 4<<20: + return 2 + default: + return 1 } - return files, plans } + +func retryDelay(ctx context.Context, attempt int) error { + timer := time.NewTimer(time.Duration(attempt) * 100 * time.Millisecond) + defer timer.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil + } +} + +const maxDownloadParts = 12 diff --git a/internal/sabr/download_benchmark_test.go b/internal/sabr/download_benchmark_test.go new file mode 100644 index 0000000..0d2c3f4 --- /dev/null +++ b/internal/sabr/download_benchmark_test.go @@ -0,0 +1,95 @@ +package sabr + +import ( + "bytes" + "context" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strconv" + "testing" +) + +func BenchmarkDownloadVideoFifteenMinutes(b *testing.B) { + benchmarkDownload(b, false, 180, 1) +} + +func BenchmarkDownloadAudioThirtyMinutes(b *testing.B) { + benchmarkDownload(b, true, 360, 1) +} + +func BenchmarkDownloadVideoFifteenMinutesMultipart(b *testing.B) { + benchmarkDownload(b, false, 180, 4) +} + +func BenchmarkDownloadAudioThirtyMinutesMultipart(b *testing.B) { + benchmarkDownload(b, true, 360, 4) +} + +func BenchmarkDownloadVideoTenHours(b *testing.B) { + benchmarkDownload(b, false, 7200, 12) +} + +func BenchmarkDownloadAudioTenHours(b *testing.B) { + benchmarkDownload(b, true, 7200, 4) +} + +func benchmarkDownload(b *testing.B, audioOnly bool, segments int, parts int) { + streams := make([][]byte, parts) + totalBytes := int64(0) + for part := range parts { + streams[part] = benchmarkStreamPart(b, audioOnly, segments, part, parts) + totalBytes += int64(len(streams[part])) + } + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + writer.Header().Set("Content-Type", downloadMediaType) + part, _ := strconv.Atoi(request.URL.Query().Get("part")) + _, _ = writer.Write(streams[part]) + })) + b.Cleanup(server.Close) + b.ReportAllocs() + b.SetBytes(totalBytes) + b.ResetTimer() + for range b.N { + dir, err := os.MkdirTemp("", "sabr-stream-benchmark-") + if err != nil { + b.Fatal(err) + } + options := Options{ + ManifestURL: server.URL, + VideoItag: 137, + AudioItag: 140, + AudioOnly: audioOnly, + VideoPath: filepath.Join(dir, "video.mp4"), + AudioPath: filepath.Join(dir, "audio.m4a"), + Parts: parts, + } + if err := Download(context.Background(), server.Client(), options, nil); err != nil { + _ = os.RemoveAll(dir) + b.Fatal(err) + } + _ = os.RemoveAll(dir) + } +} + +func benchmarkStreamPart(t testing.TB, audioOnly bool, segments int, part int, parts int) []byte { + t.Helper() + payload := bytes.Repeat([]byte{0x5a}, 16<<10) + var stream bytes.Buffer + stream.Write(downloadMagic) + writeTestFrame(t, &stream, frameInitialization, 140, 0, []byte("audio-init")) + if !audioOnly { + writeTestFrame(t, &stream, frameInitialization, 137, 0, []byte("video-init")) + } + start := segments*part/parts + 1 + end := segments * (part + 1) / parts + for sequence := start; sequence <= end; sequence++ { + writeTestFrame(t, &stream, frameMedia, 140, sequence, payload) + if !audioOnly { + writeTestFrame(t, &stream, frameMedia, 137, sequence, payload) + } + } + writeTestFrame(t, &stream, frameComplete, 0, 0, nil) + return stream.Bytes() +} diff --git a/internal/sabr/download_test.go b/internal/sabr/download_test.go index f8a10df..5f79129 100644 --- a/internal/sabr/download_test.go +++ b/internal/sabr/download_test.go @@ -2,42 +2,38 @@ package sabr import ( "context" + "encoding/binary" "errors" - "fmt" + "io" "net/http" "net/http/httptest" "os" "path/filepath" - "sync" + "sync/atomic" "testing" "time" ) -func TestDownloadPreservesOrderAndBoundsConcurrency(t *testing.T) { - var mu sync.Mutex - active := 0 - maximum := 0 - server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { +func TestDownloadWritesOrderedTracksAtomically(t *testing.T) { + server := downloadTestServer(t, func(writer http.ResponseWriter, request *http.Request) { if request.Header.Get("Authorization") != "Bearer test" { t.Errorf("authorization = %q", request.Header.Get("Authorization")) } - if request.URL.Path == "/manifest" { - writer.Header().Set("Content-Type", "application/dash+xml") - fmt.Fprint(writer, testManifest(false)) - return - } - mu.Lock() - active++ - if active > maximum { - maximum = active + if request.URL.Path != "/sabr/download/video" || + request.URL.Query().Get("audioItag") != "140" || + request.URL.Query().Get("videoItag") != "137" || + request.URL.Query().Get("part") != "0" || + request.URL.Query().Get("parts") != "1" { + t.Errorf("query = %q", request.URL.RawQuery) } - mu.Unlock() - time.Sleep(30 * time.Millisecond) - fmt.Fprint(writer, request.URL.Path) - mu.Lock() - active-- - mu.Unlock() - })) + writeTestFrame(t, writer, frameInitialization, 140, 0, []byte("ai")) + writeTestFrame(t, writer, frameInitialization, 137, 0, []byte("vi")) + writeTestFrame(t, writer, frameMedia, 137, 1, []byte("v1")) + writeTestFrame(t, writer, frameMedia, 140, 1, []byte("a1")) + writeTestFrame(t, writer, frameMedia, 140, 2, []byte("a2")) + writeTestFrame(t, writer, frameMedia, 137, 2, []byte("v2")) + writeTestFrame(t, writer, frameComplete, 0, 0, nil) + }) defer server.Close() dir := t.TempDir() @@ -45,90 +41,236 @@ func TestDownloadPreservesOrderAndBoundsConcurrency(t *testing.T) { audio := filepath.Join(dir, "audio.m4a") var progress int64 err := Download(context.Background(), server.Client(), Options{ - ManifestURL: server.URL + "/manifest", Authorization: "Bearer test", - VideoItag: 137, AudioItag: 140, Workers: 4, WorkDir: dir, - VideoPath: video, AudioPath: audio, + ManifestURL: server.URL + "/sabr/manifest/video", Authorization: "Bearer test", + VideoItag: 137, AudioItag: 140, VideoPath: video, AudioPath: audio, }, func(value int64) { progress = value }) if err != nil { t.Fatal(err) } - assertFile(t, video, "/video/init/video/1/video/2") - assertFile(t, audio, "/audio/init/audio/1/audio/2") - if maximum < 2 || maximum > 4 { - t.Fatalf("maximum concurrency = %d, want 2..4", maximum) - } - if progress != int64(len("/video/init/video/1/video/2/audio/init/audio/1/audio/2")) { - t.Fatalf("progress = %d", progress) + assertFile(t, video, "viv1v2") + assertFile(t, audio, "aia1a2") + if progress != 12 { + t.Fatalf("progress = %d, want 12", progress) } } -func TestDownloadRetriesTransientSegmentFailure(t *testing.T) { - attempts := 0 - server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { - if request.URL.Path == "/manifest" { - fmt.Fprint(writer, testManifest(true)) - return +func TestDownloadAssemblesMultipartTracksWithoutDuplicateInitialization(t *testing.T) { + server := downloadTestServer(t, func(writer http.ResponseWriter, request *http.Request) { + part := request.URL.Query().Get("part") + switch part { + case "0": + writeTestFrame(t, writer, frameInitialization, 140, 0, []byte("ai")) + writeTestFrame(t, writer, frameInitialization, 137, 0, []byte("vi")) + writeTestFrame(t, writer, frameMedia, 140, 1, []byte("a1")) + writeTestFrame(t, writer, frameMedia, 137, 1, []byte("v1")) + case "1": + writeTestFrame(t, writer, frameInitialization, 140, 0, []byte("duplicate-ai")) + writeTestFrame(t, writer, frameInitialization, 137, 0, []byte("duplicate-vi")) + writeTestFrame(t, writer, frameMedia, 140, 7, []byte("a7")) + writeTestFrame(t, writer, frameMedia, 137, 4, []byte("v4")) + default: + t.Fatalf("unexpected part %q", part) } - if request.URL.Path == "/audio/1" { - attempts++ - if attempts == 1 { - http.Error(writer, "not ready", http.StatusNotFound) - return - } - } - fmt.Fprint(writer, request.URL.Path) - })) + writeTestFrame(t, writer, frameComplete, 0, 0, nil) + }) defer server.Close() + dir := t.TempDir() - output := filepath.Join(dir, "audio.m4a") + video := filepath.Join(dir, "video.mp4") + audio := filepath.Join(dir, "audio.m4a") + var progress int64 err := Download(context.Background(), server.Client(), Options{ - ManifestURL: server.URL + "/manifest", AudioItag: 140, - AudioOnly: true, Workers: 1, WorkDir: dir, AudioPath: output, - }, nil) + ManifestURL: server.URL, VideoItag: 137, AudioItag: 140, + VideoPath: video, AudioPath: audio, Parts: 2, + }, func(value int64) { progress = value }) if err != nil { t.Fatal(err) } - assertFile(t, output, "/audio/init/audio/1/audio/2") - if attempts != 2 { - t.Fatalf("attempts = %d, want 2", attempts) + assertFile(t, video, "viv1v4") + assertFile(t, audio, "aia1a7") + if progress != 12 { + t.Fatalf("progress = %d, want 12", progress) } + assertNoFile(t, partPath(video, 0)) + assertNoFile(t, partPath(video, 1)) + assertNoFile(t, partPath(audio, 0)) + assertNoFile(t, partPath(audio, 1)) } -func TestDownloadCancellationRemovesTemporaryFiles(t *testing.T) { - server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { - if request.URL.Path == "/manifest" { - fmt.Fprint(writer, testManifest(true)) +func TestDownloadRetriesTruncatedStream(t *testing.T) { + var attempts atomic.Int64 + server := downloadTestServer(t, func(writer http.ResponseWriter, _ *http.Request) { + if attempts.Add(1) == 1 { + writeTestFrame(t, writer, frameInitialization, 140, 0, []byte("init")) return } - <-request.Context().Done() - })) + writeTestFrame(t, writer, frameInitialization, 140, 0, []byte("init")) + writeTestFrame(t, writer, frameMedia, 140, 1, []byte("media")) + writeTestFrame(t, writer, frameComplete, 0, 0, nil) + }) + defer server.Close() + + output := filepath.Join(t.TempDir(), "audio.m4a") + err := Download(context.Background(), server.Client(), Options{ + ManifestURL: server.URL, AudioItag: 140, AudioOnly: true, AudioPath: output, + }, nil) + if err != nil { + t.Fatal(err) + } + assertFile(t, output, "initmedia") + if attempts.Load() != 2 { + t.Fatalf("attempts = %d, want 2", attempts.Load()) + } +} + +func TestReporterKeepsRetryProgressMonotonic(t *testing.T) { + var updates []int64 + reporter := newReporter(func(value int64) { + updates = append(updates, value) + }) + + reporter.beginAttempt() + reporter.reportMu.Lock() + reporter.last = time.Now().Add(-time.Second) + reporter.reportMu.Unlock() + reporter.add(10) + + reporter.beginAttempt() + reporter.reportMu.Lock() + reporter.last = time.Now().Add(-time.Second) + reporter.reportMu.Unlock() + reporter.add(5) + reporter.add(10) + reporter.finish() + + if len(updates) != 2 || updates[0] != 10 || updates[1] != 15 { + t.Fatalf("updates = %v, want [10 15]", updates) + } +} + +func TestDownloadRejectsOutOfOrderMediaAndPreservesTarget(t *testing.T) { + server := downloadTestServer(t, func(writer http.ResponseWriter, _ *http.Request) { + writeTestFrame(t, writer, frameInitialization, 140, 0, []byte("init")) + writeTestFrame(t, writer, frameMedia, 140, 2, []byte("wrong")) + }) defer server.Close() + dir := t.TempDir() + output := filepath.Join(dir, "audio.m4a") + if err := os.WriteFile(output, []byte("existing"), 0o600); err != nil { + t.Fatal(err) + } + err := Download(context.Background(), server.Client(), Options{ + ManifestURL: server.URL, AudioItag: 140, AudioOnly: true, AudioPath: output, + }, nil) + if err == nil { + t.Fatal("expected out-of-order stream failure") + } + assertFile(t, output, "existing") + assertNoFile(t, output+".download") +} + +func TestDownloadCancellationRemovesTemporaryFile(t *testing.T) { + server := downloadTestServer(t, func(writer http.ResponseWriter, request *http.Request) { + writeTestFrame(t, writer, frameInitialization, 140, 0, []byte("init")) + writer.(http.Flusher).Flush() + <-request.Context().Done() + }) + defer server.Close() + + output := filepath.Join(t.TempDir(), "audio.m4a") ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) defer cancel() err := Download(ctx, server.Client(), Options{ - ManifestURL: server.URL + "/manifest", AudioItag: 140, - AudioOnly: true, Workers: 1, WorkDir: dir, AudioPath: filepath.Join(dir, "audio.m4a"), + ManifestURL: server.URL, AudioItag: 140, AudioOnly: true, AudioPath: output, }, nil) if !errors.Is(err, context.DeadlineExceeded) { t.Fatalf("error = %v, want deadline exceeded", err) } - entries, readErr := os.ReadDir(dir) - if readErr != nil { - t.Fatal(readErr) + assertNoFile(t, output) + assertNoFile(t, output+".download") +} + +func TestDownloadRejectsOversizedFrame(t *testing.T) { + server := downloadTestServer(t, func(writer http.ResponseWriter, _ *http.Request) { + writeTestHeader(t, writer, frameInitialization, 140, 0, maxFrameBytes+1) + }) + defer server.Close() + output := filepath.Join(t.TempDir(), "audio.m4a") + err := Download(context.Background(), server.Client(), Options{ + ManifestURL: server.URL, AudioItag: 140, AudioOnly: true, AudioPath: output, + }, nil) + if err == nil { + t.Fatal("expected oversized frame failure") + } + assertNoFile(t, output) + assertNoFile(t, output+".download") +} + +func TestDownloadPartCountBoundsAudioFanout(t *testing.T) { + if got := downloadPartCount(Options{AudioOnly: true, ExpectedBytes: 34 << 20}); got != 4 { + t.Fatalf("audio parts = %d, want 4", got) + } + if got := downloadPartCount(Options{ExpectedBytes: 34 << 20}); got != 4 { + t.Fatalf("video parts = %d, want 4", got) + } + if got := downloadPartCount(Options{ExpectedBytes: 416 << 20}); got != 12 { + t.Fatalf("large video parts = %d, want 12", got) } - if len(entries) != 0 { - t.Fatalf("temporary files remain: %v", entries) + if got := downloadPartCount(Options{AudioOnly: true, ExpectedBytes: 34 << 20, Parts: 6}); got != 6 { + t.Fatalf("explicit audio parts = %d, want 6", got) + } + if got := downloadPartCount(Options{ExpectedBytes: 1 << 30, Parts: 20}); got != 12 { + t.Fatalf("capped parts = %d, want 12", got) + } +} + +func downloadTestServer( + t *testing.T, + handler func(http.ResponseWriter, *http.Request), +) *httptest.Server { + t.Helper() + return httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + writer.Header().Set("Content-Type", downloadMediaType) + if _, err := writer.Write(downloadMagic); err != nil { + return + } + handler(writer, request) + })) +} + +func writeTestFrame( + t testing.TB, + writer io.Writer, + kind byte, + itag int, + sequence int, + payload []byte, +) { + t.Helper() + writeTestHeader(t, writer, kind, itag, sequence, int64(len(payload))) + if _, err := writer.Write(payload); err != nil { + t.Fatal(err) } } -func testManifest(audioOnly bool) string { - audio := `` - if audioOnly { - return `` + audio + `` +func writeTestHeader( + t testing.TB, + writer io.Writer, + kind byte, + itag int, + sequence int, + length int64, +) { + t.Helper() + var header [frameHeaderSize]byte + header[0] = kind + binary.BigEndian.PutUint32(header[1:5], uint32(itag)) + binary.BigEndian.PutUint32(header[5:9], uint32(sequence)) + binary.BigEndian.PutUint64(header[9:17], uint64(length)) + if _, err := writer.Write(header[:]); err != nil { + t.Fatal(err) } - video := `` - return `` + video + audio + `` } func assertFile(t *testing.T, path string, expected string) { @@ -141,3 +283,10 @@ func assertFile(t *testing.T, path string, expected string) { t.Fatalf("%s = %q, want %q", path, content, expected) } } + +func assertNoFile(t *testing.T, path string) { + t.Helper() + if _, err := os.Stat(path); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("%s exists or stat failed: %v", path, err) + } +} diff --git a/internal/sabr/idle_timeout.go b/internal/sabr/idle_timeout.go new file mode 100644 index 0000000..6cae754 --- /dev/null +++ b/internal/sabr/idle_timeout.go @@ -0,0 +1,70 @@ +package sabr + +import ( + "context" + "errors" + "fmt" + "io" + "sync/atomic" + "time" +) + +var errDownloadIdleTimeout = errors.New("SABR download stream timed out") + +type idleWatchdog struct { + cancel context.CancelCauseFunc + last atomic.Int64 + timeout time.Duration +} + +func newIdleWatchdog(parent context.Context, timeout time.Duration) (context.Context, *idleWatchdog) { + if timeout <= 0 { + timeout = defaultDownloadIdleTimeout + } + ctx, cancel := context.WithCancelCause(parent) + watchdog := &idleWatchdog{cancel: cancel, timeout: timeout} + watchdog.touch() + go watchdog.run(ctx) + return ctx, watchdog +} + +func (w *idleWatchdog) run(ctx context.Context) { + timer := time.NewTimer(w.timeout) + defer timer.Stop() + for { + select { + case <-ctx.Done(): + return + case now := <-timer.C: + inactive := now.Sub(time.Unix(0, w.last.Load())) + if inactive >= w.timeout { + w.cancel(fmt.Errorf("%w after %s", errDownloadIdleTimeout, w.timeout)) + return + } + timer.Reset(w.timeout - inactive) + } + } +} + +func (w *idleWatchdog) touch() { + w.last.Store(time.Now().UnixNano()) +} + +func (w *idleWatchdog) stop() { + w.cancel(nil) +} + +type activityReader struct { + reader io.Reader + touch func() +} + +func (r activityReader) Read(buffer []byte) (int, error) { + count, err := r.reader.Read(buffer) + if count > 0 { + r.touch() + } + return count, err +} + +const defaultDownloadIdleTimeout = 60 * time.Second diff --git a/internal/sabr/idle_timeout_test.go b/internal/sabr/idle_timeout_test.go new file mode 100644 index 0000000..b3197b7 --- /dev/null +++ b/internal/sabr/idle_timeout_test.go @@ -0,0 +1,63 @@ +package sabr + +import ( + "context" + "errors" + "net/http" + "path/filepath" + "testing" + "time" +) + +func TestDownloadRetriesIdleStreamThenFails(t *testing.T) { + server := downloadTestServer(t, func(writer http.ResponseWriter, request *http.Request) { + writeTestFrame(t, writer, frameInitialization, 140, 0, []byte("init")) + writer.(http.Flusher).Flush() + <-request.Context().Done() + }) + defer server.Close() + + output := filepath.Join(t.TempDir(), "audio.m4a") + started := time.Now() + err := Download(context.Background(), server.Client(), Options{ + ManifestURL: server.URL, + AudioItag: 140, + AudioOnly: true, + AudioPath: output, + IdleTimeout: 30 * time.Millisecond, + }, nil) + if !errors.Is(err, errDownloadIdleTimeout) { + t.Fatalf("error = %v, want idle timeout", err) + } + if elapsed := time.Since(started); elapsed > time.Second { + t.Fatalf("idle download took %s", elapsed) + } + assertNoFile(t, output) + assertNoFile(t, output+".download") +} + +func TestDownloadIdleTimeoutSlidesWithStreamActivity(t *testing.T) { + server := downloadTestServer(t, func(writer http.ResponseWriter, _ *http.Request) { + writeTestFrame(t, writer, frameInitialization, 140, 0, []byte("init")) + for sequence := 1; sequence <= 4; sequence++ { + time.Sleep(30 * time.Millisecond) + writeTestFrame(t, writer, frameMedia, 140, sequence, []byte("media")) + writer.(http.Flusher).Flush() + } + writeTestFrame(t, writer, frameComplete, 0, 0, nil) + }) + defer server.Close() + + output := filepath.Join(t.TempDir(), "audio.m4a") + err := Download(context.Background(), server.Client(), Options{ + ManifestURL: server.URL, + AudioItag: 140, + AudioOnly: true, + AudioPath: output, + IdleTimeout: 100 * time.Millisecond, + }, nil) + if err != nil { + t.Fatal(err) + } + assertFile(t, output, "initmediamediamediamedia") +} diff --git a/internal/sabr/manifest.go b/internal/sabr/manifest.go deleted file mode 100644 index 953672e..0000000 --- a/internal/sabr/manifest.go +++ /dev/null @@ -1,95 +0,0 @@ -package sabr - -import ( - "encoding/xml" - "fmt" - "io" - "net/url" - "strings" -) - -type manifest struct { - Period period `xml:"Period"` -} - -type period struct { - AdaptationSets []adaptationSet `xml:"AdaptationSet"` -} - -type adaptationSet struct { - MimeType string `xml:"mimeType,attr"` - Representations []representation `xml:"Representation"` -} - -type representation struct { - Segments segmentList `xml:"SegmentList"` -} - -type segmentList struct { - Initialization initialization `xml:"Initialization"` - Segments []segment `xml:"SegmentURL"` -} - -type initialization struct { - SourceURL string `xml:"sourceURL,attr"` -} - -type segment struct { - Media string `xml:"media,attr"` -} - -type Track struct { - Kind string - URLs []string -} - -func parseManifest(reader io.Reader, base *url.URL, audioOnly bool) ([]Track, error) { - var document manifest - if err := xml.NewDecoder(reader).Decode(&document); err != nil { - return nil, fmt.Errorf("decode SABR manifest: %w", err) - } - tracks := make([]Track, 0, 2) - for _, adaptation := range document.Period.AdaptationSets { - kind := strings.TrimSuffix(adaptation.MimeType, "/mp4") - if kind != "audio" && (kind != "video" || audioOnly) { - continue - } - if len(adaptation.Representations) == 0 { - return nil, fmt.Errorf("SABR manifest %s track has no representation", kind) - } - list := adaptation.Representations[0].Segments - refs := make([]string, 0, len(list.Segments)+1) - refs = append(refs, list.Initialization.SourceURL) - for _, item := range list.Segments { - refs = append(refs, item.Media) - } - urls, err := resolveURLs(base, refs) - if err != nil { - return nil, fmt.Errorf("resolve SABR %s track: %w", kind, err) - } - tracks = append(tracks, Track{Kind: kind, URLs: urls}) - } - want := 2 - if audioOnly { - want = 1 - } - if len(tracks) != want { - return nil, fmt.Errorf("SABR manifest has %d usable tracks, want %d", len(tracks), want) - } - return tracks, nil -} - -func resolveURLs(base *url.URL, refs []string) ([]string, error) { - urls := make([]string, 0, len(refs)) - for _, raw := range refs { - if strings.TrimSpace(raw) == "" { - return nil, fmt.Errorf("empty segment URL") - } - ref, err := url.Parse(raw) - if err != nil { - return nil, err - } - urls = append(urls, base.ResolveReference(ref).String()) - } - return urls, nil -} diff --git a/internal/sabr/manifest_test.go b/internal/sabr/manifest_test.go deleted file mode 100644 index a2a3c23..0000000 --- a/internal/sabr/manifest_test.go +++ /dev/null @@ -1,42 +0,0 @@ -package sabr - -import ( - "net/url" - "strings" - "testing" -) - -func TestBuildManifestURLIncludesSelectedTracks(t *testing.T) { - result, err := buildManifestURL(Options{ - ManifestURL: "http://server/api/sabr/manifest/video", - VideoItag: 137, AudioItag: 140, AudioTrackID: "fr-FR.4", - }) - if err != nil { - t.Fatal(err) - } - parsed, err := url.Parse(result) - if err != nil { - t.Fatal(err) - } - query := parsed.Query() - if query.Get("videoItag") != "137" || query.Get("audioItag") != "140" || query.Get("audioTrackId") != "fr-FR.4" { - t.Fatalf("unexpected query: %s", parsed.RawQuery) - } - if query.Get("workload") != "download" { - t.Fatalf("unexpected workload: %s", parsed.RawQuery) - } -} - -func TestParseManifestResolvesRelativeURLs(t *testing.T) { - base, err := url.Parse("https://server/api/sabr/manifest/video") - if err != nil { - t.Fatal(err) - } - tracks, err := parseManifest(strings.NewReader(testManifest(false)), base, false) - if err != nil { - t.Fatal(err) - } - if tracks[0].URLs[1] != "https://server/video/1" || tracks[1].URLs[2] != "https://server/audio/2" { - t.Fatalf("unexpected tracks: %+v", tracks) - } -} diff --git a/internal/sabr/multipart.go b/internal/sabr/multipart.go new file mode 100644 index 0000000..3015cf6 --- /dev/null +++ b/internal/sabr/multipart.go @@ -0,0 +1,83 @@ +package sabr + +import ( + "context" + "errors" + "fmt" + "net/http" + "sync" +) + +func downloadParts( + ctx context.Context, + client *http.Client, + options Options, + parts int, + progress *reporter, +) error { + ctx, cancel := context.WithCancel(ctx) + defer cancel() + failures := make(chan error, parts) + var workers sync.WaitGroup + for part := range parts { + workers.Add(1) + go func() { + defer workers.Done() + if err := downloadPart(ctx, client, options, part, parts, progress); err != nil { + select { + case failures <- err: + default: + } + cancel() + } + }() + } + workers.Wait() + close(failures) + for err := range failures { + if !errors.Is(err, context.Canceled) { + return err + } + } + if err := ctx.Err(); err != nil { + return err + } + return nil +} + +func downloadPart( + ctx context.Context, + client *http.Client, + options Options, + part int, + parts int, + progress *reporter, +) error { + partCtx, watchdog := newIdleWatchdog(ctx, options.IdleTimeout) + defer watchdog.stop() + rawURL, err := buildDownloadURL(options, part, parts) + if err != nil { + return err + } + response, err := requestDownload(partCtx, client, rawURL, options.Authorization) + if err != nil { + if cause := context.Cause(partCtx); cause != nil { + err = cause + } + return fmt.Errorf("download SABR part %d/%d: %w", part+1, parts, err) + } + defer response.Body.Close() + tracks, err := openPartTracks(options, part) + if err != nil { + return err + } + defer closeTracks(tracks) + reader := activityReader{reader: response.Body, touch: watchdog.touch} + if err := consumeDownloadStream(reader, tracks, progress, part == 0); err != nil { + if cause := context.Cause(partCtx); cause != nil { + err = cause + } + return fmt.Errorf("consume SABR part %d/%d: %w", part+1, parts, err) + } + return closeTracks(tracks) +} diff --git a/internal/sabr/network_test.go b/internal/sabr/network_test.go new file mode 100644 index 0000000..3dbce79 --- /dev/null +++ b/internal/sabr/network_test.go @@ -0,0 +1,94 @@ +package sabr + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "io" + "net/http" + "os" + "path/filepath" + "strconv" + "testing" + "time" +) + +func TestNetworkDownload(t *testing.T) { + manifestURL := os.Getenv("TYPETYPE_SABR_MANIFEST_URL") + if manifestURL == "" { + t.Skip("set TYPETYPE_SABR_MANIFEST_URL to enable the network test") + } + transport := http.DefaultTransport.(*http.Transport).Clone() + transport.DisableCompression = true + transport.ForceAttemptHTTP2 = true + transport.MaxIdleConns = 8 + transport.MaxIdleConnsPerHost = 8 + transport.MaxConnsPerHost = 8 + client := &http.Client{Transport: transport} + t.Cleanup(transport.CloseIdleConnections) + + dir := os.Getenv("TYPETYPE_SABR_OUTPUT_DIR") + if dir == "" { + dir = t.TempDir() + } else if err := os.MkdirAll(dir, 0o755); err != nil { + t.Fatal(err) + } + started := time.Now() + err := Download(context.Background(), client, Options{ + ManifestURL: manifestURL, + Authorization: os.Getenv("TYPETYPE_SABR_AUTHORIZATION"), + VideoItag: networkInt("TYPETYPE_SABR_VIDEO_ITAG", 137), + AudioItag: networkInt("TYPETYPE_SABR_AUDIO_ITAG", 140), + AudioTrackID: os.Getenv("TYPETYPE_SABR_AUDIO_TRACK_ID"), + AudioOnly: os.Getenv("TYPETYPE_SABR_AUDIO_ONLY") == "1", + VideoPath: filepath.Join(dir, "video.mp4"), + AudioPath: filepath.Join(dir, "audio.m4a"), + Parts: networkInt("TYPETYPE_SABR_PARTS", 1), + }, nil) + if err != nil { + t.Fatal(err) + } + t.Logf("duration=%s bytes=%d", time.Since(started), outputBytes(t, dir)) + for _, name := range []string{"video.mp4", "audio.m4a"} { + path := filepath.Join(dir, name) + if _, err := os.Stat(path); err == nil { + t.Logf("%s sha256=%s", name, fileSHA256(t, path)) + } + } +} + +func fileSHA256(t *testing.T, path string) string { + t.Helper() + file, err := os.Open(path) + if err != nil { + t.Fatal(err) + } + defer file.Close() + hash := sha256.New() + if _, err := io.Copy(hash, file); err != nil { + t.Fatal(err) + } + return hex.EncodeToString(hash.Sum(nil)) +} + +func networkInt(name string, fallback int) int { + value, err := strconv.Atoi(os.Getenv(name)) + if err != nil || value <= 0 { + return fallback + } + return value +} + +func outputBytes(t *testing.T, dir string) int64 { + t.Helper() + var total int64 + for _, name := range []string{"video.mp4", "audio.m4a"} { + info, err := os.Stat(filepath.Join(dir, name)) + if err == nil { + total += info.Size() + } else if !os.IsNotExist(err) { + t.Fatal(err) + } + } + return total +} diff --git a/internal/sabr/part_files.go b/internal/sabr/part_files.go new file mode 100644 index 0000000..f4e044a --- /dev/null +++ b/internal/sabr/part_files.go @@ -0,0 +1,111 @@ +package sabr + +import ( + "fmt" + "io" + "os" + "strings" +) + +func openPartTracks(options Options, part int) ([]streamTrack, error) { + specs := trackSpecs(options) + tracks := make([]streamTrack, 0, len(specs)) + for _, spec := range specs { + if spec.itag <= 0 || strings.TrimSpace(spec.path) == "" { + closeTracks(tracks) + return nil, fmt.Errorf("missing SABR %s selection or target path", spec.kind) + } + spec.path = partPath(spec.path, part) + output, err := os.Create(spec.path) + if err != nil { + closeTracks(tracks) + return nil, fmt.Errorf("create SABR %s part: %w", spec.kind, err) + } + spec.output = output + tracks = append(tracks, spec) + } + return tracks, nil +} + +func assembleDownload(options Options, parts int) error { + specs := trackSpecs(options) + for index := range specs { + if err := assembleTrack(&specs[index], parts); err != nil { + return err + } + } + for index := range specs { + if err := os.Rename(specs[index].path+".download", specs[index].path); err != nil { + return fmt.Errorf("commit SABR %s track: %w", specs[index].kind, err) + } + } + return nil +} + +func assembleTrack(track *streamTrack, parts int) error { + output, err := os.Create(track.path + ".download") + if err != nil { + return fmt.Errorf("create SABR %s track: %w", track.kind, err) + } + ok := false + defer func() { + _ = output.Close() + if !ok { + _ = os.Remove(track.path + ".download") + } + }() + for part := range parts { + input, err := os.Open(partPath(track.path, part)) + if err != nil { + return fmt.Errorf("open SABR %s part: %w", track.kind, err) + } + _, copyErr := io.Copy(output, input) + closeErr := input.Close() + if copyErr != nil { + return fmt.Errorf("assemble SABR %s track: %w", track.kind, copyErr) + } + if closeErr != nil { + return fmt.Errorf("close SABR %s part: %w", track.kind, closeErr) + } + } + if err := output.Close(); err != nil { + return fmt.Errorf("close SABR %s track: %w", track.kind, err) + } + ok = true + return nil +} + +func cleanupDownloadFiles(options Options, parts int) { + for _, track := range trackSpecs(options) { + _ = os.Remove(track.path + ".download") + for part := range parts { + _ = os.Remove(partPath(track.path, part)) + } + } +} + +func closeTracks(tracks []streamTrack) error { + var first error + for index := range tracks { + if tracks[index].output == nil { + continue + } + if err := tracks[index].output.Close(); err != nil && first == nil { + first = fmt.Errorf("close SABR %s part: %w", tracks[index].kind, err) + } + tracks[index].output = nil + } + return first +} + +func trackSpecs(options Options) []streamTrack { + specs := []streamTrack{{kind: "audio", itag: options.AudioItag, path: options.AudioPath}} + if !options.AudioOnly { + specs = append(specs, streamTrack{kind: "video", itag: options.VideoItag, path: options.VideoPath}) + } + return specs +} + +func partPath(target string, part int) string { + return fmt.Sprintf("%s.download.part-%02d", target, part) +} diff --git a/internal/sabr/progress.go b/internal/sabr/progress.go index ea557af..0c16ca7 100644 --- a/internal/sabr/progress.go +++ b/internal/sabr/progress.go @@ -2,34 +2,55 @@ package sabr import ( "sync" + "sync/atomic" "time" ) type reporter struct { - mu sync.Mutex - downloaded int64 - last time.Time + downloaded atomic.Int64 progress ProgressFunc + reportMu sync.Mutex + reported int64 + last time.Time } func newReporter(progress ProgressFunc) *reporter { - return &reporter{last: time.Now(), progress: progress} + return &reporter{progress: progress, last: time.Now()} +} + +func (r *reporter) beginAttempt() { + r.downloaded.Store(0) + r.reportMu.Lock() + r.last = time.Now() + r.reportMu.Unlock() } func (r *reporter) add(bytes int64) { - r.mu.Lock() - defer r.mu.Unlock() - r.downloaded += bytes - if r.progress != nil && time.Since(r.last) >= 250*time.Millisecond { - r.last = time.Now() - r.progress(r.downloaded) + downloaded := r.downloaded.Add(bytes) + if r.progress == nil { + return + } + now := time.Now() + r.reportMu.Lock() + defer r.reportMu.Unlock() + if downloaded <= r.reported || now.Sub(r.last) < 250*time.Millisecond { + return } + r.reported = downloaded + r.last = now + r.progress(downloaded) } func (r *reporter) finish() { - r.mu.Lock() - defer r.mu.Unlock() - if r.progress != nil { - r.progress(r.downloaded) + if r.progress == nil { + return + } + r.reportMu.Lock() + defer r.reportMu.Unlock() + downloaded := r.downloaded.Load() + if downloaded <= r.reported { + return } + r.reported = downloaded + r.progress(downloaded) } diff --git a/internal/sabr/segments.go b/internal/sabr/segments.go deleted file mode 100644 index 0919c57..0000000 --- a/internal/sabr/segments.go +++ /dev/null @@ -1,150 +0,0 @@ -package sabr - -import ( - "context" - "fmt" - "io" - "net/http" - "os" - "sync" - "time" -) - -func downloadFiles(ctx context.Context, client *http.Client, files []filePlan, authorization string, workers int, progress *reporter) error { - if workers < 1 { - workers = 1 - } - if workers > 4 { - workers = 4 - } - workerCtx, cancel := context.WithCancel(ctx) - defer cancel() - queue := make(chan filePlan, len(files)) - for _, file := range files { - queue <- file - } - close(queue) - errs := make(chan error, 1) - var group sync.WaitGroup - for range workers { - group.Add(1) - go func() { - defer group.Done() - for file := range queue { - if err := downloadFile(workerCtx, client, file, authorization); err != nil { - select { - case errs <- err: - cancel() - default: - } - return - } - info, err := os.Stat(file.Path) - if err != nil { - select { - case errs <- err: - cancel() - default: - } - return - } - progress.add(info.Size()) - } - }() - } - group.Wait() - select { - case err := <-errs: - return err - default: - return ctx.Err() - } -} - -func downloadFile(ctx context.Context, client *http.Client, file filePlan, authorization string) error { - var last error - for attempt := 1; attempt <= 4; attempt++ { - response, err := request(ctx, client, file.URL, authorization) - if err == nil { - err = writeResponse(file.Path, response) - } - if err == nil { - return nil - } - last = err - if attempt < 4 { - if err := retryDelay(ctx, attempt); err != nil { - return err - } - } - } - return fmt.Errorf("download SABR segment failed after 4 attempts: %w", last) -} - -func retryDelay(ctx context.Context, attempt int) error { - select { - case <-ctx.Done(): - return ctx.Err() - case <-time.After(time.Duration(attempt) * 250 * time.Millisecond): - return nil - } -} - -func writeResponse(path string, response *http.Response) error { - defer response.Body.Close() - tempPath := path + ".download" - output, err := os.Create(tempPath) - if err != nil { - return err - } - written, copyErr := io.Copy(output, response.Body) - closeErr := output.Close() - if copyErr != nil || closeErr != nil || written == 0 { - _ = os.Remove(tempPath) - if copyErr != nil { - return copyErr - } - if closeErr != nil { - return closeErr - } - return fmt.Errorf("empty SABR segment") - } - return os.Rename(tempPath, path) -} - -func assemble(ctx context.Context, target string, parts []string) error { - output, err := os.Create(target) - if err != nil { - return err - } - completed := false - defer func() { - if !completed { - output.Close() - _ = os.Remove(target) - } - }() - for _, path := range parts { - if err := ctx.Err(); err != nil { - return err - } - input, openErr := os.Open(path) - if openErr != nil { - return openErr - } - _, copyErr := io.Copy(output, input) - closeErr := input.Close() - if copyErr != nil || closeErr != nil { - if copyErr != nil { - return copyErr - } - return closeErr - } - _ = os.Remove(path) - } - if err := output.Close(); err != nil { - return err - } - completed = true - return nil -} diff --git a/internal/sabr/stream_protocol.go b/internal/sabr/stream_protocol.go new file mode 100644 index 0000000..17db000 --- /dev/null +++ b/internal/sabr/stream_protocol.go @@ -0,0 +1,186 @@ +package sabr + +import ( + "bytes" + "encoding/binary" + "errors" + "fmt" + "io" + "sync" +) + +func consumeDownloadStream( + reader io.Reader, + tracks []streamTrack, + progress *reporter, + writeInitialization bool, +) error { + var magic [downloadMagicSize]byte + if _, err := io.ReadFull(reader, magic[:]); err != nil { + return fmt.Errorf("read SABR stream magic: %w", err) + } + if !bytes.Equal(magic[:], downloadMagic) { + return errors.New("invalid SABR download stream magic") + } + copyBuffer := downloadBufferPool.Get().([]byte) + defer downloadBufferPool.Put(copyBuffer) + var rawHeader [frameHeaderSize]byte + for { + if _, err := io.ReadFull(reader, rawHeader[:]); err != nil { + return fmt.Errorf("read SABR frame header: %w", err) + } + header, err := decodeFrameHeader(rawHeader[:]) + if err != nil { + return err + } + if header.kind == frameComplete { + if err := validateComplete(tracks, header); err != nil { + return err + } + var trailing [1]byte + if count, err := reader.Read(trailing[:]); count != 0 || err != io.EOF { + return errors.New("SABR download stream has trailing data") + } + return nil + } + track := findTrack(tracks, header.itag) + if track == nil { + return fmt.Errorf("SABR frame references unselected itag %d", header.itag) + } + if err := consumeFrame(reader, track, header, copyBuffer, progress, writeInitialization); err != nil { + return err + } + } +} + +func consumeFrame( + reader io.Reader, + track *streamTrack, + header frameHeader, + buffer []byte, + progress *reporter, + writeInitialization bool, +) error { + output := io.Writer(track.output) + countProgress := true + switch header.kind { + case frameInitialization: + if track.initialized || header.sequence != 0 { + return fmt.Errorf("invalid SABR initialization for itag %d", track.itag) + } + track.initialized = true + if !writeInitialization { + output = io.Discard + countProgress = false + } + case frameMedia: + if track.nextSequence == 0 { + track.nextSequence = header.sequence + } + if !track.initialized || header.sequence <= 0 || header.sequence != track.nextSequence { + return fmt.Errorf( + "out-of-order SABR media for itag %d: got %d, want %d", + track.itag, + header.sequence, + track.nextSequence, + ) + } + track.nextSequence++ + track.mediaWritten = true + default: + return fmt.Errorf("unknown SABR frame type %d", header.kind) + } + written, err := copyFrame(output, reader, header.length, buffer) + if err != nil { + return fmt.Errorf("write SABR %s track: %w", track.kind, err) + } + if written != header.length { + return io.ErrUnexpectedEOF + } + if countProgress { + progress.add(written) + } + return nil +} + +func copyFrame(output io.Writer, reader io.Reader, size int64, buffer []byte) (int64, error) { + var written int64 + for written < size { + chunk := buffer + if remaining := size - written; remaining < int64(len(chunk)) { + chunk = chunk[:remaining] + } + if _, err := io.ReadFull(reader, chunk); err != nil { + return written, err + } + count, err := output.Write(chunk) + written += int64(count) + if err != nil { + return written, err + } + if count != len(chunk) { + return written, io.ErrShortWrite + } + } + return written, nil +} + +func decodeFrameHeader(raw []byte) (frameHeader, error) { + header := frameHeader{ + kind: raw[0], + itag: int(int32(binary.BigEndian.Uint32(raw[1:5]))), + sequence: int(int32(binary.BigEndian.Uint32(raw[5:9]))), + length: int64(binary.BigEndian.Uint64(raw[9:17])), + } + if header.length < 0 || header.length > maxFrameBytes { + return frameHeader{}, fmt.Errorf("invalid SABR frame length %d", header.length) + } + return header, nil +} + +func validateComplete(tracks []streamTrack, header frameHeader) error { + if header.itag != 0 || header.sequence != 0 || header.length != 0 { + return errors.New("invalid SABR completion frame") + } + for index := range tracks { + if !tracks[index].initialized || !tracks[index].mediaWritten { + return fmt.Errorf("incomplete SABR %s track", tracks[index].kind) + } + } + return nil +} + +func findTrack(tracks []streamTrack, itag int) *streamTrack { + for index := range tracks { + if tracks[index].itag == itag { + return &tracks[index] + } + } + return nil +} + +type frameHeader struct { + kind byte + itag int + sequence int + length int64 +} + +const ( + downloadMediaType = "application/vnd.typetype.sabr-download" + frameInitialization = 1 + frameMedia = 2 + frameComplete = 3 + frameHeaderSize = 17 + downloadMagicSize = 8 + streamCopyBufferSize = 256 * 1024 + maxFrameBytes = 64 << 20 +) + +var downloadMagic = []byte("TTSABR1\n") + +var downloadBufferPool = sync.Pool{ + New: func() any { + return make([]byte, streamCopyBufferSize) + }, +} diff --git a/internal/sabr/types.go b/internal/sabr/types.go index 0077aa3..c607ed4 100644 --- a/internal/sabr/types.go +++ b/internal/sabr/types.go @@ -1,5 +1,10 @@ package sabr +import ( + "os" + "time" +) + type Options struct { ManifestURL string Authorization string @@ -7,20 +12,21 @@ type Options struct { AudioItag int AudioTrackID string AudioOnly bool - Workers int - WorkDir string VideoPath string AudioPath string + ExpectedBytes int64 + Parts int + IdleTimeout time.Duration } type ProgressFunc func(downloadedBytes int64) -type filePlan struct { - URL string - Path string -} - -type trackPlan struct { - Parts []string - Target string +type streamTrack struct { + kind string + itag int + path string + output *os.File + nextSequence int + initialized bool + mediaWritten bool } diff --git a/internal/storage/monitor.go b/internal/storage/monitor.go index cb48efc..72edbb0 100644 --- a/internal/storage/monitor.go +++ b/internal/storage/monitor.go @@ -1,23 +1,30 @@ package storage import ( + "errors" "fmt" "os" + "sync" "golang.org/x/sys/unix" ) +var ErrInsufficientStorage = errors.New("insufficient storage") + type Capacity struct { TotalBytes uint64 `json:"totalBytes"` FreeBytes uint64 `json:"freeBytes"` RequiredFreeBytes uint64 `json:"requiredFreeBytes"` + ReservedBytes uint64 `json:"reservedBytes"` Available bool `json:"available"` } type Monitor struct { + mu sync.Mutex dataDir string minFreeBytes uint64 minFreePercent uint64 + reservations map[string]uint64 } func NewMonitor(dataDir string, minFreeBytes int64, minFreePercent int) (*Monitor, error) { @@ -28,21 +35,61 @@ func NewMonitor(dataDir string, minFreeBytes int64, minFreePercent int) (*Monito dataDir: dataDir, minFreeBytes: uint64(max(minFreeBytes, 0)), minFreePercent: uint64(min(max(minFreePercent, 0), 100)), + reservations: make(map[string]uint64), }, nil } func (m *Monitor) Name() string { return "disk" } func (m *Monitor) Check() (Capacity, error) { + m.mu.Lock() + defer m.mu.Unlock() + return m.check() +} + +func (m *Monitor) Reserve(id string, bytes uint64) (func(), error) { + m.mu.Lock() + defer m.mu.Unlock() + capacity, err := m.check() + if err != nil { + return nil, err + } + previous := m.reservations[id] + reserved := capacity.ReservedBytes - previous + required := saturatedAdd(capacity.RequiredFreeBytes-previous, bytes) + if capacity.FreeBytes < required { + return nil, fmt.Errorf( + "%w: need %d bytes, have %d bytes with %d bytes reserved", + ErrInsufficientStorage, + required, + capacity.FreeBytes, + reserved, + ) + } + m.reservations[id] = bytes + return func() { + m.mu.Lock() + delete(m.reservations, id) + m.mu.Unlock() + }, nil +} + +func (m *Monitor) check() (Capacity, error) { var stat unix.Statfs_t if err := unix.Statfs(m.dataDir, &stat); err != nil { return Capacity{}, err } total := stat.Blocks * uint64(stat.Bsize) free := stat.Bavail * uint64(stat.Bsize) - required := max(m.minFreeBytes, total*m.minFreePercent/100) + reserved := uint64(0) + for _, bytes := range m.reservations { + reserved = saturatedAdd(reserved, bytes) + } + minimum := max(m.minFreeBytes, total*m.minFreePercent/100) + required := saturatedAdd(minimum, reserved) return Capacity{ - TotalBytes: total, FreeBytes: free, RequiredFreeBytes: required, Available: free >= required, + TotalBytes: total, FreeBytes: free, RequiredFreeBytes: required, + ReservedBytes: reserved, Available: free >= required, }, nil } @@ -56,3 +103,10 @@ func (m *Monitor) Health() error { } return fmt.Errorf("free bytes %d below required %d", capacity.FreeBytes, capacity.RequiredFreeBytes) } + +func saturatedAdd(left uint64, right uint64) uint64 { + if ^uint64(0)-left < right { + return ^uint64(0) + } + return left + right +} diff --git a/internal/storage/monitor_test.go b/internal/storage/monitor_test.go index fda7eb1..cd4c8df 100644 --- a/internal/storage/monitor_test.go +++ b/internal/storage/monitor_test.go @@ -1,6 +1,7 @@ package storage import ( + "errors" "math" "testing" ) @@ -38,3 +39,37 @@ func TestMonitorUsesPercentageThreshold(t *testing.T) { t.Fatalf("capacity = %#v", capacity) } } + +func TestMonitorTracksAndReleasesReservations(t *testing.T) { + monitor, err := NewMonitor(t.TempDir(), 1, 1) + if err != nil { + t.Fatal(err) + } + capacity, err := monitor.Check() + if err != nil { + t.Fatal(err) + } + reserved := capacity.FreeBytes - capacity.RequiredFreeBytes + release, err := monitor.Reserve("job", reserved) + if err != nil { + t.Fatal(err) + } + capacity, err = monitor.Check() + if err != nil { + t.Fatal(err) + } + if capacity.ReservedBytes != reserved || !capacity.Available { + t.Fatalf("capacity = %#v", capacity) + } + if _, err := monitor.Reserve("second", 1); !errors.Is(err, ErrInsufficientStorage) { + t.Fatalf("reserve error = %v", err) + } + release() + capacity, err = monitor.Check() + if err != nil { + t.Fatal(err) + } + if capacity.ReservedBytes != 0 { + t.Fatalf("capacity = %#v", capacity) + } +}