httpdl: add segment work-stealing and --auto-split
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 (opt-in): derive the connection count from file size (one per min-split-size, up to 16), overriding -x/-s. default stays at 1 connection for a single server.
This commit is contained in:
@@ -27,11 +27,13 @@ type segState struct {
|
||||
|
||||
func controlPath(out string) string { return out + ".got" }
|
||||
|
||||
// snapshot builds a control record from the live segments.
|
||||
func snapshot(url string, total int64, etag, lastmod string, segs []seg) control {
|
||||
// snapshot builds a control record from the live segments. Callers hold the
|
||||
// pool lock (see pool.snapshot) so the slice and each segment's frontier are
|
||||
// stable for the read.
|
||||
func snapshot(url string, total int64, etag, lastmod string, segs []*seg) control {
|
||||
c := control{URL: url, Total: total, ETag: etag, LastModified: lastmod, Segs: make([]segState, len(segs))}
|
||||
for i := range segs {
|
||||
c.Segs[i] = segState{segs[i].start, segs[i].end, atomic.LoadInt64(&segs[i].written)}
|
||||
c.Segs[i] = segState{segs[i].start, segs[i].endOff(), atomic.LoadInt64(&segs[i].written)}
|
||||
}
|
||||
return c
|
||||
}
|
||||
|
||||
@@ -57,7 +57,8 @@ type Config struct {
|
||||
Split int // --split: total connections across all mirrors
|
||||
MaxConnPerServer int // --max-connection-per-server: per-host connection cap
|
||||
MinSplit int64
|
||||
Tries int // 0 = unlimited
|
||||
AutoSplit bool // --auto-split (got-only): pick the connection count from file size
|
||||
Tries int // 0 = unlimited
|
||||
Timeout time.Duration
|
||||
FileAlloc string // none | prealloc | trunc | falloc
|
||||
Continue bool
|
||||
@@ -119,7 +120,7 @@ type Config struct {
|
||||
// connections.
|
||||
type Download struct {
|
||||
uris []string // mirror list; uris[0] is the primary (naming + resume key)
|
||||
maxConns int // segment workers = min(split, len(uris)*max-connection-per-server)
|
||||
maxConns int // manual segment-worker cap = min(split, len(uris)*max-connection-per-server); --auto-split derives its own count per file in segmented()
|
||||
cfg Config
|
||||
client *http.Client
|
||||
limit []*rate.Limiter
|
||||
@@ -214,10 +215,17 @@ func New(uris []string, cfg Config) *Download {
|
||||
if perHost < 1 {
|
||||
perHost = 1
|
||||
}
|
||||
// The transport must allow as many concurrent connections per host as we may
|
||||
// actually open. --auto-split scales up to maxAutoConns per host once the file
|
||||
// size is known (in segmented), so raise the cap to that ceiling when it is on.
|
||||
connCap := perHost
|
||||
if cfg.AutoSplit {
|
||||
connCap = maxAutoConns
|
||||
}
|
||||
tr := &http.Transport{
|
||||
Proxy: http.ProxyFromEnvironment,
|
||||
MaxConnsPerHost: perHost,
|
||||
MaxIdleConnsPerHost: perHost,
|
||||
MaxConnsPerHost: connCap,
|
||||
MaxIdleConnsPerHost: connCap,
|
||||
ResponseHeaderTimeout: cfg.Timeout,
|
||||
TLSHandshakeTimeout: connectTimeout,
|
||||
DialContext: dialContext,
|
||||
@@ -610,7 +618,13 @@ func (d *Download) probe(ctx context.Context, uri string) (probeResult, error) {
|
||||
// segmented downloads total bytes over d.maxConns workers, with per-segment
|
||||
// resume from the control file.
|
||||
func (d *Download) segmented(ctx context.Context, out string, total int64, etag, lastmod string) error {
|
||||
segs := makeSegments(total, d.cfg.MinSplit, d.maxConns)
|
||||
// --auto-split picks the connection count from the file size now that total is
|
||||
// known; otherwise use the manual cap fixed at construction.
|
||||
conns := d.maxConns
|
||||
if d.cfg.AutoSplit {
|
||||
conns = autoConns(total, d.cfg.MinSplit)
|
||||
}
|
||||
segs := makeSegments(total, d.cfg.MinSplit, conns)
|
||||
if c := d.resumeControl(out, etag, lastmod); c != nil {
|
||||
segs = segsFromControl(c)
|
||||
} else if d.cfg.Continue && !fileExists(controlPath(out)) {
|
||||
@@ -648,36 +662,31 @@ func (d *Download) segmented(ctx context.Context, out string, total int64, etag,
|
||||
ctx, cancel := context.WithCancel(ctx)
|
||||
defer cancel()
|
||||
|
||||
p := newPool(segs, d.cfg.MinSplit)
|
||||
|
||||
// Periodically persist progress so a crash can resume.
|
||||
saveDone := make(chan struct{})
|
||||
go d.saveLoop(ctx, out, total, etag, lastmod, segs, saveDone)
|
||||
|
||||
jobs := make(chan *seg)
|
||||
go func() {
|
||||
defer close(jobs)
|
||||
for i := range segs {
|
||||
if segs[i].done() {
|
||||
continue
|
||||
}
|
||||
select {
|
||||
case jobs <- &segs[i]:
|
||||
case <-ctx.Done(): // stop feeding once we're tearing down
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
go d.saveLoop(ctx, out, total, etag, lastmod, p, saveDone)
|
||||
|
||||
// One worker per connection. A worker fetches segments until acquire runs
|
||||
// dry, which happens only once every remaining byte is owned; a worker that
|
||||
// finishes early steals the tail of a slower segment rather than sitting idle
|
||||
// while the last segment drains over a single connection.
|
||||
var (
|
||||
wg sync.WaitGroup
|
||||
errOnce sync.Once
|
||||
runErr error
|
||||
)
|
||||
for w := 0; w < d.maxConns; w++ {
|
||||
for w := 0; w < conns; w++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for s := range jobs {
|
||||
if err := d.fetchSeg(ctx, f, s, total); err != nil {
|
||||
for {
|
||||
s := p.acquire()
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
if err := d.fetchSeg(ctx, f, p, s, total); err != nil {
|
||||
errOnce.Do(func() { runErr = err; cancel() })
|
||||
return
|
||||
}
|
||||
@@ -691,7 +700,7 @@ func (d *Download) segmented(ctx context.Context, out string, total int64, etag,
|
||||
if runErr != nil {
|
||||
return runErr // keep the control file for a later -c
|
||||
}
|
||||
snapshot(d.primary(), total, etag, lastmod, segs).save(out)
|
||||
p.snapshot(d.primary(), total, etag, lastmod).save(out)
|
||||
removeControl(out)
|
||||
return nil
|
||||
}
|
||||
@@ -699,7 +708,7 @@ func (d *Download) segmented(ctx context.Context, out string, total int64, etag,
|
||||
// fetchSeg downloads one segment, retrying from its resume point on error.
|
||||
// Each attempt re-reads s.offset(), so a retry continues from the bytes already
|
||||
// written rather than restarting the range.
|
||||
func (d *Download) fetchSeg(ctx context.Context, f *os.File, s *seg, total int64) error {
|
||||
func (d *Download) fetchSeg(ctx context.Context, f *os.File, p *pool, s *seg, total int64) error {
|
||||
// Try each mirror in turn for this segment: a transient error retries the
|
||||
// same mirror under withRetries; an exhausted budget or a per-mirror permanent
|
||||
// error (a 404, a non-206, a length mismatch) falls over to the next mirror.
|
||||
@@ -716,7 +725,7 @@ func (d *Download) fetchSeg(ctx context.Context, f *os.File, s *seg, total int64
|
||||
if s.done() {
|
||||
return nil
|
||||
}
|
||||
return d.fetchOnce(ctx, f, s, uri, total)
|
||||
return d.fetchOnce(ctx, f, p, s, uri, total)
|
||||
})
|
||||
if err == nil {
|
||||
return nil
|
||||
@@ -726,10 +735,10 @@ func (d *Download) fetchSeg(ctx context.Context, f *os.File, s *seg, total int64
|
||||
return lastErr
|
||||
}
|
||||
|
||||
func (d *Download) fetchOnce(ctx context.Context, f *os.File, s *seg, uri string, total int64) error {
|
||||
func (d *Download) fetchOnce(ctx context.Context, f *os.File, p *pool, s *seg, uri string, total int64) error {
|
||||
reqCtx, cancel := context.WithCancel(ctx)
|
||||
defer cancel()
|
||||
req, err := d.request(reqCtx, uri, fmt.Sprintf("bytes=%d-%d", s.offset(), s.end))
|
||||
req, err := d.request(reqCtx, uri, fmt.Sprintf("bytes=%d-%d", s.offset(), s.endOff()))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
@@ -755,7 +764,7 @@ func (d *Download) fetchOnce(ctx context.Context, f *os.File, s *seg, uri string
|
||||
}
|
||||
body, stop := d.idleGuard(resp.Body, cancel)
|
||||
defer stop()
|
||||
return d.pump(ctx, f, s, body)
|
||||
return d.pump(ctx, f, p, s, body)
|
||||
}
|
||||
|
||||
// errIdleTimeout marks a read that stalled past the idle window, so callers can
|
||||
@@ -901,8 +910,10 @@ func statusError(ctxMsg string, code int) error {
|
||||
}
|
||||
|
||||
// pump copies the response body into the file at the segment's running offset,
|
||||
// stopping at the segment end and respecting rate limits.
|
||||
func (d *Download) pump(ctx context.Context, f *os.File, s *seg, body io.Reader) error {
|
||||
// stopping at the segment end and respecting rate limits. The loop reads
|
||||
// remaining() afresh each turn, so a steal that shrinks s mid-transfer simply
|
||||
// ends the loop early at the new end, leaving the stolen tail to its new worker.
|
||||
func (d *Download) pump(ctx context.Context, f *os.File, p *pool, s *seg, body io.Reader) error {
|
||||
buf := make([]byte, readBuf)
|
||||
for s.remaining() > 0 {
|
||||
n := int64(len(buf))
|
||||
@@ -914,7 +925,7 @@ func (d *Download) pump(ctx context.Context, f *os.File, s *seg, body io.Reader)
|
||||
if _, werr := f.WriteAt(buf[:rd], s.offset()); werr != nil {
|
||||
return werr
|
||||
}
|
||||
s.advance(int64(rd))
|
||||
p.advance(s, int64(rd))
|
||||
atomic.AddInt64(&d.completed, int64(rd))
|
||||
d.throttle(ctx, rd)
|
||||
}
|
||||
@@ -1066,7 +1077,7 @@ func (d *Download) singleOnce(ctx context.Context, out, uri string, resumeFromDi
|
||||
}
|
||||
}
|
||||
|
||||
func (d *Download) saveLoop(ctx context.Context, out string, total int64, etag, lastmod string, segs []seg, done chan<- struct{}) {
|
||||
func (d *Download) saveLoop(ctx context.Context, out string, total int64, etag, lastmod string, p *pool, done chan<- struct{}) {
|
||||
defer close(done)
|
||||
interval := d.cfg.AutoSaveInterval
|
||||
if interval <= 0 {
|
||||
@@ -1077,10 +1088,10 @@ func (d *Download) saveLoop(ctx context.Context, out string, total int64, etag,
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
snapshot(d.primary(), total, etag, lastmod, segs).save(out)
|
||||
p.snapshot(d.primary(), total, etag, lastmod).save(out)
|
||||
return
|
||||
case <-t.C:
|
||||
snapshot(d.primary(), total, etag, lastmod, segs).save(out)
|
||||
p.snapshot(d.primary(), total, etag, lastmod).save(out)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,25 +1,32 @@
|
||||
package httpdl
|
||||
|
||||
import "sync/atomic"
|
||||
import (
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
)
|
||||
|
||||
// seg is one contiguous byte range of the output file, downloaded by a single
|
||||
// ranged GET. A segment IS a byte range — a plain HTTP downloader does not need
|
||||
// the Piece/Segment/block layering that exists only to share code with
|
||||
// BitTorrent. written is updated with atomic ops so Stat() can read it while a
|
||||
// worker advances it.
|
||||
// BitTorrent. written and end are updated with atomic ops so Stat() and the
|
||||
// resume snapshot can read them while a worker advances written or a steal
|
||||
// shrinks end (see pool).
|
||||
type seg struct {
|
||||
index int
|
||||
start int64 // first byte offset, inclusive
|
||||
end int64 // last byte offset, inclusive
|
||||
end int64 // last byte offset, inclusive (a steal may shrink this)
|
||||
written int64 // bytes already written into this segment (the resume point)
|
||||
owned bool // a worker is fetching this segment; guarded by pool.mu
|
||||
}
|
||||
|
||||
func (s *seg) length() int64 { return s.end - s.start + 1 }
|
||||
func (s *seg) done() bool { return atomic.LoadInt64(&s.written) >= s.length() }
|
||||
func (s *seg) advance(n int64) { atomic.AddInt64(&s.written, n) }
|
||||
func (s *seg) progress() int64 { return atomic.LoadInt64(&s.written) }
|
||||
func (s *seg) offset() int64 { return s.start + atomic.LoadInt64(&s.written) }
|
||||
func (s *seg) remaining() int64 { return s.length() - atomic.LoadInt64(&s.written) }
|
||||
func (s *seg) endOff() int64 { return atomic.LoadInt64(&s.end) }
|
||||
func (s *seg) setEnd(v int64) { atomic.StoreInt64(&s.end, v) }
|
||||
func (s *seg) length() int64 { return s.endOff() - s.start + 1 }
|
||||
func (s *seg) done() bool { return atomic.LoadInt64(&s.written) >= s.length() }
|
||||
func (s *seg) addWritten(n int64) { atomic.AddInt64(&s.written, n) }
|
||||
func (s *seg) progress() int64 { return atomic.LoadInt64(&s.written) }
|
||||
func (s *seg) offset() int64 { return s.start + atomic.LoadInt64(&s.written) }
|
||||
func (s *seg) remaining() int64 { return s.length() - atomic.LoadInt64(&s.written) }
|
||||
|
||||
// makeSegments divides a file of total bytes into contiguous segments, using at
|
||||
// most conns of them and never splitting below minSplit. The remainder lands in
|
||||
@@ -51,3 +58,127 @@ func makeSegments(total, minSplit int64, conns int) []seg {
|
||||
}
|
||||
return segs
|
||||
}
|
||||
|
||||
// maxAutoConns is the ceiling --auto-split scales to. aria2 caps
|
||||
// max-connection-per-server (-x) at 16, so a got-chosen count honours the same
|
||||
// per-server limit rather than inventing a looser one.
|
||||
const maxAutoConns = 16
|
||||
|
||||
// autoConns picks a connection count from the file size, used only when
|
||||
// --auto-split is set: one connection per min-split-size of content, at least 1
|
||||
// and at most maxAutoConns. It is the got-only stand-in for hand-tuning -x/-s,
|
||||
// and because the count never exceeds total/minSplit, makeSegments yields exactly
|
||||
// that many balanced segments.
|
||||
func autoConns(total, minSplit int64) int {
|
||||
if minSplit < 1 {
|
||||
minSplit = 1
|
||||
}
|
||||
n := total / minSplit
|
||||
if n < 1 {
|
||||
n = 1
|
||||
}
|
||||
if n > maxAutoConns {
|
||||
n = maxAutoConns
|
||||
}
|
||||
return int(n)
|
||||
}
|
||||
|
||||
// pool is the live set of segments for one segmented download. A worker takes a
|
||||
// segment with acquire and fetches it to completion; once no fresh segment is
|
||||
// left, an idle worker steals the back half of whichever in-flight segment has
|
||||
// the most still to download, so the last slow segment is shared across the idle
|
||||
// connections instead of draining alone. aria2 does the same on-demand split via
|
||||
// its SegmentMan; without it a fixed pre-division leaves connections idle while
|
||||
// one slow mirror finishes.
|
||||
//
|
||||
// The mutex serialises a steal against the byte accounting it splits: advance (a
|
||||
// worker committing a write) takes it too, so a steal reading a victim's frontier
|
||||
// can never race the owner advancing past the chosen split point. A steal fires
|
||||
// only when the victim still has at least 2*minSplit to go and cuts at the
|
||||
// midpoint, so both halves stay >= minSplit (--min-split-size) and the half the
|
||||
// owner keeps is far larger than one read buffer — the owner's in-flight write
|
||||
// therefore can never reach into the stolen tail.
|
||||
type pool struct {
|
||||
mu sync.Mutex
|
||||
segs []*seg
|
||||
minSplit int64
|
||||
}
|
||||
|
||||
// newPool wraps the pre-divided segments in a pool. Each is copied into its own
|
||||
// allocation so a later steal can append a tail without invalidating the pointers
|
||||
// workers already hold.
|
||||
//
|
||||
// minSplit is floored at readBuf: a steal leaves the owner the front half of the
|
||||
// split, which is at least minSplit, while the owner's in-flight write is at most
|
||||
// readBuf, so minSplit >= readBuf is exactly what keeps that write out of the
|
||||
// stolen tail. The CLI already holds --min-split-size well above readBuf (>= 1
|
||||
// MiB), so the floor only guards direct callers and tests — but it keeps the
|
||||
// no-overlap invariant inside this file rather than resting on a distant flag.
|
||||
func newPool(segs []seg, minSplit int64) *pool {
|
||||
if minSplit < readBuf {
|
||||
minSplit = readBuf
|
||||
}
|
||||
p := &pool{minSplit: minSplit, segs: make([]*seg, len(segs))}
|
||||
for i := range segs {
|
||||
s := segs[i]
|
||||
p.segs[i] = &s
|
||||
}
|
||||
return p
|
||||
}
|
||||
|
||||
// acquire returns the next segment for a worker to fetch: a fresh one if any
|
||||
// remain, otherwise the tail split off the busiest in-flight segment. It returns
|
||||
// nil when every remaining byte is already owned, i.e. the download is finishing
|
||||
// and there is nothing left to steal.
|
||||
func (p *pool) acquire() *seg {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
for _, s := range p.segs {
|
||||
if !s.owned && !s.done() {
|
||||
s.owned = true
|
||||
return s
|
||||
}
|
||||
}
|
||||
return p.steal()
|
||||
}
|
||||
|
||||
// steal splits the back half off the in-flight segment with the most remaining
|
||||
// and returns it, or nil if none has enough left to be worth splitting. The
|
||||
// caller holds p.mu, which keeps every owner's advance out so each victim's
|
||||
// frontier is stable while we choose and commit the split.
|
||||
func (p *pool) steal() *seg {
|
||||
var victim *seg
|
||||
for _, s := range p.segs {
|
||||
if s.owned && !s.done() && s.remaining() >= 2*p.minSplit {
|
||||
if victim == nil || s.remaining() > victim.remaining() {
|
||||
victim = s
|
||||
}
|
||||
}
|
||||
}
|
||||
if victim == nil {
|
||||
return nil
|
||||
}
|
||||
mid := victim.offset() + victim.remaining()/2
|
||||
// index only feeds the mirror round-robin (d.mirror) and the "segment N" log
|
||||
// label; a tail's index need not be contiguous or unique, so len is fine.
|
||||
tail := &seg{index: len(p.segs), start: mid, end: victim.endOff(), owned: true}
|
||||
victim.setEnd(mid - 1)
|
||||
p.segs = append(p.segs, tail)
|
||||
return tail
|
||||
}
|
||||
|
||||
// advance commits n freshly written bytes to s under the pool lock, so a
|
||||
// concurrent steal sees a stable frontier for the segment it may split.
|
||||
func (p *pool) advance(s *seg, n int64) {
|
||||
p.mu.Lock()
|
||||
s.addWritten(n)
|
||||
p.mu.Unlock()
|
||||
}
|
||||
|
||||
// snapshot builds the resume record from the live set, including any stolen
|
||||
// tails, under the lock so it never races a steal appending a segment.
|
||||
func (p *pool) snapshot(url string, total int64, etag, lastmod string) control {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
return snapshot(url, total, etag, lastmod, p.segs)
|
||||
}
|
||||
|
||||
@@ -3,6 +3,9 @@ package httpdl
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
)
|
||||
|
||||
@@ -44,6 +47,31 @@ func TestMakeSegments(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
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 {
|
||||
@@ -52,14 +80,14 @@ func TestSegProgress(t *testing.T) {
|
||||
if s.done() {
|
||||
t.Fatalf("new segment should not be done")
|
||||
}
|
||||
s.advance(60)
|
||||
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.advance(40)
|
||||
s.addWritten(40)
|
||||
if !s.done() {
|
||||
t.Errorf("segment should be done after writing full length")
|
||||
}
|
||||
@@ -186,3 +214,151 @@ func TestParseContentRange(t *testing.T) {
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user