soft_quay/core/downloader/transfer.go

package downloader

import (
	"context"
	"fmt"
	"os"
	"time"
)

func (queue *Queue) transfer(
	runtime *runtimeTask,
	ctx context.Context,
	attempt uint64,
) (completedPath string, currentAttempt uint64, transferErr error) {
	currentAttempt = attempt
	task, valid := queue.taskForAttempt(runtime, currentAttempt)
	if !valid {
		return "", currentAttempt, context.Canceled
	}
	paths, err := ensureTransferLayout(queue.root, task.RequestID)
	if err != nil {
		return "", currentAttempt, err
	}
	offset, _, err := regularFileSize(paths.Part)
	if err != nil {
		return "", currentAttempt, err
	}

	response, err := queue.transport.Open(ctx, OpenRequest{
		URL:       task.URL,
		Offset:    offset,
		Validator: task.Validator,
	})
	if err != nil {
		return "", currentAttempt, err
	}
	body := response.Body
	defer func() {
		if body != nil {
			_ = body.Close()
		}
	}()

	if task.TotalKnown && response.TotalKnown && response.Total != task.Total {
		return "", currentAttempt, fmt.Errorf(
			"%w: response total %d, expected %d",
			ErrTransferIncomplete,
			response.Total,
			task.Total,
		)
	}
	if !task.TotalKnown && response.TotalKnown && response.Total > queue.maxUnknown {
		return "", currentAttempt, ErrTransferTooLarge
	}
	if !task.TotalKnown && offset > 0 && !response.Restart && !response.TotalKnown {
		return "", currentAttempt, fmt.Errorf(
			"%w: resumed response has unknown complete length",
			ErrTransferIncomplete,
		)
	}

	var part *os.File
	if response.Restart {
		part, err = openPart(paths, offset, true)
		if err != nil {
			return "", currentAttempt, err
		}
		if err := part.Sync(); err != nil {
			_ = part.Close()
			return "", currentAttempt, fmt.Errorf("sync restarted part: %w", err)
		}
		currentAttempt, task, err = queue.restartAttempt(
			runtime,
			currentAttempt,
			response,
		)
		if err != nil {
			_ = part.Close()
			return "", currentAttempt, err
		}
		offset = 0
	} else {
		task, err = queue.applyResponseFacts(runtime, currentAttempt, response)
		if err != nil {
			return "", currentAttempt, err
		}
	}

	var maximum int64
	var expectedKnown bool
	var expected int64
	if task.TotalKnown {
		if offset > task.Total {
			return "", currentAttempt, fmt.Errorf(
				"%w: offset exceeds total",
				ErrTaskCorrupt,
			)
		}
		maximum = task.Total - offset
		expectedKnown = true
		expected = maximum
	} else {
		if offset >= queue.maxUnknown {
			return "", currentAttempt, ErrTransferTooLarge
		}
		maximum = queue.maxUnknown - offset
		if response.ResponseLengthKnown {
			expectedKnown = true
			expected = response.ResponseLength
		}
	}
	if !task.TotalKnown && response.TotalKnown {
		if response.Total > queue.maxUnknown ||
			offset > response.Total ||
			!response.ResponseLengthKnown ||
			response.ResponseLength != response.Total-offset {
			if part != nil {
				_ = part.Close()
			}
			return "", currentAttempt, ErrTransferIncomplete
		}
	}
	if response.ResponseLengthKnown && response.ResponseLength > maximum {
		if part != nil {
			_ = part.Close()
		}
		return "", currentAttempt, ErrTransferTooLarge
	}
	if expectedKnown &&
		response.ResponseLengthKnown &&
		response.ResponseLength != expected {
		if part != nil {
			_ = part.Close()
		}
		return "", currentAttempt, ErrTransferIncomplete
	}

	if part == nil {
		part, err = openPart(paths, offset, false)
		if err != nil {
			return "", currentAttempt, err
		}
	}
	lastEventAt := queue.clock.Now()
	lastEventDone := offset
	lastPublishedDone := offset
	copyOptions := CopyOptions{
		MaxBytes:           maximum,
		ExpectedBytesKnown: expectedKnown,
		ExpectedBytes:      expected,
		Generation:         currentAttempt,
		Progress: func(progress CopyProgress) error {
			absoluteDone := offset + progress.Written
			now := queue.clock.Now()
			emit := queue.progressEvery == 0 ||
				now.Sub(lastEventAt) >= queue.progressEvery
			if err := queue.updateAttemptDone(
				runtime,
				currentAttempt,
				absoluteDone,
			); err != nil {
				return err
			}
			if !emit {
				return nil
			}
			speed := progressSpeed(lastEventAt, now, lastEventDone, absoluteDone)
			if err := queue.persistAndPublishProgress(
				runtime,
				currentAttempt,
				speed,
			); err != nil {
				return err
			}
			lastEventAt = now
			lastEventDone = absoluteDone
			lastPublishedDone = absoluteDone
			return nil
		},
	}
	written, copyErr := CopyResponse(ctx, part, body, copyOptions)
	writtenInfo, statErr := part.Stat()
	bodyCloseErr := body.Close()
	body = nil
	syncErr := syncAndClose(part)
	if syncErr != nil {
		return "", currentAttempt, syncErr
	}
	if bodyCloseErr != nil {
		return "", currentAttempt, bodyCloseErr
	}
	if statErr != nil {
		return "", currentAttempt, fmt.Errorf("stat written part: %w", statErr)
	}
	if err := verifyRegularIdentity(paths.Part, writtenInfo); err != nil {
		return "", currentAttempt, err
	}
	if copyErr != nil {
		return "", currentAttempt, copyErr
	}

	finalDone := offset + written
	if err := queue.updateAttemptDone(runtime, currentAttempt, finalDone); err != nil {
		return "", currentAttempt, err
	}
	if finalDone != lastPublishedDone {
		now := queue.clock.Now()
		speed := progressSpeed(lastEventAt, now, lastEventDone, finalDone)
		if err := queue.persistAndPublishProgress(
			runtime,
			currentAttempt,
			speed,
		); err != nil {
			return "", currentAttempt, err
		}
	}
	if task.TotalKnown && finalDone != task.Total {
		return "", currentAttempt, ErrTransferIncomplete
	}
	if err := activateCompleted(paths, writtenInfo); err != nil {
		return "", currentAttempt, err
	}
	return paths.Completed, currentAttempt, nil
}

func (queue *Queue) restartAttempt(
	runtime *runtimeTask,
	attempt uint64,
	response OpenResponse,
) (uint64, Task, error) {
	queue.mu.Lock()
	if !queue.attemptCurrent(runtime, attempt) || runtime.stop != stopNone {
		queue.mu.Unlock()
		return attempt, Task{}, context.Canceled
	}
	runtime.task.Attempt++
	runtime.task.Done = 0
	runtime.task.Validator = response.Validator
	task := runtime.task
	newAttempt := task.Attempt
	queue.mu.Unlock()
	if err := queue.store.Write(task); err != nil {
		return newAttempt, task, err
	}
	queue.publishBestEffort(Event{
		Type:       EventStarted,
		RequestID:  task.RequestID,
		AppID:      task.AppID,
		Attempt:    task.Attempt,
		Done:       0,
		TotalKnown: task.TotalKnown,
		Total:      task.Total,
	})
	return newAttempt, task, nil
}

func (queue *Queue) applyResponseFacts(
	runtime *runtimeTask,
	attempt uint64,
	response OpenResponse,
) (Task, error) {
	queue.mu.Lock()
	if !queue.attemptCurrent(runtime, attempt) || runtime.stop != stopNone {
		queue.mu.Unlock()
		return Task{}, context.Canceled
	}
	if runtime.task.Validator.Empty() {
		runtime.task.Validator = response.Validator
	}
	task := runtime.task
	queue.mu.Unlock()
	if err := queue.store.Write(task); err != nil {
		return Task{}, err
	}
	return task, nil
}

func (queue *Queue) updateAttemptDone(
	runtime *runtimeTask,
	attempt uint64,
	done int64,
) error {
	queue.mu.Lock()
	defer queue.mu.Unlock()
	if !queue.attemptCurrent(runtime, attempt) ||
		runtime.stop != stopNone ||
		done < runtime.task.Done {
		return context.Canceled
	}
	if runtime.task.TotalKnown && done > runtime.task.Total {
		return ErrTransferTooLarge
	}
	runtime.task.Done = done
	return nil
}

func (queue *Queue) persistAndPublishProgress(
	runtime *runtimeTask,
	attempt uint64,
	speed int64,
) error {
	queue.mu.Lock()
	if !queue.attemptCurrent(runtime, attempt) || runtime.stop != stopNone {
		queue.mu.Unlock()
		return context.Canceled
	}
	task := runtime.task
	queue.mu.Unlock()
	if err := queue.store.Write(task); err != nil {
		return err
	}
	queue.publishBestEffort(Event{
		Type:          EventProgress,
		RequestID:     task.RequestID,
		AppID:         task.AppID,
		Attempt:       task.Attempt,
		Done:          task.Done,
		TotalKnown:    task.TotalKnown,
		Total:         task.Total,
		SpeedBytesSec: speed,
	})
	return nil
}

func progressSpeed(start, end time.Time, startDone, endDone int64) int64 {
	elapsed := end.Sub(start)
	if elapsed <= 0 || endDone <= startDone {
		return 0
	}
	bytesPerSecond := float64(endDone-startDone) / elapsed.Seconds()
	if bytesPerSecond <= 0 {
		return 0
	}
	return int64(bytesPerSecond)
}