Skip to content

Commit 2c10592

Browse files
SuperCoolPencilMeet Mehta
SuperCoolPencil
and
Meet Mehta
authored
feat: fix some minor engine bugs and add heuristics for a more stable download (#83)
* feat: Implement hard stall detection for unresponsive workers and refine slow worker detection to correctly identify zero-speed tasks, accompanied by new custom tests. * feat: implement a new EndGame heuristic * test: add unit tests for endgame shadow work assignment and cleanup logic. --------- Co-authored-by: Meet Mehta <supercoolpencil@duck.com>
1 parent 486f435 commit 2c10592

5 files changed

Lines changed: 347 additions & 6 deletions

File tree

‎internal/engine/concurrent/downloader.go‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,50 @@ func (d *ConcurrentDownloader) getInitialConnections(fileSize int64) int {
6969
return maxConns
7070
}
7171

72+
// cleanUpShadows cancels any other workers working on the same completed chunk
73+
func (d *ConcurrentDownloader) cleanUpShadows(completedOffset int64, winnerID int) {
74+
d.activeMu.Lock()
75+
defer d.activeMu.Unlock()
76+
77+
for id, at := range d.activeTasks {
78+
// If it's a different worker but working on the same chunk offset
79+
if id != winnerID && at.Task.Offset == completedOffset {
80+
utils.Debug("EndGame: Worker %d won Chunk offset %d. Cancelling loser Worker %d.", winnerID, completedOffset, id)
81+
if at.Cancel != nil {
82+
at.Cancel()
83+
}
84+
}
85+
}
86+
}
87+
88+
// assignShadowWork assigns idle workers to shadow active tasks (End Game)
89+
func (d *ConcurrentDownloader) assignShadowWork(queue *TaskQueue) {
90+
d.activeMu.Lock()
91+
defer d.activeMu.Unlock()
92+
93+
// Count workers per offset
94+
counts := make(map[int64]int)
95+
for _, at := range d.activeTasks {
96+
counts[at.Task.Offset]++
97+
}
98+
99+
for _, at := range d.activeTasks {
100+
// If this chunk has < 2 workers (i.e. only the original one)
101+
if counts[at.Task.Offset] < 2 {
102+
// Push duplicate task to queue
103+
// Note: We shadow the ORIGINAL task range.
104+
// The shadow worker will check existence on disk/offsets, but since we overwrite it's safe.
105+
// The race is settled by cleanUpShadows.
106+
queue.Push(at.Task)
107+
108+
utils.Debug("EndGame: Double assigning Chunk offset %d (Racing against original owner)", at.Task.Offset)
109+
110+
// Only assign one shadow per tick to prevent flooding queue
111+
return
112+
}
113+
}
114+
}
115+
72116
// ReportMirrorError marks a mirror as having an error in the state
73117
func (d *ConcurrentDownloader) ReportMirrorError(url string) {
74118
if d.State == nil {
@@ -325,6 +369,9 @@ func (d *ConcurrentDownloader) Download(ctx context.Context, rawurl string, cand
325369
// Try to steal from an active worker
326370
if d.StealWork(queue) {
327371
didWork = true
372+
} else {
373+
// If stealing failed (chunks too small/balanced), try Shadowing (End Game)
374+
d.assignShadowWork(queue)
328375
}
329376
}
330377

‎internal/engine/concurrent/endgame_test.go‎

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
package concurrent
2+
3+
import (
4+
"context"
5+
"testing"
6+
7+
"github.com/stretchr/testify/assert"
8+
"github.com/surge-downloader/surge/internal/engine/types"
9+
)
10+
11+
func TestEndGame_AssignShadowWork(t *testing.T) {
12+
// Setup
13+
d := &ConcurrentDownloader{
14+
activeTasks: make(map[int]*ActiveTask),
15+
}
16+
queue := NewTaskQueue()
17+
18+
// 1. Add an active task (single worker)
19+
task := types.Task{Offset: 100, Length: 50}
20+
activeTask := &ActiveTask{
21+
Task: task,
22+
}
23+
d.activeMu.Lock()
24+
d.activeTasks[1] = activeTask
25+
d.activeMu.Unlock()
26+
27+
// 2. Run Shadow Assignment
28+
d.assignShadowWork(queue)
29+
30+
// 3. Verify that a shadow task was pushed to queue
31+
// We check if queue length is 1
32+
if queue.Len() != 1 {
33+
t.Fatalf("Expected 1 shadow task in queue, got %d", queue.Len())
34+
}
35+
36+
popped, ok := queue.Pop()
37+
assert.True(t, ok)
38+
assert.Equal(t, int64(100), popped.Offset)
39+
assert.Equal(t, int64(50), popped.Length)
40+
41+
// 4. Add a second worker for the SAME chunks (simulation of race already started)
42+
d.activeMu.Lock()
43+
d.activeTasks[2] = &ActiveTask{Task: task}
44+
d.activeMu.Unlock()
45+
46+
// 5. Run Shadow Assignment again
47+
d.assignShadowWork(queue)
48+
49+
// 6. Verify NO new shadow task (count is already 2)
50+
assert.Equal(t, 0, queue.Len(), "Should not assign shadow if 2 workers already valid")
51+
}
52+
53+
func TestEndGame_CleanUpShadows(t *testing.T) {
54+
d := &ConcurrentDownloader{
55+
activeTasks: make(map[int]*ActiveTask),
56+
}
57+
58+
// Two workers racing on same offset
59+
offset := int64(500)
60+
61+
ctx1, cancel1 := context.WithCancel(context.Background())
62+
task1 := &ActiveTask{
63+
Task: types.Task{Offset: offset},
64+
Cancel: cancel1,
65+
}
66+
67+
ctx2, cancel2 := context.WithCancel(context.Background())
68+
task2 := &ActiveTask{
69+
Task: types.Task{Offset: offset},
70+
Cancel: cancel2,
71+
}
72+
73+
// Unrelated task
74+
ctx3, cancel3 := context.WithCancel(context.Background())
75+
task3 := &ActiveTask{
76+
Task: types.Task{Offset: 999},
77+
Cancel: cancel3,
78+
}
79+
80+
d.activeMu.Lock()
81+
d.activeTasks[1] = task1
82+
d.activeTasks[2] = task2
83+
d.activeTasks[3] = task3
84+
d.activeMu.Unlock()
85+
86+
// Winner is Worker 1
87+
d.cleanUpShadows(offset, 1)
88+
89+
// Worker 2 (Loser) should be cancelled
90+
select {
91+
case <-ctx2.Done():
92+
// Success
93+
default:
94+
t.Fatal("Shadow worker (loser) was not cancelled")
95+
}
96+
97+
// Worker 1 (Winner) should NOT be cancelled (obviously, but good to check logic safety)
98+
select {
99+
case <-ctx1.Done():
100+
t.Fatal("Winner worker was cancelled incorrectly")
101+
default:
102+
// Success
103+
}
104+
105+
// Unrelated Worker 3 should NOT be cancelled
106+
select {
107+
case <-ctx3.Done():
108+
t.Fatal("Unrelated worker was cancelled incorrectly")
109+
default:
110+
// Success
111+
}
112+
}

‎internal/engine/concurrent/health.go‎

Lines changed: 31 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package concurrent
22

33
import (
4+
"sync/atomic"
45
"time"
56

67
"github.com/surge-downloader/surge/internal/utils"
@@ -16,6 +17,8 @@ func (d *ConcurrentDownloader) checkWorkerHealth() {
1617
}
1718

1819
now := time.Now()
20+
stallTimeout := d.Runtime.GetStallTimeout()
21+
gracePeriod := d.Runtime.GetSlowWorkerGracePeriod()
1922

2023
// First pass: calculate mean speed
2124
var totalSpeed float64
@@ -35,21 +38,43 @@ func (d *ConcurrentDownloader) checkWorkerHealth() {
3538
// Second pass: check for slow workers
3639
for workerID, active := range d.activeTasks {
3740

38-
// timeSinceActivity := now.Sub(lastTime)
41+
// ---------------------------------------------------------
42+
// CHECK 1: The Hard Stall (The "Zombie" Killer)
43+
// ---------------------------------------------------------
44+
// If the worker hasn't written a single byte for > StallTimeout (5s),
45+
// it is dead. Kill it immediately. Speed is irrelevant here.
46+
lastActivityNano := atomic.LoadInt64(&active.LastActivity)
47+
lastActivity := time.Unix(0, lastActivityNano)
48+
timeSinceLastActivity := now.Sub(lastActivity)
49+
50+
if timeSinceLastActivity > stallTimeout {
51+
utils.Debug("Health: Worker %d stalled. No activity for %v. Cancelling.",
52+
workerID, timeSinceLastActivity)
53+
54+
if active.Cancel != nil {
55+
active.Cancel()
56+
}
57+
continue
58+
}
59+
60+
// ---------------------------------------------------------
61+
// CHECK 2: The Relative Slowness (The "Straggler" Killer)
62+
// ---------------------------------------------------------
3963
taskDuration := now.Sub(active.StartTime)
4064

41-
// Skip workers that are still in their grace period
42-
gracePeriod := d.Runtime.GetSlowWorkerGracePeriod()
65+
// SKIP if in grace period
4366
if taskDuration < gracePeriod {
4467
continue
4568
}
4669

47-
// Check for slow worker
48-
// Only cancel if: below threshold
70+
// Only run this if we have a valid global baseline
4971
if meanSpeed > 0 {
5072
workerSpeed := active.GetSpeed()
73+
74+
// BUG FIX: Removed "workerSpeed > 0".
75+
// If speed is 0 (and not caught by stall check yet), it IS below threshold.
5176
threshold := d.Runtime.GetSlowWorkerThreshold()
52-
isBelowThreshold := workerSpeed > 0 && workerSpeed < threshold*meanSpeed
77+
isBelowThreshold := workerSpeed < threshold*meanSpeed
5378

5479
if isBelowThreshold {
5580
utils.Debug("Health: Worker %d slow (%.2f KB/s vs mean %.2f KB/s), cancelling",

‎internal/engine/concurrent/health_custom_test.go‎

Lines changed: 154 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,154 @@
1+
package concurrent
2+
3+
import (
4+
"context"
5+
"testing"
6+
"time"
7+
8+
"github.com/surge-downloader/surge/internal/engine/types"
9+
)
10+
11+
func TestHealthCheck_StallDetection(t *testing.T) {
12+
// Setup
13+
runtime := &types.RuntimeConfig{
14+
StallTimeout: 100 * time.Millisecond,
15+
SlowWorkerGracePeriod: 0,
16+
}
17+
d := &ConcurrentDownloader{
18+
Runtime: runtime,
19+
activeTasks: make(map[int]*ActiveTask),
20+
}
21+
22+
// Create a "stalled" task
23+
ctx, cancel := context.WithCancel(context.Background())
24+
now := time.Now()
25+
stalledTask := &ActiveTask{
26+
StartTime: now.Add(-1 * time.Second),
27+
LastActivity: now.Add(-500 * time.Millisecond).UnixNano(), // Past timeout (100ms)
28+
Cancel: cancel,
29+
}
30+
// Speed is 0, but check 2 doesn't matter for this test
31+
stalledTask.Speed = 0
32+
33+
d.activeMu.Lock()
34+
d.activeTasks[1] = stalledTask
35+
d.activeMu.Unlock()
36+
37+
// Run Health Check
38+
d.checkWorkerHealth()
39+
40+
// Verify cancellation
41+
select {
42+
case <-ctx.Done():
43+
// Success: context was cancelled
44+
default:
45+
t.Fatal("Stalled worker was not cancelled")
46+
}
47+
}
48+
49+
func TestHealthCheck_StragglerDetection(t *testing.T) {
50+
// Setup
51+
runtime := &types.RuntimeConfig{
52+
StallTimeout: 10 * time.Second, // Long timeout so stall check passes
53+
SlowWorkerGracePeriod: 0,
54+
SlowWorkerThreshold: 0.5,
55+
}
56+
d := &ConcurrentDownloader{
57+
Runtime: runtime,
58+
activeTasks: make(map[int]*ActiveTask),
59+
}
60+
61+
ctx1, cancel1 := context.WithCancel(context.Background())
62+
ctx2, cancel2 := context.WithCancel(context.Background())
63+
now := time.Now()
64+
65+
// Task 1: Fast (Baseline)
66+
fastTask := &ActiveTask{
67+
StartTime: now.Add(-5 * time.Second),
68+
LastActivity: now.UnixNano(),
69+
Cancel: cancel1,
70+
Speed: 1000,
71+
}
72+
73+
// Task 2: Slow (Straggler)
74+
// Speed 100 is < 50% of mean ((1000+100)/2 = 550)
75+
slowTask := &ActiveTask{
76+
StartTime: now.Add(-5 * time.Second),
77+
LastActivity: now.UnixNano(),
78+
Cancel: cancel2,
79+
Speed: 100,
80+
}
81+
82+
d.activeMu.Lock()
83+
d.activeTasks[1] = fastTask
84+
d.activeTasks[2] = slowTask
85+
d.activeMu.Unlock()
86+
87+
// Run Health Check
88+
d.checkWorkerHealth()
89+
90+
// Verify cancellation
91+
select {
92+
case <-ctx1.Done():
93+
t.Fatal("Fast worker was incorrectly cancelled")
94+
default:
95+
// OK
96+
}
97+
98+
select {
99+
case <-ctx2.Done():
100+
// OK
101+
default:
102+
t.Fatal("Slow worker was NOT cancelled")
103+
}
104+
}
105+
106+
func TestHealthCheck_ZeroSpeedDetection(t *testing.T) {
107+
// Setup
108+
runtime := &types.RuntimeConfig{
109+
StallTimeout: 10 * time.Second,
110+
SlowWorkerGracePeriod: 0,
111+
SlowWorkerThreshold: 0.5,
112+
}
113+
d := &ConcurrentDownloader{
114+
Runtime: runtime,
115+
activeTasks: make(map[int]*ActiveTask),
116+
}
117+
118+
_, cancel1 := context.WithCancel(context.Background())
119+
ctx2, cancel2 := context.WithCancel(context.Background())
120+
now := time.Now()
121+
122+
// Task 1: Fast
123+
fastTask := &ActiveTask{
124+
StartTime: now.Add(-5 * time.Second),
125+
LastActivity: now.UnixNano(),
126+
Cancel: cancel1,
127+
Speed: 1000,
128+
}
129+
130+
// Task 2: Zero Speed (but not stalled by time yet)
131+
// Should be caught by Check 2 because 0 < threshold * mean
132+
zeroTask := &ActiveTask{
133+
StartTime: now.Add(-5 * time.Second),
134+
LastActivity: now.UnixNano(),
135+
Cancel: cancel2,
136+
Speed: 0,
137+
}
138+
139+
d.activeMu.Lock()
140+
d.activeTasks[1] = fastTask
141+
d.activeTasks[2] = zeroTask
142+
d.activeMu.Unlock()
143+
144+
// Run Health Check
145+
d.checkWorkerHealth()
146+
147+
// Verify cancellation
148+
select {
149+
case <-ctx2.Done():
150+
// Success
151+
default:
152+
t.Fatal("Zero speed worker was NOT cancelled by relative check")
153+
}
154+
}

0 commit comments

Comments
 (0)