first commit
This commit is contained in:
npc0-hue
@@ -0,0 +1,384 @@
|
||||
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,
|
||||
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)
|
||||
cancelled := make(chan protocol.RunJobCancelPollResponse, 1)
|
||||
go func() {
|
||||
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()
|
||||
cancelled <- cancelPoll
|
||||
return
|
||||
}
|
||||
cancelled <- protocol.RunJobCancelPollResponse{Accepted: true}
|
||||
}()
|
||||
execution := worker.executor.ExecuteContext(jobCtx, assignment)
|
||||
cancel()
|
||||
select {
|
||||
case poll := <-cancelled:
|
||||
if poll.HasCancel && execution.State == lifecycleResultStateSucceeded {
|
||||
execution = LifecycleExecutionResult{
|
||||
State: lifecycleResultStateCancelled,
|
||||
Progress: protocol.RunJobProgressReport{Percent: 100, Message: "cancelled by platform"},
|
||||
Message: "cancelled by platform",
|
||||
ErrorCode: "lifecycle_cancelled",
|
||||
}
|
||||
}
|
||||
default:
|
||||
}
|
||||
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) 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
|
||||
}
|
||||
Reference in New Issue
Block a user