fix(hls): isolate staging per task and finalize output atomically

- Give every task a random staging identity persisted in the saved
  state: two tasks downloading the same playlist into the same
  directory get separate segment folders and journals instead of
  deleting each other's progress
- Bind the journal to a fingerprint of the full download plan (media
  URL, ordered segments, ranges, keys); foreign or unreadable journals
  are discarded rather than applied to the wrong segments, and the
  ledger is rewritten through a temp file + rename
- Legacy states without a staging identity migrate to a fresh folder
  instead of reusing the shared URL-hash folder
- Merge into a unique temporary file next to the target, check write
  and close errors, then remove-then-rename over the final name: a
  failed merge preserves the staged segments for a retry and re-merging
  after a crash stays idempotent
- Drop the staging folder when a URL patch invalidates the plan
This commit is contained in:
sun committed 2026-09-15 11:27:26 +08:00
1 parent b6d18a267b
commit 675028c64a
3 files changed
+471 -45

No files matched your search

+65 -45
View File
@@ -3,7 +3,6 @@ package hls
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"io"
@@ -46,11 +45,17 @@ func copyReqExtra(extra *ReqExtra) *ReqExtra {
return copied
}
// fetcherState is the resolved download plan persisted for resume.
// fetcherState is the resolved download plan persisted for resume. StagingID
// gives the task its own staging folder (two tasks downloading the same URL
// into the same directory never share segment files) and survives restarts
// through the persisted state. StagingDir pins the absolute folder so user
// renames or moves of the task output do not orphan the staged segments.
type fetcherState struct {
MediaURL string `json:"mediaURL"`
OutputName string `json:"outputName"`
Segments []*Segment `json:"segments"`
StagingID string `json:"stagingID,omitempty"`
StagingDir string `json:"stagingDir,omitempty"`
}
// Stats is the protocol task statistics snapshot.
@@ -86,6 +91,8 @@ type hlsRun struct {
keyCache map[string][]byte
keyCacheMu sync.Mutex
stagingID string // identity of the staging folder, bound into the journal
planHash string // fingerprint of the download plan, validated on resume
journal map[int64]int64
journalPath string
journalMu sync.Mutex
@@ -184,7 +191,12 @@ func (f *Fetcher) Resolve(req *base.Request, opts *base.Options) (err error) {
if err != nil {
return err
}
state.StagingID = randomStagingID()
f.state = state
// Pin the staging folder now so the very first persisted state already
// carries it: a crash right after create still resumes into the same
// folder.
f.resolveStagingDirLocked()
size := planTotalSize(state.Segments)
f.meta.Res = &base.Resource{
Range: false,
@@ -352,8 +364,13 @@ func mergeReqExtra(cur *base.Request, incoming *base.Request) error {
return nil
}
// dropPlanLocked discards the resolved plan after a URL change.
// dropPlanLocked discards the resolved plan after a URL change. The staging
// folder of the dropped plan is removed too: this only runs while the
// fetcher is quiescent, so no run can still be writing into it.
func (f *Fetcher) dropPlanLocked() {
if f.state != nil && f.state.StagingDir != "" {
os.RemoveAll(f.state.StagingDir)
}
f.state = nil
f.meta.Res = nil
}
@@ -421,7 +438,9 @@ func (f *Fetcher) Start() error {
extra: copyReqExtra(f.extra),
client: f.client,
outputPath: f.meta.SingleFilepath(),
tempDir: f.stagingDirLocked(),
stagingID: f.state.StagingID,
planHash: planFingerprint(f.state),
tempDir: f.resolveStagingDirLocked(),
keyCache: make(map[string][]byte),
journal: make(map[int64]int64),
}
@@ -464,7 +483,7 @@ func (f *Fetcher) Close() error {
run := f.activeRun
tempDir := ""
if f.state != nil {
tempDir = f.stagingDirLocked()
tempDir = f.resolveStagingDirLocked()
}
f.mu.Unlock()
@@ -757,15 +776,20 @@ func (f *Fetcher) fetchKey(run *hlsRun, key *Key) ([]byte, error) {
}
// merge concatenates the downloaded segments in plan order into the final
// output file. Segments are stored decrypted, so this is pure I/O. The final
// size is published through Progress: the engine backfills Res.Size from
// Progress().TotalDownloaded() on completion.
// output file. Segments are stored decrypted, so this is pure I/O. The merge
// writes to a unique temporary file next to the target, checks write and
// close errors, and only then replaces the target: a failed merge leaves the
// staged segments intact so it can be retried, and the unique part name keeps
// concurrent tasks that share an output path from clobbering each other. The
// final size is published through Progress: the engine backfills Res.Size
// from Progress().TotalDownloaded() on completion.
func (f *Fetcher) merge(run *hlsRun) error {
target := run.outputPath
if err := os.MkdirAll(filepath.Dir(target), 0777); err != nil {
return err
}
out, err := os.Create(target)
partPath := target + "." + tempDirPrefix + run.stagingID + ".part"
out, err := os.Create(partPath)
if err != nil {
return err
}
@@ -774,15 +798,29 @@ func (f *Fetcher) merge(run *hlsRun) error {
data, err := os.ReadFile(segmentPath(run.tempDir, int64(i)))
if err != nil {
out.Close()
os.Remove(partPath)
return fmt.Errorf("read segment %d failed: %w", i, err)
}
if _, err := out.Write(data); err != nil {
out.Close()
os.Remove(partPath)
return err
}
total += int64(len(data))
}
if err := out.Close(); err != nil {
os.Remove(partPath)
return err
}
// Remove-then-rename keeps this idempotent: if a previous merge already
// renamed its output before a crash lost the bookkeeping, re-merging
// must still succeed.
if err := os.Remove(target); err != nil && !os.IsNotExist(err) {
os.Remove(partPath)
return err
}
if err := os.Rename(partPath, target); err != nil {
os.Remove(partPath)
return err
}
run.downloaded.Store(total)
@@ -807,45 +845,27 @@ func segmentPath(tempDir string, idx int64) string {
return filepath.Join(tempDir, "seg-"+strconv.FormatInt(idx, 10))
}
// stagingDirLocked derives the staging folder for the current state. It must
// be called with f.mu held.
func (f *Fetcher) stagingDirLocked() string {
return filepath.Join(filepath.Dir(f.meta.SingleFilepath()), tempDirPrefix+sha1Hex(f.state.MediaURL)[:12])
}
func loadJournal(run *hlsRun) error {
data, err := os.ReadFile(run.journalPath)
if err != nil {
if os.IsNotExist(err) {
return nil
// resolveStagingDirLocked returns the staging folder of the current plan,
// deriving and pinning it on first use so later output renames or moves do
// not orphan the staged segments. It lives next to the task output, which is
// exactly opts.Path for a single-file task. Must be called with f.mu held.
func (f *Fetcher) resolveStagingDirLocked() string {
if f.state.StagingDir == "" {
if f.state.StagingID == "" {
// Legacy state (or a hand-crafted plan): migrate to a fresh task
// owned identity. The old URL-hash folder is deliberately not
// reused: its contents cannot be proven to belong to this plan,
// and sharing it with another same-URL task is exactly the bug
// being fixed. It is left behind as orphaned temp data.
f.state.StagingID = randomStagingID()
}
return err
}
stored := make(map[string]int64)
if err := json.Unmarshal(data, &stored); err != nil {
// A broken journal must not kill a restartable download.
return nil
}
for key, size := range stored {
idx, err := strconv.ParseInt(key, 10, 64)
if err != nil {
continue
dir := ""
if f.meta.Opts != nil {
dir = f.meta.Opts.Path
}
run.journal[idx] = size
f.state.StagingDir = filepath.Join(dir, tempDirPrefix+f.state.StagingID)
}
return nil
}
func saveJournal(run *hlsRun) error {
stored := make(map[string]int64, len(run.journal))
for idx, size := range run.journal {
stored[strconv.FormatInt(idx, 10)] = size
}
data, err := json.Marshal(stored)
if err != nil {
return err
}
return os.WriteFile(run.journalPath, data, 0666)
return f.state.StagingDir
}
// prefetchSize probes segment Content-Length via HEAD. Any failure resets all
+110
View File
@@ -0,0 +1,110 @@
package hls
import (
"crypto/rand"
"crypto/sha1"
"encoding/hex"
"encoding/json"
"fmt"
"io"
"os"
"strconv"
)
// journalFile is the persisted completed-segment ledger. It is bound to the
// exact download plan: a journal whose fingerprint does not match the current
// plan is discarded instead of being applied to the wrong segments.
type journalFile struct {
StagingID string `json:"stagingID"`
PlanHash string `json:"planHash"`
SegmentCount int `json:"segmentCount"`
Sizes map[string]int64 `json:"sizes"`
}
// randomStagingID returns a random identifier used to give every task its own
// staging folder, stable across app restarts through the persisted state.
func randomStagingID() string {
buf := make([]byte, 6)
if _, err := rand.Read(buf); err != nil {
// crypto/rand failing is catastrophic anyway; fall back to address
// entropy rather than crashing the fetch.
return fmt.Sprintf("%x", sha1.Sum([]byte(fmt.Sprintf("%p", buf))))[:12]
}
return hex.EncodeToString(buf)
}
// planFingerprint hashes everything that defines the download plan: the media
// URL, the ordered segments with their sequence numbers, byte ranges, keys
// and init markers. Two tasks resolve to different fingerprints even when
// they point at the same playlist URL.
func planFingerprint(state *fetcherState) string {
h := sha1.New()
io.WriteString(h, state.MediaURL+"\n")
io.WriteString(h, strconv.Itoa(len(state.Segments))+"\n")
for _, seg := range state.Segments {
fmt.Fprintf(h, "%d|%s|", seg.Sequence, seg.URI)
if seg.Byterange != nil {
fmt.Fprintf(h, "%d@%d|", seg.Byterange.Offset, seg.Byterange.Length)
}
if seg.Key != nil {
fmt.Fprintf(h, "key=%s|", seg.Key.URI)
}
if seg.IsInit {
io.WriteString(h, "init|")
}
io.WriteString(h, "\n")
}
return hex.EncodeToString(h.Sum(nil))[:16]
}
// loadJournal reads the journal of a run, discarding anything that cannot be
// proven to belong to the current plan: foreign or legacy journals, corrupted
// files, and entries whose recorded size no longer matches the plan.
func loadJournal(run *hlsRun) error {
data, err := os.ReadFile(run.journalPath)
if err != nil {
if os.IsNotExist(err) {
return nil
}
return err
}
var file journalFile
if err := json.Unmarshal(data, &file); err != nil {
// A broken or legacy journal must not kill a restartable download.
return nil
}
if file.StagingID != run.stagingID || file.PlanHash != run.planHash || file.SegmentCount != len(run.segs) {
return nil
}
for key, size := range file.Sizes {
idx, err := strconv.ParseInt(key, 10, 64)
if err != nil || idx < 0 || idx >= int64(len(run.segs)) || size <= 0 {
continue
}
run.journal[idx] = size
}
return nil
}
// saveJournal atomically rewrites the journal through a temp file + rename so
// a crash never leaves a half-written ledger behind.
func saveJournal(run *hlsRun) error {
file := journalFile{
StagingID: run.stagingID,
PlanHash: run.planHash,
SegmentCount: len(run.segs),
Sizes: make(map[string]int64, len(run.journal)),
}
for idx, size := range run.journal {
file.Sizes[strconv.FormatInt(idx, 10)] = size
}
data, err := json.Marshal(file)
if err != nil {
return err
}
tmp := run.journalPath + ".tmp"
if err := os.WriteFile(tmp, data, 0666); err != nil {
return err
}
return os.Rename(tmp, run.journalPath)
}
+296
View File
@@ -0,0 +1,296 @@
package hls
import (
"bytes"
"context"
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"strings"
"sync/atomic"
"testing"
"time"
"github.com/GopeedLab/gopeed/pkg/base"
)
func TestFetcherE2E_SameURLTwoTasks(t *testing.T) {
const total = 6
mux := http.NewServeMux()
var playlist strings.Builder
playlist.WriteString("#EXTM3U\n")
for i := 0; i < total; i++ {
playlist.WriteString(fmt.Sprintf("#EXTINF:4,\nseg-%d.ts\n", i))
}
playlist.WriteString("#EXT-X-ENDLIST\n")
mux.HandleFunc("/video.m3u8", func(w http.ResponseWriter, r *http.Request) {
w.Write([]byte(playlist.String()))
})
segments := make(map[int][]byte)
for i := 0; i < total; i++ {
segments[i] = []byte(fmt.Sprintf("SHARED-SEGMENT-%02d", i))
idx := i
mux.HandleFunc(fmt.Sprintf("/seg-%d.ts", idx), func(w http.ResponseWriter, r *http.Request) {
time.Sleep(120 * time.Millisecond)
w.Write(segments[idx])
})
}
server := httptest.NewServer(mux)
defer server.Close()
cfg := testConfig()
cfg.SegmentConnections = 2
cfg.PrefetchContentLength = false
dir := t.TempDir()
// Two tasks for the same playlist URL into the same directory, like a
// user duplicating a task. The engine would rename the second output
// (AutoRename); mirror that here.
f1 := newTestFetcher(t, cfg)
if err := f1.Resolve(&base.Request{URL: server.URL + "/video.m3u8"}, &base.Options{Path: dir}); err != nil {
t.Fatal(err)
}
f2 := newTestFetcher(t, cfg)
if err := f2.Resolve(&base.Request{URL: server.URL + "/video.m3u8"}, &base.Options{Path: dir, Name: "video-2.ts"}); err != nil {
t.Fatal(err)
}
if f1.state.StagingDir == "" || f1.state.StagingDir == f2.state.StagingDir {
t.Fatalf("staging dirs must be distinct per task: %q vs %q", f1.state.StagingDir, f2.state.StagingDir)
}
// Start task 1, let it stage some segments, then pause it.
if err := f1.Start(); err != nil {
t.Fatal(err)
}
time.Sleep(250 * time.Millisecond)
if err := f1.Pause(); err != nil {
t.Fatal(err)
}
staged1, err := os.ReadDir(f1.state.StagingDir)
if err != nil || len(staged1) == 0 {
t.Fatalf("task 1 should have staged segments, dir: %s", f1.state.StagingDir)
}
// Task 2 must complete without touching task 1's staging data...
if err := f2.Start(); err != nil {
t.Fatal(err)
}
if err := f2.Wait(); err != nil {
t.Fatal(err)
}
stillThere, err := os.ReadDir(f1.state.StagingDir)
if err != nil || len(stillThere) != len(staged1) {
t.Fatalf("task 2 disturbed task 1 staging: had %d entries, now %d", len(staged1), len(stillThere))
}
// ...and task 1 resumes to a complete, correct output of its own.
if err := f1.Start(); err != nil {
t.Fatal(err)
}
if err := f1.Wait(); err != nil {
t.Fatal(err)
}
var want []byte
for i := 0; i < total; i++ {
want = append(want, segments[i]...)
}
if !bytes.Equal(readOutput(t, f1), want) || !bytes.Equal(readOutput(t, f2), want) {
t.Error("same-URL task outputs mismatch")
}
entries, _ := os.ReadDir(dir)
if len(entries) != 2 {
names := make([]string, 0, len(entries))
for _, e := range entries {
names = append(names, e.Name())
}
t.Errorf("expected only the two outputs to remain, got %v", names)
}
}
func TestFetcherE2E_ForeignJournalDiscarded(t *testing.T) {
const total = 3
var served atomic.Int64
mux := http.NewServeMux()
var playlist strings.Builder
playlist.WriteString("#EXTM3U\n")
for i := 0; i < total; i++ {
playlist.WriteString(fmt.Sprintf("#EXTINF:4,\nseg-%d.ts\n", i))
}
playlist.WriteString("#EXT-X-ENDLIST\n")
mux.HandleFunc("/video.m3u8", func(w http.ResponseWriter, r *http.Request) {
w.Write([]byte(playlist.String()))
})
for i := 0; i < total; i++ {
mux.HandleFunc(fmt.Sprintf("/seg-%d.ts", i), func(w http.ResponseWriter, r *http.Request) {
served.Add(1)
w.Write([]byte("DATA"))
})
}
server := httptest.NewServer(mux)
defer server.Close()
cfg := testConfig()
cfg.PrefetchContentLength = false
f := newTestFetcher(t, cfg)
if err := f.Resolve(&base.Request{URL: server.URL + "/video.m3u8"}, &base.Options{Path: t.TempDir()}); err != nil {
t.Fatal(err)
}
// Plant a journal that does not belong to this plan (wrong fingerprint).
tempDir := f.resolveStagingDirLocked()
if err := os.MkdirAll(tempDir, 0777); err != nil {
t.Fatal(err)
}
foreign := journalFile{
StagingID: "some-other-task",
PlanHash: "deadbeef",
SegmentCount: total,
Sizes: map[string]int64{"0": 4, "1": 4, "2": 4},
}
data, _ := json.Marshal(foreign)
if err := os.WriteFile(filepath.Join(tempDir, "journal.json"), data, 0666); err != nil {
t.Fatal(err)
}
if err := f.Start(); err != nil {
t.Fatal(err)
}
if err := f.Wait(); err != nil {
t.Fatal(err)
}
if got := served.Load(); got != total {
t.Errorf("foreign journal must be discarded, want %d downloads, got %d", total, got)
}
}
func TestFetcherE2E_MergeFailurePreservesStaging(t *testing.T) {
mux := http.NewServeMux()
mux.HandleFunc("/video.m3u8", func(w http.ResponseWriter, r *http.Request) {
w.Write([]byte("#EXTM3U\n#EXTINF:4,\na.ts\n#EXTINF:4,\nb.ts\n#EXT-X-ENDLIST\n"))
})
mux.HandleFunc("/a.ts", func(w http.ResponseWriter, r *http.Request) { w.Write([]byte("AAAA")) })
mux.HandleFunc("/b.ts", func(w http.ResponseWriter, r *http.Request) { w.Write([]byte("BBBB")) })
server := httptest.NewServer(mux)
defer server.Close()
f := newTestFetcher(t, testConfig())
if err := f.Resolve(&base.Request{URL: server.URL + "/video.m3u8"}, &base.Options{Path: t.TempDir()}); err != nil {
t.Fatal(err)
}
run := &hlsRun{
ctx: context.Background(),
cancel: func() {},
done: make(chan struct{}),
segs: f.state.Segments,
outputPath: f.meta.SingleFilepath(),
stagingID: f.state.StagingID,
planHash: planFingerprint(f.state),
tempDir: f.resolveStagingDirLocked(),
keyCache: make(map[string][]byte),
journal: make(map[int64]int64),
}
run.journalPath = filepath.Join(run.tempDir, "journal.json")
if err := os.MkdirAll(run.tempDir, 0777); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(segmentPath(run.tempDir, 0), []byte("AAAA"), 0666); err != nil {
t.Fatal(err)
}
// Segment 1 is missing on disk: the merge must fail without deleting the
// staged data or leaving a partial output behind.
if err := f.merge(run); err == nil {
t.Fatal("merge with a missing segment should fail")
}
if _, err := os.Stat(segmentPath(run.tempDir, 0)); err != nil {
t.Fatal("merge failure must preserve staged segments")
}
entries, _ := os.ReadDir(filepath.Dir(run.outputPath))
for _, e := range entries {
if strings.HasSuffix(e.Name(), ".part") {
t.Errorf("failed merge left a partial output behind: %s", e.Name())
}
}
// Recover the segment: the retried merge succeeds.
if err := os.WriteFile(segmentPath(run.tempDir, 1), []byte("BBBB"), 0666); err != nil {
t.Fatal(err)
}
if err := f.merge(run); err != nil {
t.Fatalf("retried merge failed: %v", err)
}
got, err := os.ReadFile(run.outputPath)
if err != nil {
t.Fatal(err)
}
if string(got) != "AAAABBBB" {
t.Errorf("retried merge output mismatch: %q", got)
}
}
func TestFetcherE2E_ReMergeOverExistingOutput(t *testing.T) {
const total = 2
mux := http.NewServeMux()
var playlist strings.Builder
playlist.WriteString("#EXTM3U\n")
for i := 0; i < total; i++ {
playlist.WriteString(fmt.Sprintf("#EXTINF:4,\nseg-%d.ts\n", i))
}
playlist.WriteString("#EXT-X-ENDLIST\n")
mux.HandleFunc("/video.m3u8", func(w http.ResponseWriter, r *http.Request) {
w.Write([]byte(playlist.String()))
})
for i := 0; i < total; i++ {
mux.HandleFunc(fmt.Sprintf("/seg-%d.ts", i), func(w http.ResponseWriter, r *http.Request) {
w.Write([]byte("DATA"))
})
}
server := httptest.NewServer(mux)
defer server.Close()
f := newTestFetcher(t, testConfig())
if err := f.Resolve(&base.Request{URL: server.URL + "/video.m3u8"}, &base.Options{Path: t.TempDir()}); err != nil {
t.Fatal(err)
}
tempDir := f.resolveStagingDirLocked()
if err := os.MkdirAll(tempDir, 0777); err != nil {
t.Fatal(err)
}
// Simulate the crash window after a completed rename but before the task
// was marked done: a stale output file, a complete journal and all
// segments staged.
stale := filepath.Join(filepath.Dir(f.meta.SingleFilepath()), f.state.OutputName)
if err := os.WriteFile(stale, []byte("STALE-OUTPUT"), 0666); err != nil {
t.Fatal(err)
}
jr := &hlsRun{
segs: f.state.Segments,
stagingID: f.state.StagingID,
planHash: planFingerprint(f.state),
journalPath: filepath.Join(tempDir, "journal.json"),
journal: map[int64]int64{0: 4, 1: 4},
}
if err := saveJournal(jr); err != nil {
t.Fatal(err)
}
for i := 0; i < total; i++ {
if err := os.WriteFile(segmentPath(tempDir, int64(i)), []byte("DATA"), 0666); err != nil {
t.Fatal(err)
}
}
if err := f.Start(); err != nil {
t.Fatal(err)
}
if err := f.Wait(); err != nil {
t.Fatalf("re-merge over an existing output failed: %v", err)
}
got, err := os.ReadFile(stale)
if err != nil {
t.Fatal(err)
}
if string(got) != "DATADATA" {
t.Errorf("stale output was not replaced: %q", got)
}
}