diff --git a/internal/runner/api/jobconn.go b/internal/runner/api/jobconn.go index ee3fbbf..5b72cc1 100644 --- a/internal/runner/api/jobconn.go +++ b/internal/runner/api/jobconn.go @@ -18,6 +18,7 @@ type JobConnection struct { writeMu sync.Mutex stopPing chan struct{} stopHeartbeat chan struct{} + stopOnce sync.Once isConnected bool connMu sync.RWMutex } @@ -132,13 +133,12 @@ func (j *JobConnection) pingLoop() { // Heartbeat sends a heartbeat message over WebSocket to keep runner online. func (j *JobConnection) Heartbeat() { + j.writeMu.Lock() + defer j.writeMu.Unlock() if j.conn == nil { return } - j.writeMu.Lock() - defer j.writeMu.Unlock() - msg := map[string]interface{}{ "type": "runner_heartbeat", "timestamp": time.Now().Unix(), @@ -178,27 +178,34 @@ func (j *JobConnection) heartbeatLoop() { } } +// stopLoops signals ping/heartbeat goroutines to exit. +// Channels are closed once and left non-nil so loops can receive without racing Close. +func (j *JobConnection) stopLoops() { + j.stopOnce.Do(func() { + if j.stopHeartbeat != nil { + close(j.stopHeartbeat) + } + if j.stopPing != nil { + close(j.stopPing) + } + }) +} + // Close closes the WebSocket connection. func (j *JobConnection) Close() { + j.stopLoops() + + j.writeMu.Lock() + defer j.writeMu.Unlock() + j.connMu.Lock() j.isConnected = false + conn := j.conn + j.conn = nil j.connMu.Unlock() - // Stop heartbeat goroutine - if j.stopHeartbeat != nil { - close(j.stopHeartbeat) - j.stopHeartbeat = nil - } - - // Stop ping goroutine - if j.stopPing != nil { - close(j.stopPing) - j.stopPing = nil - } - - if j.conn != nil { - j.conn.Close() - j.conn = nil + if conn != nil { + conn.Close() } } @@ -211,13 +218,12 @@ func (j *JobConnection) IsConnected() bool { // Log sends a log entry to the manager. func (j *JobConnection) Log(taskID int64, level types.LogLevel, message string) { + j.writeMu.Lock() + defer j.writeMu.Unlock() if j.conn == nil { return } - j.writeMu.Lock() - defer j.writeMu.Unlock() - msg := map[string]interface{}{ "type": "log_entry", "data": map[string]interface{}{ @@ -242,13 +248,12 @@ func (j *JobConnection) Log(taskID int64, level types.LogLevel, message string) // Progress sends a progress update to the manager. func (j *JobConnection) Progress(taskID int64, progress float64) { + j.writeMu.Lock() + defer j.writeMu.Unlock() if j.conn == nil { return } - j.writeMu.Lock() - defer j.writeMu.Unlock() - msg := map[string]interface{}{ "type": "progress", "data": map[string]interface{}{ @@ -272,13 +277,12 @@ func (j *JobConnection) Progress(taskID int64, progress float64) { // OutputUploaded notifies that an output file was uploaded. func (j *JobConnection) OutputUploaded(taskID int64, fileName string) { + j.writeMu.Lock() + defer j.writeMu.Unlock() if j.conn == nil { return } - j.writeMu.Lock() - defer j.writeMu.Unlock() - msg := map[string]interface{}{ "type": "output_uploaded", "data": map[string]interface{}{ @@ -307,14 +311,13 @@ func (j *JobConnection) OutputUploaded(taskID int64, fileName string) { // freeRequeue asks the manager to requeue a failure without incrementing retry_count // (used when this attempt newly armed GPU lockout). func (j *JobConnection) Complete(taskID int64, success bool, errorMsg error, freeRequeue bool) { + j.writeMu.Lock() + defer j.writeMu.Unlock() if j.conn == nil { log.Printf("Cannot send task complete: WebSocket connection is nil") return } - j.writeMu.Lock() - defer j.writeMu.Unlock() - data := map[string]interface{}{ "task_id": taskID, "success": success,