393 lines
12 KiB
Go
393 lines
12 KiB
Go
package runtime
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"runtime"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"browser.local/run/config"
|
|
"browser.local/run/protocol"
|
|
"browser.local/run/spool"
|
|
)
|
|
|
|
type WorkerClient interface {
|
|
Hello(context.Context, protocol.RunHelloRequest) (protocol.RunHelloResponse, error)
|
|
Heartbeat(context.Context, protocol.RunHeartbeatRequest) (protocol.RunHeartbeatResponse, error)
|
|
ClaimJob(context.Context, protocol.RunJobClaimRequest) (protocol.RunJobClaimResponse, error)
|
|
AckJob(context.Context, protocol.RunJobAckRequest) (protocol.RunJobAckResponse, error)
|
|
UpdateJobProgress(context.Context, protocol.RunJobProgressRequest) (protocol.RunJobProgressResponse, error)
|
|
CompleteJob(context.Context, protocol.RunJobResultRequest) (protocol.RunJobResultResponse, error)
|
|
PollJobCancel(context.Context, protocol.RunJobCancelPollRequest) (protocol.RunJobCancelPollResponse, error)
|
|
ReconcileJobs(context.Context, protocol.RunJobReconcileRequest) (protocol.RunJobReconcileResponse, error)
|
|
}
|
|
|
|
type Worker struct {
|
|
cfg config.Config
|
|
client WorkerClient
|
|
executor LifecycleExecutor
|
|
state WorkerState
|
|
journal *JobJournal
|
|
}
|
|
|
|
type WorkerState struct {
|
|
RunEndpointID string
|
|
SessionToken string
|
|
Capabilities []string
|
|
Capacity protocol.RunCapacityReport
|
|
LastHeartbeat time.Time
|
|
Sequence uint64
|
|
}
|
|
|
|
func NewWorker(cfg config.Config, client WorkerClient, options ...LifecycleExecutorOption) (*Worker, error) {
|
|
if client == nil {
|
|
return nil, fmt.Errorf("worker client is required")
|
|
}
|
|
if cfg.RunEndpointID == "" {
|
|
cfg.RunEndpointID = config.DefaultEndpointID
|
|
}
|
|
if cfg.DisplayName == "" {
|
|
cfg.DisplayName = config.DefaultDisplayName
|
|
}
|
|
if cfg.Version == "" {
|
|
cfg.Version = config.DefaultVersion
|
|
}
|
|
if cfg.MaxJobs <= 0 {
|
|
cfg.MaxJobs = 1
|
|
}
|
|
executorOptions := append([]LifecycleExecutorOption{
|
|
WithLifecycleWorkspaceRoot(cfg.WorkspaceRoot),
|
|
}, options...)
|
|
return &Worker{
|
|
cfg: cfg,
|
|
client: client,
|
|
executor: NewLifecycleExecutor(executorOptions...),
|
|
state: WorkerState{
|
|
RunEndpointID: cfg.RunEndpointID,
|
|
Capabilities: SupportedRunCapabilities(),
|
|
Capacity: protocol.RunCapacityReport{MaxJobs: cfg.MaxJobs},
|
|
},
|
|
journal: NewJobJournal(),
|
|
}, nil
|
|
}
|
|
|
|
func (worker *Worker) Register(ctx context.Context) error {
|
|
response, err := worker.client.Hello(ctx, protocol.RunHelloRequest{
|
|
RegistrationToken: worker.cfg.RegistrationToken,
|
|
RunEndpointID: worker.cfg.RunEndpointID,
|
|
ServerInstanceID: worker.cfg.ServerInstanceID,
|
|
PluginID: worker.cfg.PluginID,
|
|
ComponentKind: worker.cfg.ComponentKind,
|
|
ComponentKey: worker.cfg.ComponentKey,
|
|
KeyGeneration: worker.cfg.KeyGeneration,
|
|
DisplayName: worker.cfg.DisplayName,
|
|
Version: worker.cfg.Version,
|
|
Status: "online",
|
|
Platform: runtime.GOOS,
|
|
CapabilityReport: protocol.RunCapabilityReport{
|
|
Capabilities: worker.state.Capabilities,
|
|
Fingerprint: capabilityFingerprint(worker.state.Capabilities),
|
|
},
|
|
Capacity: worker.capacityReport(),
|
|
})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if !response.Accepted || response.SessionToken == "" {
|
|
return fmt.Errorf("run hello was not accepted")
|
|
}
|
|
worker.state.SessionToken = response.SessionToken
|
|
if sink, ok := worker.executor.logSink.(*SpoolLogSink); ok {
|
|
sink.RunEndpointID = worker.state.RunEndpointID
|
|
sink.SessionToken = worker.state.SessionToken
|
|
}
|
|
if hook, ok := worker.executor.artifactHook.(*QueueArtifactHook); ok {
|
|
hook.RunEndpointID = worker.state.RunEndpointID
|
|
hook.SessionToken = worker.state.SessionToken
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (worker *Worker) HeartbeatOnce(ctx context.Context) error {
|
|
if worker.state.SessionToken == "" {
|
|
return fmt.Errorf("worker is not registered")
|
|
}
|
|
response, err := worker.client.Heartbeat(ctx, protocol.RunHeartbeatRequest{
|
|
RunEndpointID: worker.state.RunEndpointID,
|
|
SessionToken: worker.state.SessionToken,
|
|
Version: worker.cfg.Version,
|
|
Status: "online",
|
|
CapabilityFingerprint: capabilityFingerprint(worker.state.Capabilities),
|
|
Capacity: worker.capacityReport(),
|
|
})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if !response.Accepted {
|
|
return fmt.Errorf("heartbeat was not accepted")
|
|
}
|
|
worker.state.LastHeartbeat = response.ServerTime
|
|
return nil
|
|
}
|
|
|
|
func (worker *Worker) ClaimAndRunOnce(ctx context.Context) (bool, error) {
|
|
if worker.state.SessionToken == "" {
|
|
return false, fmt.Errorf("worker is not registered")
|
|
}
|
|
claim, err := worker.client.ClaimJob(ctx, protocol.RunJobClaimRequest{
|
|
RunEndpointID: worker.state.RunEndpointID,
|
|
SessionToken: worker.state.SessionToken,
|
|
Capabilities: worker.state.Capabilities,
|
|
Capacity: worker.capacityReport(),
|
|
})
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
if !claim.Accepted || !claim.HasJob || claim.Job == nil {
|
|
return false, nil
|
|
}
|
|
assignment := *claim.Job
|
|
worker.journal.MarkActive(assignment)
|
|
ack, err := worker.client.AckJob(ctx, protocol.RunJobAckRequest{
|
|
RunEndpointID: worker.state.RunEndpointID,
|
|
SessionToken: worker.state.SessionToken,
|
|
JobID: assignment.JobID,
|
|
LeaseToken: assignment.LeaseToken,
|
|
Attempt: assignment.Attempt,
|
|
Message: "job accepted by run worker",
|
|
})
|
|
if err != nil {
|
|
return true, err
|
|
}
|
|
assignment = ack.Job
|
|
worker.journal.MarkActive(assignment)
|
|
worker.state.Sequence++
|
|
if _, err := worker.client.UpdateJobProgress(ctx, protocol.RunJobProgressRequest{
|
|
RunEndpointID: worker.state.RunEndpointID,
|
|
SessionToken: worker.state.SessionToken,
|
|
JobID: assignment.JobID,
|
|
LeaseToken: assignment.LeaseToken,
|
|
Attempt: assignment.Attempt,
|
|
Progress: protocol.RunJobProgressReport{Percent: 10, Message: "lifecycle execution started"},
|
|
Sequence: worker.state.Sequence,
|
|
}); err != nil {
|
|
return true, err
|
|
}
|
|
jobCtx, cancel := context.WithCancel(ctx)
|
|
cancelPoll, pollErr := worker.client.PollJobCancel(ctx, protocol.RunJobCancelPollRequest{
|
|
RunEndpointID: worker.state.RunEndpointID,
|
|
SessionToken: worker.state.SessionToken,
|
|
JobID: assignment.JobID,
|
|
LeaseToken: assignment.LeaseToken,
|
|
})
|
|
if pollErr == nil && cancelPoll.HasCancel {
|
|
cancel()
|
|
}
|
|
execution := worker.executeAssignment(jobCtx, assignment)
|
|
cancel()
|
|
if pollErr == nil && cancelPoll.HasCancel && execution.State == lifecycleResultStateSucceeded {
|
|
execution = LifecycleExecutionResult{
|
|
State: lifecycleResultStateCancelled,
|
|
Progress: protocol.RunJobProgressReport{Percent: 100, Message: "cancelled by platform"},
|
|
Message: "cancelled by platform",
|
|
ErrorCode: "lifecycle_cancelled",
|
|
}
|
|
}
|
|
if _, err := worker.client.CompleteJob(ctx, LifecycleResultRequest(assignment, worker.state.SessionToken, execution)); err != nil {
|
|
return true, err
|
|
}
|
|
worker.journal.MarkTerminal(assignment.JobID)
|
|
return true, nil
|
|
}
|
|
|
|
func (worker *Worker) executeAssignment(ctx context.Context, assignment protocol.RunJobAssignment) LifecycleExecutionResult {
|
|
if isSupportedLifecycleCapability(assignment.Capability) {
|
|
return worker.executor.ExecuteContext(ctx, assignment)
|
|
}
|
|
if isSupportedDistributionCapability(assignment.Capability) {
|
|
return ExecuteDistributionJob(ctx, assignment)
|
|
}
|
|
if isSupportedRemoteCapability(assignment.Capability) {
|
|
return ExecuteRemoteAccessJob(ctx, assignment)
|
|
}
|
|
return lifecycleFailure("unsupported_run_capability", "unsupported run capability")
|
|
}
|
|
|
|
func (worker *Worker) ReconcileOnce(ctx context.Context) error {
|
|
if worker.state.SessionToken == "" {
|
|
return fmt.Errorf("worker is not registered")
|
|
}
|
|
response, err := worker.client.ReconcileJobs(ctx, protocol.RunJobReconcileRequest{
|
|
RunEndpointID: worker.state.RunEndpointID,
|
|
SessionToken: worker.state.SessionToken,
|
|
ActiveJobIDs: worker.journal.ActiveJobIDs(),
|
|
})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
for _, job := range response.ActiveJobs {
|
|
worker.journal.MarkActive(job)
|
|
}
|
|
for _, unknown := range response.UnknownJobIDs {
|
|
worker.journal.MarkTerminal(unknown)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (worker *Worker) Run(ctx context.Context) error {
|
|
if err := worker.Register(ctx); err != nil {
|
|
return err
|
|
}
|
|
heartbeatInterval := durationOrDefault(worker.cfg.HeartbeatInterval, 15*time.Second)
|
|
jobInterval := durationOrDefault(worker.cfg.PollInterval, 2*time.Second)
|
|
heartbeatTicker := time.NewTicker(heartbeatInterval)
|
|
jobTicker := time.NewTicker(jobInterval)
|
|
defer heartbeatTicker.Stop()
|
|
defer jobTicker.Stop()
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
case <-heartbeatTicker.C:
|
|
if err := worker.HeartbeatOnce(ctx); err != nil {
|
|
heartbeatTicker.Reset(boundedRetryBackoff(worker.cfg.RetryBackoff))
|
|
continue
|
|
}
|
|
heartbeatTicker.Reset(heartbeatInterval)
|
|
case <-jobTicker.C:
|
|
if _, err := worker.ClaimAndRunOnce(ctx); err != nil {
|
|
jobTicker.Reset(boundedRetryBackoff(worker.cfg.RetryBackoff))
|
|
continue
|
|
}
|
|
jobTicker.Reset(jobInterval)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (worker *Worker) capacityReport() protocol.RunCapacityReport {
|
|
return protocol.RunCapacityReport{
|
|
MaxJobs: worker.state.Capacity.MaxJobs,
|
|
RunningJobs: worker.journal.ActiveCount(),
|
|
QueuedJobs: 0,
|
|
Summary: "worker control active; job capacity reported separately",
|
|
}
|
|
}
|
|
|
|
func (worker *Worker) State() WorkerState {
|
|
state := worker.state
|
|
state.Capabilities = append([]string(nil), state.Capabilities...)
|
|
return state
|
|
}
|
|
|
|
type JobJournal struct {
|
|
mu sync.Mutex
|
|
active map[string]protocol.RunJobAssignment
|
|
}
|
|
|
|
func NewJobJournal() *JobJournal {
|
|
return &JobJournal{active: map[string]protocol.RunJobAssignment{}}
|
|
}
|
|
|
|
func (journal *JobJournal) MarkActive(job protocol.RunJobAssignment) {
|
|
journal.mu.Lock()
|
|
defer journal.mu.Unlock()
|
|
journal.active[job.JobID] = job
|
|
}
|
|
|
|
func (journal *JobJournal) MarkTerminal(jobID string) {
|
|
journal.mu.Lock()
|
|
defer journal.mu.Unlock()
|
|
delete(journal.active, jobID)
|
|
}
|
|
|
|
func (journal *JobJournal) ActiveJobIDs() []string {
|
|
journal.mu.Lock()
|
|
defer journal.mu.Unlock()
|
|
ids := make([]string, 0, len(journal.active))
|
|
for id := range journal.active {
|
|
ids = append(ids, id)
|
|
}
|
|
return ids
|
|
}
|
|
|
|
func (journal *JobJournal) ActiveCount() int {
|
|
journal.mu.Lock()
|
|
defer journal.mu.Unlock()
|
|
return len(journal.active)
|
|
}
|
|
|
|
type SpoolLogSink struct {
|
|
RunEndpointID string
|
|
SessionToken string
|
|
Spool spool.LogSpool
|
|
seq uint64
|
|
}
|
|
|
|
func (sink *SpoolLogSink) Append(_ context.Context, assignment protocol.RunJobAssignment, stream string, line string) error {
|
|
sink.seq++
|
|
entry := protocol.LogEntry{Seq: sink.seq, Timestamp: time.Now().UTC(), Level: "info", Line: RedactText(line), Redacted: line != RedactText(line)}
|
|
logStreamID := fmt.Sprintf("job.%s.%s", assignment.JobID, stream)
|
|
return sink.Spool.Enqueue(protocol.LogBatchIngestRequest{
|
|
RunEndpointID: sink.RunEndpointID,
|
|
SessionToken: sink.SessionToken,
|
|
LogStreamID: logStreamID,
|
|
ServerInstanceID: assignment.ServerInstanceID,
|
|
StreamKey: stream,
|
|
Source: "process",
|
|
FirstSeq: sink.seq,
|
|
LastSeq: sink.seq,
|
|
Checksum: checksumForText(entry.Line),
|
|
Entries: []protocol.LogEntry{entry},
|
|
})
|
|
}
|
|
|
|
type QueueArtifactHook struct {
|
|
RunEndpointID string
|
|
SessionToken string
|
|
Queue spool.ArtifactQueue
|
|
}
|
|
|
|
func (hook QueueArtifactHook) QueueLifecycleResult(_ context.Context, assignment protocol.RunJobAssignment, result ProcessResult) (string, error) {
|
|
ref := fmt.Sprintf("artifact://jobs/%s/lifecycle-result", assignment.JobID)
|
|
payload := []byte(RedactText(result.Stdout + result.Stderr))
|
|
if len(payload) == 0 {
|
|
payload = []byte("lifecycle result metadata")
|
|
}
|
|
artifactID := "artifact-" + assignment.JobID + "-lifecycle"
|
|
if err := hook.Queue.Enqueue(protocol.ArtifactChunkUploadRequest{
|
|
RunEndpointID: hook.RunEndpointID,
|
|
SessionToken: hook.SessionToken,
|
|
TransferID: "transfer-" + assignment.JobID,
|
|
ArtifactID: artifactID,
|
|
ChunkIndex: 0,
|
|
Offset: 0,
|
|
SizeBytes: len(payload),
|
|
Checksum: checksumForText(string(payload)),
|
|
Payload: payload,
|
|
}); err != nil {
|
|
return "", err
|
|
}
|
|
return ref, nil
|
|
}
|
|
|
|
func capabilityFingerprint(capabilities []string) string {
|
|
return checksumForText(strings.Join(capabilities, ","))
|
|
}
|
|
|
|
func durationOrDefault(value time.Duration, fallback time.Duration) time.Duration {
|
|
if value <= 0 {
|
|
return fallback
|
|
}
|
|
return value
|
|
}
|
|
|
|
func boundedRetryBackoff(value time.Duration) time.Duration {
|
|
value = durationOrDefault(value, time.Second)
|
|
if value > 30*time.Second {
|
|
return 30 * time.Second
|
|
}
|
|
return value
|
|
}
|