A self-hosted orchestration system crafted for power users. composia.xyz
Something went wrong. Try again.
5.5 kB · 150 lines
Go
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151package controller
import ( "context" "errors" "time"
"connectrpc.com/connect"
agentv1 "forgejo.alexma.top/alexma233/composia/gen/go/proto/composia/agent/v1" "forgejo.alexma.top/alexma233/composia/internal/platform/rpcutil" "forgejo.alexma.top/alexma233/composia/internal/platform/store" "google.golang.org/protobuf/types/known/timestamppb")
type agentTaskServer struct { db *store.DB taskQueue *taskQueueNotifier dockerQueries *dockerQueryBroker maxWait time.Duration retryInterval time.Duration}
func (server *agentTaskServer) PullNextTask(ctx context.Context, req *connect.Request[agentv1.PullNextTaskRequest]) (*connect.Response[agentv1.PullNextTaskResponse], error) { if req.Msg.GetNodeId() == "" { return nil, connect.NewError(connect.CodeInvalidArgument, errors.New("node_id is required")) } if req.Msg.GetExecutionProtocol() != 1 { return nil, connect.NewError(connect.CodeFailedPrecondition, errors.New("agent must support task execution protocol 1")) } authenticatedNodeID, ok := rpcutil.BearerSubject(ctx) if !ok || authenticatedNodeID != req.Msg.GetNodeId() { return nil, connect.NewError(connect.CodeUnauthenticated, errors.New("node_id does not match bearer token")) }
waitCh := server.taskQueue.Subscribe() defer server.taskQueue.Unsubscribe(waitCh) deadline := time.Now().Add(server.longPollMaxWait())
for { now := time.Now().UTC() record, err := server.db.ClaimNextPendingTaskForNode(ctx, req.Msg.GetNodeId(), now, now.Add(taskExecutionLease)) if err == nil { params, err := taskParams(record.ParamsJSON) if err != nil { return nil, connect.NewError(connect.CodeInternal, err) } logControllerAssignedTask(record) response := &agentv1.PullNextTaskResponse{ HasTask: true, Task: &agentv1.AgentTask{ TaskId: record.TaskID, Type: protoAgentTaskType(record.Type), ServiceName: record.ServiceName, NodeId: record.NodeID, RepoRevision: record.RepoRevision, ServiceDir: params.ServiceDir, DataNames: params.DataNames, ParamsJson: record.ParamsJSON, ExecutionId: record.ExecutionID, LeaseExpiresAt: timestamppb.New(*record.LeaseExpiresAt), }, } return connect.NewResponse(response), nil } if !errors.Is(err, store.ErrNoPendingTask) { return nil, connect.NewError(connect.CodeInternal, err) }
remaining := time.Until(deadline) if remaining <= 0 { return connect.NewResponse(&agentv1.PullNextTaskResponse{HasTask: false}), nil } waitFor := minDuration(remaining, server.longPollRetryInterval()) timer := time.NewTimer(waitFor) select { case <-ctx.Done(): timer.Stop() return nil, ctx.Err() case <-waitCh: timer.Stop() case <-timer.C: } }}
func (server *agentTaskServer) AcknowledgeTask(ctx context.Context, req *connect.Request[agentv1.AcknowledgeTaskRequest]) (*connect.Response[agentv1.AcknowledgeTaskResponse], error) { nodeID, err := validateTaskExecutionRequest(ctx, req.Msg.GetTaskId(), req.Msg.GetExecutionId()) if err != nil { return nil, err } detail, err := server.db.GetTask(ctx, req.Msg.GetTaskId()) if err != nil { return nil, connect.NewError(connect.CodeNotFound, err) } if detail.Record.NodeID != nodeID { return nil, connect.NewError(connect.CodePermissionDenied, errors.New("task does not belong to authenticated node")) } now := time.Now().UTC() expiresAt := now.Add(taskExecutionLease) if err := server.db.AcknowledgeTaskExecution(ctx, req.Msg.GetTaskId(), req.Msg.GetExecutionId(), now, expiresAt); err != nil { return nil, connect.NewError(connect.CodeFailedPrecondition, err) } return connect.NewResponse(&agentv1.AcknowledgeTaskResponse{LeaseExpiresAt: timestamppb.New(expiresAt)}), nil}
func (server *agentTaskServer) RenewTaskLease(ctx context.Context, req *connect.Request[agentv1.RenewTaskLeaseRequest]) (*connect.Response[agentv1.RenewTaskLeaseResponse], error) { nodeID, err := validateTaskExecutionRequest(ctx, req.Msg.GetTaskId(), req.Msg.GetExecutionId()) if err != nil { return nil, err } detail, err := server.db.GetTask(ctx, req.Msg.GetTaskId()) if err != nil { return nil, connect.NewError(connect.CodeNotFound, err) } if detail.Record.NodeID != nodeID { return nil, connect.NewError(connect.CodePermissionDenied, errors.New("task does not belong to authenticated node")) } now := time.Now().UTC() expiresAt := now.Add(taskExecutionLease) if err := server.db.RenewTaskExecution(ctx, req.Msg.GetTaskId(), req.Msg.GetExecutionId(), now, expiresAt); err != nil { return nil, connect.NewError(connect.CodeFailedPrecondition, err) } return connect.NewResponse(&agentv1.RenewTaskLeaseResponse{LeaseExpiresAt: timestamppb.New(expiresAt)}), nil}
func validateTaskExecutionRequest(ctx context.Context, taskID, executionID string) (string, error) { if taskID == "" || executionID == "" { return "", connect.NewError(connect.CodeInvalidArgument, errors.New("task_id and execution_id are required")) } nodeID, ok := rpcutil.BearerSubject(ctx) if !ok || nodeID == "" { return "", connect.NewError(connect.CodeUnauthenticated, errors.New("missing bearer subject")) } return nodeID, nil}
func (server *agentTaskServer) longPollMaxWait() time.Duration { if server.maxWait > 0 { return server.maxWait } return pullNextTaskMaxWait}
func (server *agentTaskServer) longPollRetryInterval() time.Duration { if server.retryInterval > 0 { return server.retryInterval } return pullNextTaskRetryWait}