Fix JobConnection Close races with ping/heartbeat loops.
CI / check-and-test (pull_request) Successful in 1m2s
CI / check-and-test (pull_request) Successful in 1m2s
Stop nil-ing stop channels after close, close them once, and serialize conn teardown so go test -race stays green.
This commit is contained in:
@@ -18,6 +18,7 @@ type JobConnection struct {
|
|||||||
writeMu sync.Mutex
|
writeMu sync.Mutex
|
||||||
stopPing chan struct{}
|
stopPing chan struct{}
|
||||||
stopHeartbeat chan struct{}
|
stopHeartbeat chan struct{}
|
||||||
|
stopOnce sync.Once
|
||||||
isConnected bool
|
isConnected bool
|
||||||
connMu sync.RWMutex
|
connMu sync.RWMutex
|
||||||
}
|
}
|
||||||
@@ -132,13 +133,12 @@ func (j *JobConnection) pingLoop() {
|
|||||||
|
|
||||||
// Heartbeat sends a heartbeat message over WebSocket to keep runner online.
|
// Heartbeat sends a heartbeat message over WebSocket to keep runner online.
|
||||||
func (j *JobConnection) Heartbeat() {
|
func (j *JobConnection) Heartbeat() {
|
||||||
|
j.writeMu.Lock()
|
||||||
|
defer j.writeMu.Unlock()
|
||||||
if j.conn == nil {
|
if j.conn == nil {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
j.writeMu.Lock()
|
|
||||||
defer j.writeMu.Unlock()
|
|
||||||
|
|
||||||
msg := map[string]interface{}{
|
msg := map[string]interface{}{
|
||||||
"type": "runner_heartbeat",
|
"type": "runner_heartbeat",
|
||||||
"timestamp": time.Now().Unix(),
|
"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.
|
// Close closes the WebSocket connection.
|
||||||
func (j *JobConnection) Close() {
|
func (j *JobConnection) Close() {
|
||||||
|
j.stopLoops()
|
||||||
|
|
||||||
|
j.writeMu.Lock()
|
||||||
|
defer j.writeMu.Unlock()
|
||||||
|
|
||||||
j.connMu.Lock()
|
j.connMu.Lock()
|
||||||
j.isConnected = false
|
j.isConnected = false
|
||||||
|
conn := j.conn
|
||||||
|
j.conn = nil
|
||||||
j.connMu.Unlock()
|
j.connMu.Unlock()
|
||||||
|
|
||||||
// Stop heartbeat goroutine
|
if conn != nil {
|
||||||
if j.stopHeartbeat != nil {
|
conn.Close()
|
||||||
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
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -211,13 +218,12 @@ func (j *JobConnection) IsConnected() bool {
|
|||||||
|
|
||||||
// Log sends a log entry to the manager.
|
// Log sends a log entry to the manager.
|
||||||
func (j *JobConnection) Log(taskID int64, level types.LogLevel, message string) {
|
func (j *JobConnection) Log(taskID int64, level types.LogLevel, message string) {
|
||||||
|
j.writeMu.Lock()
|
||||||
|
defer j.writeMu.Unlock()
|
||||||
if j.conn == nil {
|
if j.conn == nil {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
j.writeMu.Lock()
|
|
||||||
defer j.writeMu.Unlock()
|
|
||||||
|
|
||||||
msg := map[string]interface{}{
|
msg := map[string]interface{}{
|
||||||
"type": "log_entry",
|
"type": "log_entry",
|
||||||
"data": map[string]interface{}{
|
"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.
|
// Progress sends a progress update to the manager.
|
||||||
func (j *JobConnection) Progress(taskID int64, progress float64) {
|
func (j *JobConnection) Progress(taskID int64, progress float64) {
|
||||||
|
j.writeMu.Lock()
|
||||||
|
defer j.writeMu.Unlock()
|
||||||
if j.conn == nil {
|
if j.conn == nil {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
j.writeMu.Lock()
|
|
||||||
defer j.writeMu.Unlock()
|
|
||||||
|
|
||||||
msg := map[string]interface{}{
|
msg := map[string]interface{}{
|
||||||
"type": "progress",
|
"type": "progress",
|
||||||
"data": map[string]interface{}{
|
"data": map[string]interface{}{
|
||||||
@@ -272,13 +277,12 @@ func (j *JobConnection) Progress(taskID int64, progress float64) {
|
|||||||
|
|
||||||
// OutputUploaded notifies that an output file was uploaded.
|
// OutputUploaded notifies that an output file was uploaded.
|
||||||
func (j *JobConnection) OutputUploaded(taskID int64, fileName string) {
|
func (j *JobConnection) OutputUploaded(taskID int64, fileName string) {
|
||||||
|
j.writeMu.Lock()
|
||||||
|
defer j.writeMu.Unlock()
|
||||||
if j.conn == nil {
|
if j.conn == nil {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
j.writeMu.Lock()
|
|
||||||
defer j.writeMu.Unlock()
|
|
||||||
|
|
||||||
msg := map[string]interface{}{
|
msg := map[string]interface{}{
|
||||||
"type": "output_uploaded",
|
"type": "output_uploaded",
|
||||||
"data": map[string]interface{}{
|
"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
|
// freeRequeue asks the manager to requeue a failure without incrementing retry_count
|
||||||
// (used when this attempt newly armed GPU lockout).
|
// (used when this attempt newly armed GPU lockout).
|
||||||
func (j *JobConnection) Complete(taskID int64, success bool, errorMsg error, freeRequeue bool) {
|
func (j *JobConnection) Complete(taskID int64, success bool, errorMsg error, freeRequeue bool) {
|
||||||
|
j.writeMu.Lock()
|
||||||
|
defer j.writeMu.Unlock()
|
||||||
if j.conn == nil {
|
if j.conn == nil {
|
||||||
log.Printf("Cannot send task complete: WebSocket connection is nil")
|
log.Printf("Cannot send task complete: WebSocket connection is nil")
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
j.writeMu.Lock()
|
|
||||||
defer j.writeMu.Unlock()
|
|
||||||
|
|
||||||
data := map[string]interface{}{
|
data := map[string]interface{}{
|
||||||
"task_id": taskID,
|
"task_id": taskID,
|
||||||
"success": success,
|
"success": success,
|
||||||
|
|||||||
Reference in New Issue
Block a user