Hojun-Cho
3d9ea7ce8b
an idle worker now steals the back half of the busiest in-flight segment instead of exiting, so one slow segment no longer drains alone over a single connection. control file persists stolen tails; resume re-tiles. add --auto-split (got-only, opt-in): derive the connection count from file size (one per min-split-size, up to 16), overriding -x/-s. default stays aria2-faithful at 1 connection for a single server.
365 lines
11 KiB
Go
365 lines
11 KiB
Go
package httpdl
|
|
|
|
import (
|
|
"os"
|
|
"path/filepath"
|
|
"sort"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
)
|
|
|
|
func TestMakeSegments(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
total int64
|
|
minSplit int64
|
|
conns int
|
|
wantN int
|
|
}{
|
|
{"even split", 1000, 100, 5, 5},
|
|
{"min-split caps count", 1000, 400, 5, 2},
|
|
{"tiny file is one segment", 50, 100, 5, 1},
|
|
{"single connection", 1000, 1, 1, 1},
|
|
{"min-split exact", 1000, 200, 5, 5},
|
|
}
|
|
for _, tc := range tests {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
segs := makeSegments(tc.total, tc.minSplit, tc.conns)
|
|
if len(segs) != tc.wantN {
|
|
t.Fatalf("got %d segments, want %d", len(segs), tc.wantN)
|
|
}
|
|
// Segments must be contiguous and cover the whole file exactly.
|
|
var next int64
|
|
for i, s := range segs {
|
|
if s.start != next {
|
|
t.Errorf("segment %d starts at %d, want %d", i, s.start, next)
|
|
}
|
|
if s.end < s.start {
|
|
t.Errorf("segment %d has end %d < start %d", i, s.end, s.start)
|
|
}
|
|
next = s.end + 1
|
|
}
|
|
if next != tc.total {
|
|
t.Errorf("segments cover %d bytes, want %d", next, tc.total)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestAutoConns(t *testing.T) {
|
|
const m = int64(20 << 20) // 20 MiB min-split, aria2's default
|
|
tests := []struct {
|
|
name string
|
|
total int64
|
|
minSplit int64
|
|
want int
|
|
}{
|
|
{"under one min-split is a single connection", 5 << 20, m, 1},
|
|
{"exactly one min-split", m, m, 1},
|
|
{"five min-splits", 100 << 20, m, 5},
|
|
{"caps at maxAutoConns", 10 << 30, m, maxAutoConns},
|
|
{"rounds down to whole pieces", m*3 + 1, m, 3},
|
|
{"tiny min-split still capped", 1 << 30, 1 << 20, maxAutoConns},
|
|
{"zero total is one connection", 0, m, 1},
|
|
}
|
|
for _, tc := range tests {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
if got := autoConns(tc.total, tc.minSplit); got != tc.want {
|
|
t.Errorf("autoConns(%d, %d) = %d, want %d", tc.total, tc.minSplit, got, tc.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestSegProgress(t *testing.T) {
|
|
s := seg{start: 100, end: 199} // length 100
|
|
if s.length() != 100 {
|
|
t.Fatalf("length = %d", s.length())
|
|
}
|
|
if s.done() {
|
|
t.Fatalf("new segment should not be done")
|
|
}
|
|
s.addWritten(60)
|
|
if s.offset() != 160 {
|
|
t.Errorf("offset = %d, want 160", s.offset())
|
|
}
|
|
if s.remaining() != 40 {
|
|
t.Errorf("remaining = %d, want 40", s.remaining())
|
|
}
|
|
s.addWritten(40)
|
|
if !s.done() {
|
|
t.Errorf("segment should be done after writing full length")
|
|
}
|
|
}
|
|
|
|
func TestSplitExt(t *testing.T) {
|
|
tests := []struct {
|
|
in string
|
|
stem, ext string
|
|
}{
|
|
{"foo.txt", "foo", ".txt"},
|
|
{"foo", "foo", ""},
|
|
{"/dir/foo.tar.gz", "/dir/foo.tar", ".gz"},
|
|
{"/dir/.bashrc", "/dir/.bashrc", ""}, // dotfile, no extension
|
|
{".bashrc", ".bashrc", ""},
|
|
{"archive.", "archive", "."},
|
|
}
|
|
for _, tc := range tests {
|
|
stem, ext := splitExt(tc.in)
|
|
if stem != tc.stem || ext != tc.ext {
|
|
t.Errorf("splitExt(%q) = %q,%q; want %q,%q", tc.in, stem, ext, tc.stem, tc.ext)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestUniqueName(t *testing.T) {
|
|
dir := t.TempDir()
|
|
|
|
// foo.txt taken -> foo.1.txt (.N inserted before the extension, item [13]).
|
|
taken := filepath.Join(dir, "foo.txt")
|
|
if err := os.WriteFile(taken, nil, 0o644); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
got, err := uniqueName(taken)
|
|
if err != nil {
|
|
t.Fatalf("uniqueName: %v", err)
|
|
}
|
|
if want := filepath.Join(dir, "foo.1.txt"); got != want {
|
|
t.Errorf("uniqueName(%q) = %q, want %q", taken, got, want)
|
|
}
|
|
|
|
// No-extension base -> base.1.
|
|
plain := filepath.Join(dir, "data")
|
|
if err := os.WriteFile(plain, nil, 0o644); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
got, err = uniqueName(plain)
|
|
if err != nil {
|
|
t.Fatalf("uniqueName: %v", err)
|
|
}
|
|
if want := filepath.Join(dir, "data.1"); got != want {
|
|
t.Errorf("uniqueName(%q) = %q, want %q", plain, got, want)
|
|
}
|
|
|
|
// foo.txt and foo.1.txt both taken -> foo.2.txt.
|
|
if err := os.WriteFile(filepath.Join(dir, "foo.1.txt"), nil, 0o644); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
got, err = uniqueName(taken)
|
|
if err != nil {
|
|
t.Fatalf("uniqueName: %v", err)
|
|
}
|
|
if want := filepath.Join(dir, "foo.2.txt"); got != want {
|
|
t.Errorf("uniqueName second pass = %q, want %q", got, want)
|
|
}
|
|
|
|
// A candidate that exists but has a .got control file is acceptable (it is a
|
|
// resumable partial).
|
|
ctrlBase := filepath.Join(dir, "r.bin")
|
|
cand1 := filepath.Join(dir, "r.1.bin")
|
|
if err := os.WriteFile(ctrlBase, nil, 0o644); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := os.WriteFile(cand1, nil, 0o644); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := os.WriteFile(controlPath(cand1), nil, 0o644); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
got, err = uniqueName(ctrlBase)
|
|
if err != nil {
|
|
t.Fatalf("uniqueName: %v", err)
|
|
}
|
|
if got != cand1 {
|
|
t.Errorf("uniqueName with control sidecar = %q, want %q", got, cand1)
|
|
}
|
|
}
|
|
|
|
func TestNameFromContentDisposition(t *testing.T) {
|
|
tests := []struct {
|
|
cd, want string
|
|
}{
|
|
{`attachment; filename="report.pdf"`, "report.pdf"},
|
|
{`attachment; filename=report.pdf`, "report.pdf"},
|
|
{`inline; filename*=UTF-8''na%C3%AFve.txt`, "naïve.txt"},
|
|
{`attachment; filename="../../etc/passwd"`, "passwd"}, // path stripped
|
|
{`attachment`, ""},
|
|
{`garbage`, ""},
|
|
}
|
|
for _, tc := range tests {
|
|
if got := nameFromContentDisposition(tc.cd); got != tc.want {
|
|
t.Errorf("nameFromContentDisposition(%q) = %q, want %q", tc.cd, got, tc.want)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestSeedSegmentsFromPrefix(t *testing.T) {
|
|
segs := makeSegments(1000, 100, 5) // five 200-byte segments
|
|
seedSegmentsFromPrefix(segs, 450) // 2 full segments + 50 bytes of the third
|
|
wants := []int64{200, 200, 50, 0, 0}
|
|
for i := range segs {
|
|
if segs[i].written != wants[i] {
|
|
t.Errorf("segment %d written = %d, want %d", i, segs[i].written, wants[i])
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestParseContentRange(t *testing.T) {
|
|
start, end, total, ok := parseContentRange("bytes 0-0/12345")
|
|
if !ok || start != 0 || end != 0 || total != 12345 {
|
|
t.Errorf("parseContentRange = %d,%d,%d,%v", start, end, total, ok)
|
|
}
|
|
if _, _, _, ok := parseContentRange("garbage"); ok {
|
|
t.Errorf("parseContentRange(garbage) ok = true, want false")
|
|
}
|
|
}
|
|
|
|
// assertTiles checks that the pool's segments still cover [0,total) exactly:
|
|
// sorted by start they must be contiguous with no gap and no overlap.
|
|
func assertTiles(t *testing.T, p *pool, total int64) {
|
|
t.Helper()
|
|
p.mu.Lock()
|
|
type rng struct{ start, end int64 }
|
|
rs := make([]rng, len(p.segs))
|
|
for i, s := range p.segs {
|
|
rs[i] = rng{s.start, s.endOff()}
|
|
}
|
|
p.mu.Unlock()
|
|
sort.Slice(rs, func(i, j int) bool { return rs[i].start < rs[j].start })
|
|
var next int64
|
|
for _, r := range rs {
|
|
if r.start != next {
|
|
t.Fatalf("segment gap/overlap: next byte %d, got start %d (ranges %v)", next, r.start, rs)
|
|
}
|
|
next = r.end + 1
|
|
}
|
|
if next != total {
|
|
t.Fatalf("segments cover %d bytes, want %d", next, total)
|
|
}
|
|
}
|
|
|
|
func TestPoolAcquireThenSteal(t *testing.T) {
|
|
// Sizes are in units of readBuf because that is what newPool floors minSplit
|
|
// to and what the no-overlap invariant is stated against.
|
|
const total = 8 * readBuf
|
|
p := newPool(makeSegments(total, readBuf, 1), readBuf) // one segment [0,total)
|
|
|
|
first := p.acquire()
|
|
if first == nil || first.start != 0 || first.endOff() != total-1 {
|
|
t.Fatalf("first acquire = %+v, want the whole [0,%d] segment", first, total-1)
|
|
}
|
|
if p.acquire() == nil {
|
|
// nothing fresh left, so this must steal first's back half
|
|
t.Fatal("second acquire returned nil, want a stolen tail")
|
|
}
|
|
// first kept the front half, a tail took the back half; together they still tile.
|
|
if first.endOff() != total/2-1 {
|
|
t.Errorf("victim end = %d, want %d after midpoint split", first.endOff(), total/2-1)
|
|
}
|
|
assertTiles(t, p, total)
|
|
}
|
|
|
|
func TestPoolNoStealBelowThreshold(t *testing.T) {
|
|
// remaining is below 2*minSplit, so the lone segment is not worth splitting
|
|
// and an idle worker is told there is nothing to do.
|
|
p := newPool(makeSegments(readBuf+readBuf/2, readBuf, 1), readBuf)
|
|
if s := p.acquire(); s == nil {
|
|
t.Fatal("first acquire returned nil, want the only segment")
|
|
}
|
|
if s := p.acquire(); s != nil {
|
|
t.Fatalf("second acquire = %+v, want nil (tail would be below min-split)", s)
|
|
}
|
|
}
|
|
|
|
// TestPoolConcurrentCoverage drives the pool the way real workers do — acquire a
|
|
// segment, copy it in small chunks, commit each with p.advance, repeat — across
|
|
// more workers than initial segments so stealing is forced. The coverage array
|
|
// proves the core invariant: every byte is written exactly once, so a steal
|
|
// never overlaps the owner's writes and never leaves a gap.
|
|
func TestPoolConcurrentCoverage(t *testing.T) {
|
|
const (
|
|
total = int64(64 * readBuf) // 2 MiB
|
|
minSplit = int64(readBuf) // steal threshold is 2*this
|
|
chunk = int64(readBuf / 4) // one "read", well below minSplit
|
|
workers = 8
|
|
)
|
|
p := newPool(makeSegments(total, minSplit, 1), minSplit) // start with a single segment
|
|
cover := make([]uint32, total)
|
|
|
|
var wg sync.WaitGroup
|
|
for w := 0; w < workers; w++ {
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
for {
|
|
s := p.acquire()
|
|
if s == nil {
|
|
return
|
|
}
|
|
for s.remaining() > 0 {
|
|
n := chunk
|
|
if s.remaining() < n {
|
|
n = s.remaining()
|
|
}
|
|
off := s.offset()
|
|
for j := off; j < off+n; j++ {
|
|
atomic.AddUint32(&cover[j], 1)
|
|
}
|
|
p.advance(s, n)
|
|
}
|
|
}
|
|
}()
|
|
}
|
|
wg.Wait()
|
|
|
|
for i, c := range cover {
|
|
if c != 1 {
|
|
t.Fatalf("byte %d written %d times, want exactly 1", i, c)
|
|
}
|
|
}
|
|
assertTiles(t, p, total)
|
|
if len(p.segs) == 1 {
|
|
t.Error("no stealing happened: still one segment after 8 workers drained it")
|
|
}
|
|
}
|
|
|
|
// TestPoolSnapshotAfterStealRoundTrips checks that a steal which happens before a
|
|
// crash survives the control file: the snapshot records the stolen tail, and a
|
|
// resume rebuilds a segment set that still tiles the file and keeps the bytes
|
|
// already written. This is the persistence path an interrupt-and-resume relies on
|
|
// but cannot deterministically trigger from the outside.
|
|
func TestPoolSnapshotAfterStealRoundTrips(t *testing.T) {
|
|
const total = 8 * readBuf
|
|
p := newPool(makeSegments(total, readBuf, 1), readBuf)
|
|
a := p.acquire() // the whole file
|
|
p.advance(a, readBuf) // owner makes some progress, then idles out
|
|
b := p.acquire() // steals a's back half
|
|
if b == nil {
|
|
t.Fatal("expected a stolen tail")
|
|
}
|
|
p.advance(b, readBuf) // the tail's worker makes progress too
|
|
|
|
c := p.snapshot("http://x", total, "", "")
|
|
if len(c.Segs) != 2 {
|
|
t.Fatalf("snapshot has %d segments, want 2 (original + stolen tail)", len(c.Segs))
|
|
}
|
|
|
|
rebuilt := segsFromControl(&c)
|
|
sort.Slice(rebuilt, func(i, j int) bool { return rebuilt[i].start < rebuilt[j].start })
|
|
var next, written int64
|
|
for _, s := range rebuilt {
|
|
if s.start != next {
|
|
t.Fatalf("rebuilt gap/overlap: next byte %d, got start %d", next, s.start)
|
|
}
|
|
next = s.endOff() + 1
|
|
written += s.progress()
|
|
}
|
|
if next != total {
|
|
t.Fatalf("rebuilt covers %d bytes, want %d", next, total)
|
|
}
|
|
if written != 2*readBuf {
|
|
t.Errorf("rebuilt written = %d, want %d (bytes must survive the round trip)", written, 2*readBuf)
|
|
}
|
|
}
|