got/httpdl/segment.go

package httpdl

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 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 (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) 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
// the last segment. With conns==1 (the default) this yields one segment.
func makeSegments(total, minSplit int64, conns int) []seg {
	if conns < 1 {
		conns = 1
	}
	if minSplit < 1 {
		minSplit = 1
	}
	n := int64(conns)
	if max := total / minSplit; max < n {
		n = max
	}
	if n < 1 {
		n = 1
	}
	chunk := total / n
	segs := make([]seg, n)
	var start int64
	for i := int64(0); i < n; i++ {
		end := start + chunk - 1
		if i == n-1 {
			end = total - 1
		}
		segs[i] = seg{index: int(i), start: start, end: end}
		start = end + 1
	}
	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)
}