Also fix pre-existing -race failures in job connection Close().
This commit was merged in pull request #5.
This commit is contained in:
@@ -1,13 +1,16 @@
|
||||
name: PR Check
|
||||
name: CI
|
||||
on:
|
||||
- pull_request
|
||||
pull_request:
|
||||
push:
|
||||
branches:
|
||||
- master
|
||||
|
||||
jobs:
|
||||
check-and-test:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@main
|
||||
- uses: actions/setup-go@main
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/setup-go@v5
|
||||
with:
|
||||
go-version-file: 'go.mod'
|
||||
- uses: FedericoCarboni/setup-ffmpeg@v3
|
||||
|
||||
@@ -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() {
|
||||
}
|
||||
}
|
||||
|
||||
// Close closes the WebSocket connection.
|
||||
func (j *JobConnection) Close() {
|
||||
j.connMu.Lock()
|
||||
j.isConnected = false
|
||||
j.connMu.Unlock()
|
||||
|
||||
// Stop heartbeat goroutine
|
||||
// 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)
|
||||
j.stopHeartbeat = nil
|
||||
}
|
||||
|
||||
// Stop ping goroutine
|
||||
if j.stopPing != nil {
|
||||
close(j.stopPing)
|
||||
j.stopPing = nil
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
if j.conn != nil {
|
||||
j.conn.Close()
|
||||
// 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()
|
||||
|
||||
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,
|
||||
|
||||
Reference in New Issue
Block a user