202 lines
9.3 KiB
Go
202 lines
9.3 KiB
Go
package dispatcher
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"sync"
|
|
"time"
|
|
|
|
agentv1 "git.ipao.vip/rogee/go-sip/gen/agent/v1"
|
|
"google.golang.org/protobuf/proto"
|
|
)
|
|
|
|
var ErrRemoteResultUnknown = errors.New("remote Agent result is unknown; reconciliation required")
|
|
|
|
type AgentSession struct {
|
|
AgentID string
|
|
CellID string
|
|
BootID string
|
|
DispatcherEpoch string
|
|
SessionGeneration uint64
|
|
}
|
|
|
|
type DispatchResult struct {
|
|
Permit *agentv1.ExecutionPermit
|
|
Receipt *agentv1.OperationReceipt
|
|
State agentv1.ExecutionState
|
|
Unknown bool
|
|
Snapshot *agentv1.ExecutionSnapshot
|
|
}
|
|
|
|
// AgentCoordinator owns session binding and the no-retry-after-unknown rule.
|
|
// Quotas and reservations remain in Dispatcher SQLite; this type only turns a
|
|
// durable reservation into a fenced Agent permit and execution request.
|
|
type AgentCoordinator struct {
|
|
mu sync.Mutex
|
|
now func() time.Time
|
|
clients map[string]agentv1.AgentControlServiceClient
|
|
sessions map[string]AgentSession
|
|
}
|
|
|
|
func NewAgentCoordinator(now func() time.Time) *AgentCoordinator {
|
|
if now == nil {
|
|
now = time.Now
|
|
}
|
|
return &AgentCoordinator{now: now, clients: make(map[string]agentv1.AgentControlServiceClient), sessions: make(map[string]AgentSession)}
|
|
}
|
|
|
|
func (c *AgentCoordinator) Register(agentID string, client agentv1.AgentControlServiceClient) error {
|
|
if agentID == "" || client == nil {
|
|
return errors.New("agent ID and client are required")
|
|
}
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
c.clients[agentID] = client
|
|
return nil
|
|
}
|
|
|
|
// Probe performs the pre-activation R01 status read. The response is limited
|
|
// to deployment status; Activate must still bind the returned boot ID before
|
|
// any session-authorized operation is attempted.
|
|
func (c *AgentCoordinator) Probe(ctx context.Context, agentID, cellID string) (*agentv1.AgentStatus, error) {
|
|
if agentID == "" || cellID == "" {
|
|
return nil, errors.New("agent ID and cell ID are required")
|
|
}
|
|
client, err := c.client(agentID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
operationID := "status:" + agentID
|
|
response, err := client.GetAgentStatus(ctx, &agentv1.GetAgentStatusRequest{
|
|
Meta: &agentv1.RequestMeta{
|
|
ProtocolVersion: "agent.v1",
|
|
RequestId: operationID + ":request",
|
|
TraceId: operationID,
|
|
OperationId: operationID,
|
|
AgentId: agentID,
|
|
CellId: cellID,
|
|
},
|
|
Target: &agentv1.AgentBinding{AgentId: agentID, CellId: cellID},
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if response == nil || response.Status == nil {
|
|
return nil, errors.New("Agent status response is empty")
|
|
}
|
|
if response.Status.AgentId != agentID || response.Status.CellId != cellID || response.Status.BootId == "" {
|
|
return nil, errors.New("Agent status identity or boot ID is invalid")
|
|
}
|
|
return response.Status, nil
|
|
}
|
|
|
|
func (c *AgentCoordinator) Activate(ctx context.Context, agentID, cellID, bootID, epoch string, generation uint64) (AgentSession, error) {
|
|
client, err := c.client(agentID)
|
|
if err != nil {
|
|
return AgentSession{}, err
|
|
}
|
|
if cellID == "" || bootID == "" || epoch == "" {
|
|
return AgentSession{}, errors.New("cell, boot and dispatcher epoch are required")
|
|
}
|
|
operationID := fmt.Sprintf("activate:%s:%s", agentID, bootID)
|
|
meta := &agentv1.RequestMeta{ProtocolVersion: "agent.v1", RequestId: operationID + ":request", TraceId: operationID, OperationId: operationID, DispatcherEpoch: epoch, AgentId: agentID, CellId: cellID, BootId: bootID}
|
|
response, err := client.ActivateAgent(ctx, &agentv1.ActivateAgentRequest{Meta: meta, Binding: &agentv1.AgentBinding{AgentId: agentID, CellId: cellID, ExpectedBootId: bootID, DispatcherEpoch: epoch, SessionGeneration: generation}, ActivationOperationId: operationID})
|
|
if err != nil {
|
|
return AgentSession{}, err
|
|
}
|
|
if response.Session == nil || response.State != agentv1.ActivationState_ACTIVATION_STATE_ACTIVE {
|
|
return AgentSession{}, errors.New("Agent activation was not active")
|
|
}
|
|
session := AgentSession{AgentID: agentID, CellID: cellID, BootID: bootID, DispatcherEpoch: epoch, SessionGeneration: response.Session.SessionGeneration}
|
|
c.mu.Lock()
|
|
c.sessions[agentID] = session
|
|
c.mu.Unlock()
|
|
return session, nil
|
|
}
|
|
|
|
func (c *AgentCoordinator) ExecuteRaw(ctx context.Context, agentID string, binding *agentv1.ExecutionBinding, raw []byte, reservationID, configSHA256 string) (DispatchResult, error) {
|
|
if binding == nil || binding.ExecutionId == "" || reservationID == "" {
|
|
return DispatchResult{}, errors.New("execution binding and reservation are required")
|
|
}
|
|
client, session, err := c.clientAndSession(agentID)
|
|
if err != nil {
|
|
return DispatchResult{}, err
|
|
}
|
|
permitRequest := &agentv1.GetExecutionPermitRequest{Meta: c.meta(session, "permit:"+binding.ExecutionId, "permit:"+binding.ExecutionId), Binding: proto.Clone(binding).(*agentv1.ExecutionBinding), ResourceReservationId: reservationID, ExpectedTaskRevision: binding.TaskRevision, ConfigSha256: configSHA256}
|
|
permitResponse, err := client.GetExecutionPermit(ctx, permitRequest)
|
|
if err != nil {
|
|
return c.reconcileUnknown(ctx, client, session, binding, err)
|
|
}
|
|
if permitResponse != nil && permitResponse.Permit == nil && permitResponse.Receipt != nil && permitResponse.Receipt.Result == agentv1.ResultCode_RESULT_CODE_APPLIED {
|
|
// A durable receipt without the permit body is an incomplete read, not a
|
|
// reason to originate. Re-read the same reservation under a new read key.
|
|
permitRequest.Meta = c.meta(session, "permit-reconcile:"+binding.ExecutionId, "permit-reconcile:"+binding.ExecutionId)
|
|
permitResponse, err = client.GetExecutionPermit(ctx, permitRequest)
|
|
if err != nil {
|
|
return c.reconcileUnknown(ctx, client, session, binding, err)
|
|
}
|
|
}
|
|
if permitResponse == nil || permitResponse.Permit == nil || permitResponse.Receipt == nil || permitResponse.Receipt.Result == agentv1.ResultCode_RESULT_CODE_REJECTED || permitResponse.Receipt.Result == agentv1.ResultCode_RESULT_CODE_CONFLICT {
|
|
if permitResponse == nil {
|
|
return DispatchResult{}, fmt.Errorf("Agent did not grant execution permit: empty response")
|
|
}
|
|
return DispatchResult{}, fmt.Errorf("Agent did not grant execution permit: result=%s failure=%v permit=%v", permitResponse.Receipt.GetResult().String(), permitResponse.Receipt.GetFailure(), permitResponse.Permit)
|
|
}
|
|
executeMeta := c.meta(session, "execute:"+binding.ExecutionId, "execute:"+binding.ExecutionId)
|
|
executeResponse, err := client.Execute(ctx, &agentv1.ExecuteRequest{Meta: executeMeta, Binding: proto.Clone(binding).(*agentv1.ExecutionBinding), CallExecuteJson: append([]byte(nil), raw...), ConfigSha256: configSHA256, PermitId: permitResponse.Permit.PermitId})
|
|
if err != nil {
|
|
return c.reconcileUnknown(ctx, client, session, binding, err)
|
|
}
|
|
if executeResponse == nil || executeResponse.Receipt == nil {
|
|
return c.reconcileUnknown(ctx, client, session, binding, ErrRemoteResultUnknown)
|
|
}
|
|
return DispatchResult{Permit: permitResponse.Permit, Receipt: executeResponse.Receipt, State: executeResponse.State}, nil
|
|
}
|
|
|
|
func (c *AgentCoordinator) Control(ctx context.Context, agentID string, binding *agentv1.ExecutionBinding, action agentv1.ControlAction, policy agentv1.ActiveCallPolicy) (*agentv1.ApplyTaskControlResponse, error) {
|
|
client, session, err := c.clientAndSession(agentID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
meta := c.meta(session, fmt.Sprintf("control:%s:%d", binding.ExecutionId, binding.TaskRevision), fmt.Sprintf("control:%s:%d", binding.ExecutionId, binding.TaskRevision))
|
|
return client.ApplyTaskControl(ctx, &agentv1.ApplyTaskControlRequest{Meta: meta, Binding: proto.Clone(binding).(*agentv1.ExecutionBinding), Action: action, ActiveCallPolicy: policy, ExpectedTaskRevision: binding.TaskRevision})
|
|
}
|
|
|
|
func (c *AgentCoordinator) client(agentID string) (agentv1.AgentControlServiceClient, error) {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
client := c.clients[agentID]
|
|
if client == nil {
|
|
return nil, fmt.Errorf("Agent %q is not registered", agentID)
|
|
}
|
|
return client, nil
|
|
}
|
|
|
|
func (c *AgentCoordinator) clientAndSession(agentID string) (agentv1.AgentControlServiceClient, AgentSession, error) {
|
|
client, err := c.client(agentID)
|
|
if err != nil {
|
|
return nil, AgentSession{}, err
|
|
}
|
|
c.mu.Lock()
|
|
session, ok := c.sessions[agentID]
|
|
c.mu.Unlock()
|
|
if !ok {
|
|
return nil, AgentSession{}, fmt.Errorf("Agent %q is not activated", agentID)
|
|
}
|
|
return client, session, nil
|
|
}
|
|
|
|
func (c *AgentCoordinator) meta(session AgentSession, operationID, idempotencyKey string) *agentv1.RequestMeta {
|
|
return &agentv1.RequestMeta{ProtocolVersion: "agent.v1", RequestId: operationID + ":request", TraceId: operationID, OperationId: operationID, IdempotencyKey: idempotencyKey, DispatcherEpoch: session.DispatcherEpoch, AgentId: session.AgentID, CellId: session.CellID, BootId: session.BootID, SessionGeneration: session.SessionGeneration}
|
|
}
|
|
|
|
func (c *AgentCoordinator) reconcileUnknown(ctx context.Context, client agentv1.AgentControlServiceClient, session AgentSession, binding *agentv1.ExecutionBinding, cause error) (DispatchResult, error) {
|
|
queryMeta := c.meta(session, "query:"+binding.ExecutionId, "query:"+binding.ExecutionId)
|
|
response, err := client.QueryExecution(ctx, &agentv1.QueryExecutionRequest{Meta: queryMeta, Binding: proto.Clone(binding).(*agentv1.ExecutionBinding)})
|
|
if err == nil && response != nil && response.Snapshot != nil {
|
|
return DispatchResult{Unknown: true, Snapshot: response.Snapshot}, fmt.Errorf("%w: %v", ErrRemoteResultUnknown, cause)
|
|
}
|
|
return DispatchResult{Unknown: true}, fmt.Errorf("%w: %v", ErrRemoteResultUnknown, cause)
|
|
}
|