ben comment

Author: ianhodge

README.md

@@ -52,7 +52,7 @@ By default, the worker runs tasks without a concurrency limit. Use `--max-concur
 oz-agent-worker --api-key "wk-abc123" --worker-id "my-worker" --max-concurrent-tasks 4
 ```
 
-When the limit is reached, incoming tasks are queued and will execute once a slot becomes available. A value of `0` (the default) means unlimited.
+When the limit is reached, incoming tasks are rejected so the server can reassign them to other available workers. A value of `0` (the default) means unlimited.
 
 ## Environment Variables for Task Containers
 

go.mod

@@ -50,6 +50,7 @@ require (
 	go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.39.0 // indirect
 	go.opentelemetry.io/otel/metric v1.39.0 // indirect
 	go.opentelemetry.io/otel/trace v1.39.0 // indirect
+	golang.org/x/sync v0.19.0 // indirect
 	golang.org/x/sys v0.39.0 // indirect
 	golang.org/x/time v0.14.0 // indirect
 	google.golang.org/protobuf v1.36.10 // indirect

go.sum

@@ -216,6 +216,8 @@ golang.org/x/sync v0.0.0-20181108010431-42b317875d0f/go.mod h1:RxMgew5VJxzue5/jJ
 golang.org/x/sync v0.0.0-20181221193216-37e7f081c4d4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20201207232520-09787c993a3a/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.19.0 h1:vV+1eWNmZ5geRlYjzm2adRgW2/mcpevXNg50YZtPCE4=
+golang.org/x/sync v0.19.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI=
 golang.org/x/sys v0.0.0-20180905080454-ebe1bf3edb33/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20181116152217-5ac8a444bdc5/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=

internal/worker/worker.go

@@ -13,6 +13,7 @@ import (
 
 	"github.com/distribution/reference"
 	cliconfig "github.com/docker/cli/cli/config"
+	"golang.org/x/sync/semaphore"
 	"github.com/docker/docker/api/types/container"
 	"github.com/docker/docker/api/types/image"
 	"github.com/docker/docker/api/types/volume"
@@ -58,7 +59,7 @@ type Worker struct {
 	sendChan       chan []byte
 	activeTasks    map[string]context.CancelFunc
 	tasksMutex     sync.Mutex
-	taskSemaphore  chan struct{} // nil when unlimited; buffered channel used as counting semaphore
+	taskSemaphore  *semaphore.Weighted // nil when unlimited
 	dockerClient   *client.Client
 	platform       string // Docker daemon platform (e.g., "linux/amd64" or "linux/arm64")
 }
@@ -107,9 +108,9 @@ func New(ctx context.Context, config Config) (*Worker, error) {
 
 	log.Debugf(ctx, "Docker daemon is reachable, platform: %s", platform)
 
-	var taskSemaphore chan struct{}
+	var taskSemaphore *semaphore.Weighted
 	if config.MaxConcurrentTasks > 0 {
-		taskSemaphore = make(chan struct{}, config.MaxConcurrentTasks)
+		taskSemaphore = semaphore.NewWeighted(int64(config.MaxConcurrentTasks))
 		log.Infof(ctx, "Concurrency limit set to %d", config.MaxConcurrentTasks)
 	}
 
@@ -334,6 +335,18 @@ func (w *Worker) handleMessage(message []byte) {
 func (w *Worker) handleTaskAssignment(assignment *types.TaskAssignmentMessage) {
 	log.Infof(w.ctx, "Received task assignment: taskID=%s, title=%s", assignment.TaskID, assignment.Task.Title)
 
+	// If a concurrency limit is configured, try to acquire a slot without blocking.
+	// Failing fast lets the server reroute the task to another worker.
+	if w.taskSemaphore != nil {
+		if !w.taskSemaphore.TryAcquire(1) {
+			log.Warnf(w.ctx, "At max concurrency (%d), rejecting task: taskID=%s", w.config.MaxConcurrentTasks, assignment.TaskID)
+			if err := w.sendTaskFailed(assignment.TaskID, "worker at maximum concurrency"); err != nil {
+				log.Errorf(w.ctx, "Failed to send task failed message: %v", err)
+			}
+			return
+		}
+	}
+
 	// It's important to update the task state to claimed as the task lifecycle treats this as a dependency to advance to further states.
 	if err := w.sendTaskClaimed(assignment.TaskID); err != nil {
 		log.Errorf(w.ctx, "Failed to send task claimed message: %v", err)
@@ -349,29 +362,17 @@ func (w *Worker) handleTaskAssignment(assignment *types.TaskAssignmentMessage) {
 }
 
 func (w *Worker) executeTask(ctx context.Context, assignment *types.TaskAssignmentMessage) {
-	acquiredSlot := false
 	defer func() {
 		w.tasksMutex.Lock()
 		delete(w.activeTasks, assignment.TaskID)
 		w.tasksMutex.Unlock()
 
-		// Release the semaphore slot if we acquired one.
-		if w.taskSemaphore != nil && acquiredSlot {
-			<-w.taskSemaphore
+		// Release the semaphore slot if concurrency is limited.
+		if w.taskSemaphore != nil {
+			w.taskSemaphore.Release(1)
 		}
 	}()
 
-	// If a concurrency limit is configured, wait for an available slot.
-	if w.taskSemaphore != nil {
-		log.Infof(ctx, "Waiting for concurrency slot: taskID=%s", assignment.TaskID)
-		select {
-		case w.taskSemaphore <- struct{}{}:
-			acquiredSlot = true
-		case <-ctx.Done():
-			return
-		}
-	}
-
 	taskID := assignment.TaskID
 	log.Infof(ctx, "Starting task execution: taskID=%s, title=%s", taskID, assignment.Task.Title)