129 lines
2.8 KiB
Go
129 lines
2.8 KiB
Go
package fetch
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"os"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"ddc/internal/diskspace"
|
|
)
|
|
|
|
const (
|
|
checkpointInterval = 10 * time.Second
|
|
spaceCheckInterval = int64(256) << 20
|
|
)
|
|
|
|
type session struct {
|
|
downloader *Downloader
|
|
source *remoteFile
|
|
file *os.File
|
|
request Request
|
|
state State
|
|
segments []*liveSegment
|
|
sinceSpaceCheck atomic.Int64
|
|
}
|
|
|
|
func newSession(downloader *Downloader, source *remoteFile, file *os.File, request Request, state State) *session {
|
|
segments := make([]*liveSegment, len(state.Segments))
|
|
for index, segment := range state.Segments {
|
|
segments[index] = newLiveSegment(segment)
|
|
}
|
|
return &session{
|
|
downloader: downloader,
|
|
source: source,
|
|
file: file,
|
|
request: request,
|
|
state: state,
|
|
segments: segments,
|
|
}
|
|
}
|
|
|
|
func (session *session) run(ctx context.Context, probe *http.Response) error {
|
|
session.saveCheckpoint()
|
|
stopSaving := make(chan struct{})
|
|
savingStopped := make(chan struct{})
|
|
go session.saveCheckpoints(stopSaving, savingStopped)
|
|
|
|
err := session.runWorkers(ctx, probe)
|
|
close(stopSaving)
|
|
<-savingStopped
|
|
session.saveCheckpoint()
|
|
return err
|
|
}
|
|
|
|
func (session *session) runWorkers(ctx context.Context, probe *http.Response) error {
|
|
ctx, cancel := context.WithCancelCause(ctx)
|
|
defer cancel(nil)
|
|
|
|
var workers sync.WaitGroup
|
|
var firstFailure sync.Once
|
|
var failure error
|
|
for index, segment := range session.segments {
|
|
if segment.complete() {
|
|
continue
|
|
}
|
|
var initial *http.Response
|
|
if index == 0 {
|
|
initial, probe = probe, nil
|
|
}
|
|
worker := &segmentWorker{session: session, segment: segment}
|
|
workers.Add(1)
|
|
go func() {
|
|
defer workers.Done()
|
|
if err := worker.run(ctx, initial); err != nil {
|
|
firstFailure.Do(func() {
|
|
failure = err
|
|
cancel(err)
|
|
})
|
|
}
|
|
}()
|
|
}
|
|
closeBody(probe)
|
|
workers.Wait()
|
|
return failure
|
|
}
|
|
|
|
func (session *session) record(count int64) error {
|
|
session.request.Progress.Add(count)
|
|
if session.sinceSpaceCheck.Add(count) < spaceCheckInterval {
|
|
return nil
|
|
}
|
|
session.sinceSpaceCheck.Store(0)
|
|
return diskspace.EnsureFree(session.request.Directory, 0, session.request.ReservedBytes)
|
|
}
|
|
|
|
func (session *session) snapshot() State {
|
|
state := session.state
|
|
state.Segments = make([]Segment, len(session.segments))
|
|
for index, segment := range session.segments {
|
|
state.Segments[index] = segment.snapshot()
|
|
}
|
|
return state
|
|
}
|
|
|
|
func (session *session) saveCheckpoints(stop <-chan struct{}, stopped chan<- struct{}) {
|
|
defer close(stopped)
|
|
ticker := time.NewTicker(checkpointInterval)
|
|
defer ticker.Stop()
|
|
for {
|
|
select {
|
|
case <-stop:
|
|
return
|
|
case <-ticker.C:
|
|
session.saveCheckpoint()
|
|
}
|
|
}
|
|
}
|
|
|
|
func (session *session) saveCheckpoint() {
|
|
if session.request.Checkpoint == nil {
|
|
return
|
|
}
|
|
state := session.snapshot()
|
|
session.file.Sync()
|
|
session.request.Checkpoint(state)
|
|
}
|