134 lines
3.4 KiB
Go
134 lines
3.4 KiB
Go
package agent
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"sync"
|
|
)
|
|
|
|
var (
|
|
ErrTaskAdmissionClosed = errors.New("Agent task admission is paused")
|
|
ErrTaskStopped = errors.New("Agent task is stopped")
|
|
)
|
|
|
|
type TaskIdentity struct {
|
|
DispatcherID string
|
|
TenantID int64
|
|
TaskID string
|
|
}
|
|
|
|
type activeTaskCall struct {
|
|
cancel context.CancelFunc
|
|
done chan struct{}
|
|
}
|
|
|
|
type taskCalls struct {
|
|
state string
|
|
active map[string]*activeTaskCall
|
|
}
|
|
|
|
// TaskCalls keeps an in-process execution barrier for task-level controls.
|
|
// Dispatcher owns durable task state; an Agent restart never authorizes a new
|
|
// call without a fresh Dispatcher instruction. The caller must authenticate
|
|
// the active Dispatcher session before invoking Register or Apply.
|
|
type TaskCalls struct {
|
|
mu sync.Mutex
|
|
tasks map[TaskIdentity]*taskCalls
|
|
}
|
|
|
|
func (r *TaskCalls) Register(task TaskIdentity, callID string, cancel context.CancelFunc) (func(), error) {
|
|
if r == nil || !validTaskIdentity(task) || callID == "" || cancel == nil {
|
|
return nil, errors.New("Agent call requires a task identity, call ID and cancellation")
|
|
}
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
entry := r.entry(task)
|
|
switch entry.state {
|
|
case "stopped":
|
|
return nil, ErrTaskStopped
|
|
case "paused":
|
|
return nil, ErrTaskAdmissionClosed
|
|
}
|
|
if _, exists := entry.active[callID]; exists {
|
|
return nil, errors.New("Agent call is already active")
|
|
}
|
|
call := &activeTaskCall{cancel: cancel, done: make(chan struct{})}
|
|
entry.active[callID] = call
|
|
var once sync.Once
|
|
return func() {
|
|
once.Do(func() {
|
|
r.mu.Lock()
|
|
delete(entry.active, callID)
|
|
close(call.done)
|
|
r.mu.Unlock()
|
|
})
|
|
}, nil
|
|
}
|
|
|
|
// Apply closes admission before requesting hangup or waiting for drain. Its
|
|
// success means all calls active when the control arrived have finished; a
|
|
// timed-out wait leaves the pause/stop barrier in place for re-delivery.
|
|
func (r *TaskCalls) Apply(ctx context.Context, task TaskIdentity, action, policy string) error {
|
|
if r == nil || ctx == nil || !validTaskIdentity(task) {
|
|
return errors.New("Agent task control requires a context and task identity")
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return err
|
|
}
|
|
if (action != "pause" && action != "stop" && action != "resume") ||
|
|
(action == "resume" && policy != "") ||
|
|
(action != "resume" && policy != "hangup" && policy != "drain") {
|
|
return errors.New("Agent task control action or active-call policy is invalid")
|
|
}
|
|
r.mu.Lock()
|
|
entry := r.entry(task)
|
|
if entry.state == "stopped" && action != "stop" {
|
|
r.mu.Unlock()
|
|
return ErrTaskStopped
|
|
}
|
|
if action == "resume" {
|
|
entry.state = ""
|
|
r.mu.Unlock()
|
|
return nil
|
|
}
|
|
entry.state = "paused"
|
|
if action == "stop" {
|
|
entry.state = "stopped"
|
|
}
|
|
active := make([]*activeTaskCall, 0, len(entry.active))
|
|
for _, call := range entry.active {
|
|
active = append(active, call)
|
|
}
|
|
r.mu.Unlock()
|
|
if policy == "hangup" {
|
|
for _, call := range active {
|
|
call.cancel()
|
|
}
|
|
}
|
|
for _, call := range active {
|
|
select {
|
|
case <-call.done:
|
|
case <-ctx.Done():
|
|
return fmt.Errorf("Agent task control waiting for active calls: %w", ctx.Err())
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (r *TaskCalls) entry(task TaskIdentity) *taskCalls {
|
|
if r.tasks == nil {
|
|
r.tasks = make(map[TaskIdentity]*taskCalls)
|
|
}
|
|
entry := r.tasks[task]
|
|
if entry == nil {
|
|
entry = &taskCalls{active: make(map[string]*activeTaskCall)}
|
|
r.tasks[task] = entry
|
|
}
|
|
return entry
|
|
}
|
|
|
|
func validTaskIdentity(task TaskIdentity) bool {
|
|
return task.DispatcherID != "" && task.TenantID > 0 && task.TaskID != ""
|
|
}
|