goget/internal/protocol/http/parallel_download.go

//go:build linux || freebsd
// +build linux freebsd

package http

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"sync"
	"time"

	"codeberg.org/petrbalvin/goget/internal/auth"
	"codeberg.org/petrbalvin/goget/internal/core"
	"codeberg.org/petrbalvin/goget/internal/output"
)

func (c *Client) downloadParallel(ctx context.Context, req *core.DownloadRequest, totalSize int64, connections int, resumeChunks map[int]int64, username, password string) (*core.DownloadResult, error) {
	startTime := time.Now()

	chunks := createChunks(totalSize, connections, c.config.Parallel)
	if c.config.Output.Verbose {
		fmt.Fprintf(os.Stderr, "[parallel] Split into %d chunks\n", len(chunks))
	}

	tempDir, err := os.MkdirTemp("", "goget-chunks-*")
	if err != nil {
		return nil, core.NewFileError("failed to create temp dir", err)
	}
	defer os.RemoveAll(tempDir)

	// completedChunks tracks per-chunk downloaded bytes for resume metadata.
	completedChunks := make(map[int]int64)
	var completedChunksMu sync.Mutex

	for id, n := range resumeChunks {
		completedChunks[id] = n
	}

	progressChan := make(chan int64, len(chunks)*2)
	var wg sync.WaitGroup
	var firstError error
	var errorMu sync.Mutex

	for i := range chunks {
		chunk := &chunks[i]
		wg.Add(1)

		go func(chunk *core.ChunkInfo) {
			defer wg.Done()

			if resumeChunks != nil {
				if done, ok := resumeChunks[chunk.ID]; ok && done >= (chunk.End-chunk.Start+1) {
					progressChan <- done
					chunk.Status = 2
					return
				}
			}

			chunkPath := filepath.Join(tempDir, fmt.Sprintf("chunk_%04d", chunk.ID))
			err := c.downloadChunk(ctx, req, chunk, chunkPath, progressChan, username, password)
			if err != nil {
				errorMu.Lock()
				if firstError == nil {
					firstError = err
				}
				errorMu.Unlock()
				chunk.Status = 3
				chunk.Error = err
				return
			}
			// Save per-chunk progress for resume on interruption.
			completedChunksMu.Lock()
			completedChunks[chunk.ID] = chunk.End - chunk.Start + 1
			saveParallelResumeMetadata(req, totalSize, completedChunks)
			completedChunksMu.Unlock()
			chunk.Status = 2
		}(chunk)
	}

	go func() {
		wg.Wait()
		close(progressChan)
	}()

	var totalDownloaded int64
	var cancelled bool
aggregationLoop:
	for {
		select {
		case progress, ok := <-progressChan:
			if !ok {
				break aggregationLoop
			}
			totalDownloaded += progress
			if req.ProgressCallback != nil {
				duration := time.Since(startTime)
				speed := float64(totalDownloaded) / duration.Seconds()
				req.ProgressCallback(totalDownloaded, totalSize, speed)
			}
		case <-ctx.Done():
			cancelled = true
			break aggregationLoop
		}
	}

	if cancelled {
		// Context was cancelled (SIGINT/SIGTERM). Return the error so
		// the caller can show "cancelled" instead of "complete", and
		// the progress bar renders the interrupted state. Partial
		// chunks are NOT merged — they remain in temp files for a
		// potential resume later.
		return nil, ctx.Err()
	}

	if firstError != nil {
		return nil, firstError
	}

	if req.Output != "" && req.Output != "-" {
		if err := mergeChunks(req.Output, chunks, tempDir); err != nil {
			return nil, core.NewFileError("failed to merge chunks", err)
		}
	}

	if req.Output != "" {
		output.DeleteResumeMetadata(req.Output)
	}

	duration := time.Since(startTime)
	speed := float64(0)
	if duration.Seconds() > 0 {
		speed = float64(totalSize) / duration.Seconds()
	}

	return &core.DownloadResult{
		BytesDownloaded: totalSize,
		TotalSize:       totalSize,
		Duration:        duration,
		Speed:           speed,
		Protocol:        "HTTP/HTTPS",
		IPVersion:       6,
		OutputPath:      req.Output,
		Parallel:        true,
		ChunksCount:     len(chunks),
	}, nil
}

// saveParallelResumeMetadata saves per-chunk progress to the resume metadata file.
func saveParallelResumeMetadata(req *core.DownloadRequest, totalSize int64, chunks map[int]int64) {
	if req.Output == "" || req.Output == "-" {
		return
	}
	// Build a snapshot of the current chunk progress.
	snapshot := make(map[int]int64, len(chunks))
	for id, n := range chunks {
		snapshot[id] = n
	}
	meta := &output.ResumeMetadata{
		URL:        req.URL.String(),
		Total:      totalSize,
		Downloaded: totalChunkBytes(snapshot),
		LastWrite:  time.Now(),
		Version:    "1.0",
		Chunks:     snapshot,
	}
	_ = meta.Save(req.Output)
}

// totalChunkBytes sums the byte count of all completed chunks.
func totalChunkBytes(chunks map[int]int64) int64 {
	var total int64
	for _, n := range chunks {
		total += n
	}
	return total
}

func createChunks(totalSize int64, count int, cfg *core.ParallelConfig) []core.ChunkInfo {
	if count <= 1 || totalSize <= cfg.MinSize {
		return []core.ChunkInfo{
			{ID: 0, Start: 0, End: totalSize - 1},
		}
	}

	// Divide the file into `count` equal-sized chunks. The last chunk
	// absorbs any remainder. MaxChunkSize is NOT used to further
	// subdivide — the user explicitly asked for N parallel connections
	// via --parallel, not N×M tiny range requests. Each chunk becomes
	// one HTTP Range request handled by one of the N worker goroutines.
	chunkSize := totalSize / int64(count)
	if chunkSize < cfg.MinChunkSize {
		chunkSize = cfg.MinChunkSize
		count = int((totalSize + chunkSize - 1) / chunkSize)
	}

	chunks := make([]core.ChunkInfo, 0, count)
	for i := 0; i < count; i++ {
		start := int64(i) * chunkSize
		end := start + chunkSize - 1
		if i == count-1 || end >= totalSize {
			end = totalSize - 1
		}
		chunks = append(chunks, core.ChunkInfo{
			ID:    i,
			Start: start,
			End:   end,
		})
	}

	return chunks
}

func (c *Client) downloadChunk(ctx context.Context, req *core.DownloadRequest, chunk *core.ChunkInfo, outputPath string, progressChan chan<- int64, username, password string) error {
	chunk.Status = 1

	httpReq, err := http.NewRequestWithContext(ctx, "GET", req.URL.String(), nil)
	if err != nil {
		return core.NewNetworkError("failed to create chunk request", err, req.URL.String())
	}

	httpReq.Header.Set("User-Agent", c.config.UserAgent)
	httpReq.Header.Set("Range", fmt.Sprintf("bytes=%d-%d", chunk.Start, chunk.End))

	if req.ResumeInfo != nil {
		if req.ResumeInfo.ETag != "" {
			httpReq.Header.Set("If-Range", req.ResumeInfo.ETag)
		} else if req.ResumeInfo.LastModified != "" {
			httpReq.Header.Set("If-Range", req.ResumeInfo.LastModified)
		}
	}

	if username != "" {
		auth.BasicAuth(httpReq, username, password)
	}

	for key, value := range req.Headers {
		httpReq.Header.Set(key, value)
	}

	resp, err := c.httpClient.Do(httpReq)
	if err != nil {
		return core.NewNetworkError("failed to fetch chunk", err, req.URL.String())
	}
	defer resp.Body.Close()
	c.updateHSTS(resp)

	if resp.StatusCode != http.StatusPartialContent {
		return core.NewProtocolError(
			fmt.Sprintf("chunk HTTP error: %s", resp.Status),
			nil,
			req.URL.String(),
		)
	}

	file, err := os.Create(outputPath)
	if err != nil {
		return core.NewFileError("failed to create chunk file", err)
	}
	defer file.Close()

	reader := io.Reader(resp.Body)
	if c.config.RateLimiter != nil {
		reader = c.config.RateLimiter.WrapReader(reader)
	}

	buf := make([]byte, c.config.Transport.BufferSize)
	var downloaded int64

	for {
		if ctx.Err() != nil {
			return ctx.Err()
		}

		n, err := reader.Read(buf)
		if n > 0 {
			_, wErr := file.Write(buf[:n])
			if wErr != nil {
				return core.NewFileError("failed to write chunk", wErr)
			}
			downloaded += int64(n)
			if progressChan != nil {
				progressChan <- int64(n)
			}
		}

		if err == io.EOF {
			break
		}
		if err != nil {
			return core.NewNetworkError("failed to read chunk", err, req.URL.String())
		}
	}

	chunk.Downloaded = downloaded
	return nil
}

func mergeChunks(outputPath string, chunks []core.ChunkInfo, tempDir string) error {
	outFile, err := os.Create(outputPath)
	if err != nil {
		return err
	}
	defer outFile.Close()

	for _, chunk := range chunks {
		chunkPath := filepath.Join(tempDir, fmt.Sprintf("chunk_%04d", chunk.ID))
		chunkFile, err := os.Open(chunkPath)
		if err != nil {
			return fmt.Errorf("failed to open chunk %d: %w", chunk.ID, err)
		}
		_, err = io.Copy(outFile, chunkFile)
		chunkFile.Close()
		if err != nil {
			return fmt.Errorf("failed to copy chunk %d: %w", chunk.ID, err)
		}
	}

	return nil
}