fix: minor issues with pause/resume, eta calculation, elapsed calculation, task_queue (#50)

* feat: Implement elapsed time tracking for downloads and introduce an intermediate "pausing" state for better UX.

* fix: add idempotency check to prevent re-pausing if the download is already in a pausing or paused state.

* feat: Add checks to prevent resuming a download that is still pausing or already running.

* feat: Distinguish between total and session elapsed time in progress tracking for accurate speed calculation and reporting.

* fix: while traversing task queue, we used range q.tasks, this tries to split tasks that have already been split

* fix: Handle file rename race condition by checking if the destination file already exists with the correct size.

* test: add tests to ensure all the bugs mentioned dont happen again

* feat: Add temporary database setup and a worker pool to `TestStateSync` and adjust sleep durations for improved state synchronization testing.

---------

Co-authored-by: Meet Mehta <supercoolpencil@duck.com>

Author: SuperCoolPencil

internal/download/manager.go

@@ -178,10 +178,15 @@ func TUIDownload(ctx context.Context, cfg *types.DownloadConfig) error {
 
 	isPaused := cfg.State != nil && cfg.State.IsPaused()
 	if downloadErr == nil && !isPaused && cfg.ProgressCh != nil {
+		elapsed := time.Since(start)
+		// For resumed downloads, add previously saved elapsed time
+		if cfg.State != nil && cfg.State.SavedElapsed > 0 {
+			elapsed += cfg.State.SavedElapsed
+		}
 		cfg.ProgressCh <- events.DownloadCompleteMsg{
 			DownloadID: cfg.ID,
 			Filename:   finalFilename,
-			Elapsed:    time.Since(start),
+			Elapsed:    elapsed,
 			Total:      probe.FileSize,
 		}
 	}

internal/download/pool.go

@@ -3,10 +3,12 @@ package download
 import (
 	"context"
 	"sync"
+	"time"
 
 	"github.com/surge-downloader/surge/internal/engine/events"
 	"github.com/surge-downloader/surge/internal/engine/state"
 	"github.com/surge-downloader/surge/internal/engine/types"
+	"github.com/surge-downloader/surge/internal/utils"
 )
 
 // activeDownload tracks a download that's currently running
@@ -84,6 +86,11 @@ func (p *WorkerPool) Pause(downloadID string) {
 
 	// Set paused flag and cancel context
 	if ad.config.State != nil {
+		// Idempotency: If already pausing or paused, do nothing
+		if ad.config.State.IsPausing() || ad.config.State.IsPaused() {
+			return
+		}
+		ad.config.State.SetPausing(true) // Mark as transitioning to pause
 		ad.config.State.Pause()
 	}
 
@@ -105,8 +112,8 @@ func (p *WorkerPool) PauseAll() {
 	p.mu.RLock()
 	ids := make([]string, 0, len(p.downloads)) //This stores the uuids of the downloads to be paused
 	for id, ad := range p.downloads {
-		// Only pause downloads that are actually active (not already paused or done)
-		if ad != nil && ad.config.State != nil && !ad.config.State.IsPaused() && !ad.config.State.Done.Load() {
+		// Only pause downloads that are actually active (not already paused or done or pausing)
+		if ad != nil && ad.config.State != nil && !ad.config.State.IsPaused() && !ad.config.State.Done.Load() && !ad.config.State.IsPausing() {
 			ids = append(ids, id)
 		}
 	}
@@ -151,6 +158,18 @@ func (p *WorkerPool) Resume(downloadID string) {
 		return
 	}
 
+	// Prevent race: Don't resume if still pausing
+	if ad.config.State != nil && ad.config.State.IsPausing() {
+		utils.Debug("Resume ignored: download %s is still pausing", downloadID)
+		return
+	}
+
+	// Idempotency: If already running (not paused), do nothing
+	if ad.config.State != nil && !ad.config.State.IsPaused() {
+		utils.Debug("Resume ignored: download %s is already running", downloadID)
+		return
+	}
+
 	// Clear paused flag and reset session start to avoid speed spikes/dips checks
 	if ad.config.State != nil {
 		ad.config.State.Resume()
@@ -187,10 +206,21 @@ func (p *WorkerPool) worker() {
 
 		err := TUIDownload(ctx, &ad.config)
 
-		// Check if this was a pause (not an error)
-		isPaused := cfg.State != nil && cfg.State.IsPaused()
+		// Logic:
+		// 1. If Pause() was called: State.IsPaused() is true. We keep the task in p.downloads (so it can be resumed).
+		// 2. If finished/error: We remove from p.downloads.
+
+		isPaused := ad.config.State != nil && ad.config.State.IsPaused()
+
+		// Clear "Pausing" transition state now that worker has exited
+		if ad.config.State != nil {
+			ad.config.State.SetPausing(false)
+		}
 
-		if err != nil && !isPaused {
+		if isPaused {
+			utils.Debug("WorkerPool: Download %s paused cleanly", cfg.ID)
+			// If paused, we keep it in downloads map for potential resume
+		} else if err != nil {
 			if cfg.State != nil {
 				cfg.State.SetError(err)
 			}
@@ -255,7 +285,9 @@ func (p *WorkerPool) GetStatus(id string) *types.DownloadStatus {
 		Status:     "downloading",
 	}
 
-	if state.IsPaused() {
+	if ad.config.State.IsPausing() {
+		status.Status = "pausing"
+	} else if ad.config.State.IsPaused() {
 		status.Status = "paused"
 	} else if state.Done.Load() {
 		status.Status = "completed"
@@ -272,10 +304,10 @@ func (p *WorkerPool) GetStatus(id string) *types.DownloadStatus {
 	}
 
 	// Calculate speed (MB/s)
-	downloaded, _, elapsed, _, sessionStart := state.GetProgress()
+	downloaded, _, _, sessionElapsed, _, sessionStart := state.GetProgress()
 	sessionDownloaded := downloaded - sessionStart
-	if elapsed.Seconds() > 0 && sessionDownloaded > 0 {
-		bytesPerSec := float64(sessionDownloaded) / elapsed.Seconds()
+	if sessionElapsed.Seconds() > 0 && sessionDownloaded > 0 {
+		bytesPerSec := float64(sessionDownloaded) / sessionElapsed.Seconds()
 		status.Speed = bytesPerSec / (1024 * 1024)
 	}
 
@@ -286,5 +318,38 @@ func (p *WorkerPool) GetStatus(id string) *types.DownloadStatus {
 func (p *WorkerPool) GracefulShutdown() {
 	// ... existing implementation
 	p.PauseAll()
+
+	// Wait for any downloads in "Pausing" state to finish transitioning
+	// This ensures we don't exit while a database write is pending/active
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	ticker := time.NewTicker(100 * time.Millisecond)
+	defer ticker.Stop()
+
+	for {
+		p.mu.RLock()
+		stillPausing := false
+		for _, ad := range p.downloads {
+			if ad.config.State != nil && ad.config.State.IsPausing() {
+				stillPausing = true
+				break
+			}
+		}
+		p.mu.RUnlock()
+
+		if !stillPausing {
+			break
+		}
+
+		select {
+		case <-ctx.Done():
+			utils.Debug("GracefulShutdown: timed out waiting for downloads to pause")
+			return // Return from function, loop will exit
+		case <-ticker.C:
+			continue
+		}
+	}
+
 	p.wg.Wait() // Blocks until all workers call Done()
 }

internal/download/pool_test.go

@@ -500,6 +500,9 @@ func TestWorkerPool_Resume_RequeuesDownload(t *testing.T) {
 		State: state,
 	}
 
+	// Resume checks idempotency, so we MUST start in Paused state
+	state.Paused.Store(true)
+
 	pool.mu.Lock()
 	pool.downloads["test-id"] = &activeDownload{
 		config: cfg,
@@ -547,7 +550,21 @@ func TestWorkerPool_GracefulShutdown_PausesAll(t *testing.T) {
 	pool.mu.Unlock()
 
 	// GracefulShutdown should call PauseAll
+	// PauseAll will set IsPausing() = true
+	// GracefulShutdown waits for IsPausing() = false
+	// We verify that PauseAll was called by checking state.IsPausing()
+	// Then we clear it to unblock shutdown
+
 	done := make(chan bool)
+	go func() {
+		// Wait for PauseAll to be called (Pausing=true)
+		for !state.IsPausing() {
+			time.Sleep(10 * time.Millisecond)
+		}
+		// Simulate worker finishing pause transition
+		state.SetPausing(false)
+	}()
+
 	go func() {
 		pool.GracefulShutdown()
 		done <- true
@@ -628,7 +645,76 @@ func TestWorkerPool_HasDownload(t *testing.T) {
 		t.Error("Expected HasDownload to return false for missing download")
 	}
 
-	// 3. Persistent Check is harder to test without mocking 'state' package or setting up DB.
-	// However, we can trust the integration test for that part or mock it if we refactored.
 	// For now, this unit test covers the memory-check part of HasDownload which was the critical logic add.
 }
+
+func TestWorkerPool_PauseResume_Idempotency(t *testing.T) {
+	ch := make(chan any, 10)
+	pool := NewWorkerPool(ch, 3)
+
+	state := types.NewProgressState("idempotent-test", 1000)
+
+	pool.mu.Lock()
+	pool.downloads["idempotent-test"] = &activeDownload{
+		config: types.DownloadConfig{
+			ID:    "idempotent-test",
+			State: state,
+		},
+	}
+	pool.mu.Unlock()
+
+	// 1. First Pause
+	pool.Pause("idempotent-test")
+
+	// Should be Pausing
+	if !state.IsPausing() {
+		t.Error("Expected state to be Pausing after first Pause")
+	}
+
+	// Consume message
+	select {
+	case <-ch:
+		// OK
+	case <-time.After(100 * time.Millisecond):
+		t.Error("Expected pause message")
+	}
+
+	// 2. Second Pause (Idempotent)
+	// Should NOT send another message or change state significantly (still Pausing/Paused)
+	pool.Pause("idempotent-test")
+
+	select {
+	case <-ch:
+		t.Error("Did not expect second pause message")
+	default:
+		// OK
+	}
+
+	// Manually transition to Paused (simulating worker finish)
+	state.SetPausing(false)
+	state.Pause()
+
+	// 3. Resume
+	pool.Resume("idempotent-test")
+	if state.IsPaused() {
+		t.Error("Expected state to be Resumed (not paused)")
+	}
+
+	// Consume resume message
+	select {
+	case <-ch:
+		// OK
+	case <-time.After(100 * time.Millisecond):
+		t.Error("Expected resume message")
+	}
+
+	// 4. Second Resume (Idempotent)
+	pool.Resume("idempotent-test")
+
+	select {
+	case <-ch:
+		t.Error("Did not expect second resume message")
+	default:
+		// OK
+	}
+}

internal/engine/concurrent/downloader.go

@@ -204,6 +204,8 @@ func (d *ConcurrentDownloader) Download(ctx context.Context, rawurl, destPath st
 		tasks = savedState.Tasks
 		if d.State != nil {
 			d.State.Downloaded.Store(savedState.Downloaded)
+			// Restore elapsed time from previous sessions
+			d.State.SetSavedElapsed(time.Duration(savedState.Elapsed))
 			// Fix speed spike: sync session start so we don't count previous bytes as new speed
 			d.State.SyncSessionStart()
 		}
@@ -348,6 +350,14 @@ func (d *ConcurrentDownloader) Download(ctx context.Context, rawurl, destPath st
 		}
 		computedDownloaded := fileSize - remainingBytes
 
+		// Calculate total elapsed time
+		var totalElapsed time.Duration
+		if d.State != nil {
+			totalElapsed = d.State.SavedElapsed + time.Since(startTime)
+		} else {
+			totalElapsed = time.Since(startTime)
+		}
+
 		// Save state for resume (use computed value for consistency)
 		s := &types.DownloadState{
 			URL:        d.URL,
@@ -357,6 +367,7 @@ func (d *ConcurrentDownloader) Download(ctx context.Context, rawurl, destPath st
 			Downloaded: computedDownloaded,
 			Tasks:      remainingTasks,
 			Filename:   filepath.Base(destPath),
+			Elapsed:    totalElapsed.Nanoseconds(),
 		}
 		if err := state.SaveState(d.URL, destPath, s); err != nil {
 			utils.Debug("Failed to save pause state: %v", err)
@@ -387,6 +398,15 @@ func (d *ConcurrentDownloader) Download(ctx context.Context, rawurl, destPath st
 
 	// Rename from .surge to final destination
 	if err := os.Rename(workingPath, destPath); err != nil {
+		// Check for race condition: did someone else already rename it?
+		if os.IsNotExist(err) {
+			if info, statErr := os.Stat(destPath); statErr == nil && info.Size() == fileSize {
+				utils.Debug("Race condition detected: File already exists and has correct size. Treating as success.")
+				// Clean up state just in case, though usually done by caller
+				_ = state.DeleteState(d.ID, d.URL, destPath)
+				return nil
+			}
+		}
 		return fmt.Errorf("failed to rename completed file: %w", err)
 	}
 

internal/engine/concurrent/task_queue.go

@@ -106,7 +106,8 @@ func (q *TaskQueue) SplitLargestIfNeeded() bool {
 	// Find largest queued task
 	idx := -1
 	var maxLen int64 = 0
-	for i, t := range q.tasks {
+	for i := q.head; i < len(q.tasks); i++ {
+		t := q.tasks[i]
 		if t.Length > maxLen && t.Length > 2*types.MinChunk {
 			maxLen = t.Length
 			idx = i

internal/engine/concurrent/task_queue_test.go

@@ -132,3 +132,59 @@ func TestAlignedSplitSize(t *testing.T) {
 		}
 	}
 }
+func TestTaskQueue_SplitLargestIfNeeded_Regression(t *testing.T) {
+	q := NewTaskQueue()
+
+	// Push 2 tasks
+	q.Push(types.Task{Offset: 0, Length: types.MinChunk})       // Task 0
+	q.Push(types.Task{Offset: 100, Length: 4 * types.MinChunk}) // Task 1 (should be split)
+
+	// Pop Task 0 (Head moves to 1)
+	if _, ok := q.Pop(); !ok {
+		t.Fatal("Failed to pop task 0")
+	}
+
+	// Now Head is at 1 (Task 1). Task 0 is "dead" space in slice.
+	// The bug was that SplitLargestIfNeeded started iteration at 0.
+	// If it starts at 0, it might enable splitting of Task 0 again if it was large enough,
+	// or mess up indexing.
+
+	// In this specific regression setup:
+	// We want to ensure it splits Task 1 (index 1), NOT Task 0 (index 0).
+	// Although Task 0 is small here, if the loop ignores head, it processes garbage/old data.
+
+	if !q.SplitLargestIfNeeded() {
+		t.Error("Should have split Task 1")
+	}
+
+	// Check queue state
+	// Should have: Task 1 (split right half) at head, and Task 1 (split left half) appended.
+	// Len should be 2.
+	if q.Len() != 2 {
+		t.Errorf("Len = %d, want 2", q.Len())
+	}
+
+	// Check that we didn't resurrect Task 0 or duplicate it.
+	t1, _ := q.Pop()
+	t2, _ := q.Pop()
+
+	// Original Task 1 was Offset 100, Len 4*MinChunk
+	// Split: Right half replaces it, Left half appended.
+	// Left: Offset 100, Len 2*MinChunk
+	// Right: Offset 100+2*MinChunk, Len 2*MinChunk
+
+	// Implementation detail: task_queue.go:132: q.tasks[idx] = right; q.tasks = append(q.tasks, left)
+	// So next pop (t1) should be "right" half.
+	// And t2 should be "left" half.
+
+	// Verify T1 (Right Half)
+	expectedRightOffset := int64(100 + alignedSplitSize(4*types.MinChunk))
+	if t1.Offset != expectedRightOffset {
+		t.Errorf("First popped task offset %d, want %d", t1.Offset, expectedRightOffset)
+	}
+
+	// Verify T2 (Left Half)
+	if t2.Offset != 100 {
+		t.Errorf("Second popped task offset %d, want 100", t2.Offset)
+	}
+}

internal/engine/events/events.go

@@ -12,6 +12,7 @@ type ProgressMsg struct {
 	Downloaded        int64
 	Total             int64
 	Speed             float64 // bytes per second
+	Elapsed           time.Duration
 	ActiveConnections int
 }