From 79966763c6d37a29e3adbd03e92c0f874170a18a Mon Sep 17 00:00:00 2001 From: Jason Song Date: Tue, 11 Oct 2022 18:14:23 +0800 Subject: [PATCH] feat: distribute tasks to runner --- models/bots/run.go | 52 +++++++++++++++++++++++++ models/bots/run_job.go | 47 +++++++++++++++++++++++ models/bots/task.go | 63 +++++++++++++++++++++---------- models/bots/task_step.go | 3 +- routers/api/bots/runner/runner.go | 36 ++++++++++++------ routers/api/bots/runner/unary.go | 21 +++++++---- 6 files changed, 182 insertions(+), 40 deletions(-) diff --git a/models/bots/run.go b/models/bots/run.go index ffece5f3e7..d8eac83297 100644 --- a/models/bots/run.go +++ b/models/bots/run.go @@ -5,6 +5,7 @@ package bots import ( + "context" "fmt" "hash/fnv" @@ -50,6 +51,34 @@ func (Run) TableName() string { return "bots_run" } +// LoadAttributes load Repo TriggerUser if not loaded +func (r *Run) LoadAttributes(ctx context.Context) error { + if r == nil { + return nil + } + + if r.Repo == nil { + repo, err := repo_model.GetRepositoryByIDCtx(ctx, r.RepoID) + if err != nil { + return err + } + r.Repo = repo + } + if err := r.Repo.LoadAttributes(ctx); err != nil { + return err + } + + if r.TriggerUser == nil { + u, err := user_model.GetUserByIDCtx(ctx, r.TriggerUserID) + if err != nil { + return err + } + r.TriggerUser = u + } + + return nil +} + // InsertRun inserts a bot run func InsertRun(run *Run, jobs []*jobparser.SingleWorkflow) error { var groupId int64 @@ -103,6 +132,29 @@ func InsertRun(run *Run, jobs []*jobparser.SingleWorkflow) error { return nil } +// ErrRunNotExist represents an error for bot run not exist +type ErrRunNotExist struct { + ID int64 +} + +func (err ErrRunNotExist) Error() string { + return fmt.Sprintf("run [%d] is not exist", err.ID) +} + +func GetRunByID(ctx context.Context, id int64) (*Run, error) { + var run Run + has, err := db.GetEngine(ctx).Where("id=?", id).Get(&run) + if err != nil { + return nil, err + } else if !has { + return nil, ErrRunNotExist{ + ID: id, + } + } + + return &run, nil +} + type RunIndex db.ResourceIndex func (RunIndex) TableName() string { diff --git a/models/bots/run_job.go b/models/bots/run_job.go index 3a477eaa27..6e55088135 100644 --- a/models/bots/run_job.go +++ b/models/bots/run_job.go @@ -5,6 +5,9 @@ package bots import ( + "context" + "fmt" + "code.gitea.io/gitea/core" "code.gitea.io/gitea/models/db" "code.gitea.io/gitea/modules/timeutil" @@ -14,6 +17,7 @@ import ( type RunJob struct { ID int64 RunID int64 + Run *Run `xorm:"-"` Name string Ready bool // ready to be executed Attempt int64 @@ -36,3 +40,46 @@ func init() { func (RunJob) TableName() string { return "bots_run_job" } + +// LoadAttributes load Run if not loaded +func (job *RunJob) LoadAttributes(ctx context.Context) error { + if job == nil { + return nil + } + + if job.Run == nil { + run, err := GetRunByID(ctx, job.RunID) + if err != nil { + return err + } + job.Run = run + } + if err := job.Run.LoadAttributes(ctx); err != nil { + return err + } + + return nil +} + +// ErrRunJobNotExist represents an error for bot run job not exist +type ErrRunJobNotExist struct { + ID int64 +} + +func (err ErrRunJobNotExist) Error() string { + return fmt.Sprintf("run job [%d] is not exist", err.ID) +} + +func GetRunJobByID(ctx context.Context, id int64) (*RunJob, error) { + var job RunJob + has, err := db.GetEngine(ctx).Where("id=?", id).Get(&job) + if err != nil { + return nil, err + } else if !has { + return nil, ErrRunNotExist{ + ID: id, + } + } + + return &job, nil +} diff --git a/models/bots/task.go b/models/bots/task.go index efad43168e..fe705497d4 100644 --- a/models/bots/task.go +++ b/models/bots/task.go @@ -5,20 +5,24 @@ package bots import ( + "context" + "code.gitea.io/gitea/core" "code.gitea.io/gitea/models/db" "code.gitea.io/gitea/modules/timeutil" + runnerv1 "gitea.com/gitea/proto-go/runner/v1" ) // Task represents a distribution of job type Task struct { ID int64 JobID int64 + Job *RunJob `xorm:"-"` Attempt int64 RunnerID int64 `xorm:"index"` LogToFile bool // read log from database or from storage LogUrl string // url of the log file in storage - Result int64 // TODO: use runnerv1.Result + Result runnerv1.Result Started timeutil.TimeStamp Stopped timeutil.TimeStamp Created timeutil.TimeStamp `xorm:"created"` @@ -33,27 +37,49 @@ func (Task) TableName() string { return "bots_task" } -func CreateTask(runner *Runner) (task *Task, job *RunJob, run *Run, ok bool, err error) { +// LoadAttributes load Job if not loaded +func (task *Task) LoadAttributes(ctx context.Context) error { + if task == nil { + return nil + } + + if task.Job == nil { + job, err := GetRunJobByID(ctx, task.JobID) + if err != nil { + return err + } + task.Job = job + } + if err := task.Job.LoadAttributes(ctx); err != nil { + return err + } + + return nil +} + +func CreateTask(runner *Runner) (*Task, bool, error) { ctx, commiter, err := db.TxContext() if err != nil { - return + return nil, false, err } defer commiter.Close() var jobs []*RunJob - if err = db.GetEngine(ctx).Where("task_id = 0 AND ready = true").OrderBy("id").Find(jobs); err != nil { - return + if err := db.GetEngine(ctx).Where("task_id=? AND ready=?", 0, true).OrderBy("id").Find(&jobs); err != nil { + return nil, false, err } + // TODO: a more efficient way to filter labels + var job *RunJob labels := append(runner.AgentLabels, runner.CustomLabels...) for _, v := range jobs { - if isSubset(v.RunsOn, labels) { + if isSubset(labels, v.RunsOn) { job = v break } } if job == nil { - return + return nil, false, nil } now := timeutil.TimeStampNow() @@ -61,33 +87,32 @@ func CreateTask(runner *Runner) (task *Task, job *RunJob, run *Run, ok bool, err job.Started = now job.Status = core.StatusRunning - task = &Task{ + task := &Task{ JobID: job.ID, Attempt: job.Attempt, RunnerID: runner.ID, Started: now, } - if err = db.Insert(ctx, task); err != nil { - return + if err := db.Insert(ctx, task); err != nil { + return nil, false, err } job.TaskID = task.ID - if _, err = db.GetEngine(ctx).ID(job.ID).Update(job); err != nil { - return + if _, err := db.GetEngine(ctx).ID(job.ID).Update(job); err != nil { + return nil, false, err } - run = &Run{} - if _, err = db.GetEngine(ctx).ID(job.RunID).Get(run); err != nil { - return + task.Job = job + if err := task.Job.LoadAttributes(ctx); err != nil { + return nil, false, err } - if err = commiter.Commit(); err != nil { - return + if err := commiter.Commit(); err != nil { + return nil, false, err } - ok = true - return + return task, true, nil } func isSubset(set, subset []string) bool { diff --git a/models/bots/task_step.go b/models/bots/task_step.go index f8319cbde0..c54cb6f913 100644 --- a/models/bots/task_step.go +++ b/models/bots/task_step.go @@ -7,6 +7,7 @@ package bots import ( "code.gitea.io/gitea/models/db" "code.gitea.io/gitea/modules/timeutil" + runnerv1 "gitea.com/gitea/proto-go/runner/v1" ) // TaskStep represents a step of Task @@ -14,7 +15,7 @@ type TaskStep struct { ID int64 TaskID int64 Number int64 - Result int64 // TODO: use runnerv1.Result + Result runnerv1.Result LogIndex int64 LogLength int64 Started timeutil.TimeStamp diff --git a/routers/api/bots/runner/runner.go b/routers/api/bots/runner/runner.go index aadd00bf6a..4ae385f83e 100644 --- a/routers/api/bots/runner/runner.go +++ b/routers/api/bots/runner/runner.go @@ -10,6 +10,7 @@ import ( "fmt" "net/url" "strings" + "time" "code.gitea.io/gitea/core" bots_model "code.gitea.io/gitea/models/bots" @@ -140,6 +141,17 @@ func (s *Service) FetchTask( task = t } + // avoid crazy retry + if task == nil { + duration := 2 * time.Second + if deadline, ok := ctx.Deadline(); ok { + if d := time.Until(deadline) - time.Second; d < duration { + duration = d + } + } + time.Sleep(duration) + } + res := connect.NewResponse(&runnerv1.FetchTaskResponse{ Task: task, }) @@ -165,7 +177,7 @@ func (s *Service) UpdateLog( } func (s *Service) pickTask(ctx context.Context, runner *bots_model.Runner) (*runnerv1.Task, bool, error) { - t, job, run, ok, err := bots_model.CreateTask(runner) + t, ok, err := bots_model.CreateTask(runner) if err != nil { return nil, false, fmt.Errorf("CreateTask: %w", err) } @@ -174,31 +186,31 @@ func (s *Service) pickTask(ctx context.Context, runner *bots_model.Runner) (*run } event := map[string]interface{}{} - _ = json.Unmarshal([]byte(run.EventPayload), &event) + _ = json.Unmarshal([]byte(t.Job.Run.EventPayload), &event) // TODO: more context in https://docs.github.com/cn/actions/learn-github-actions/contexts#github-context taskContext, _ := structpb.NewStruct(map[string]interface{}{ "event": event, - "run_id": fmt.Sprint(run.ID), - "run_number": fmt.Sprint(run.Index), - "run_attempt": fmt.Sprint(job.Attempt), - "actor": fmt.Sprint(run.TriggerUser.Name), - "repository": fmt.Sprint(run.Repo.Name), - "event_name": fmt.Sprint(run.Event.Event()), - "sha": fmt.Sprint(run.CommitSHA), - "ref": fmt.Sprint(run.Ref), + "run_id": fmt.Sprint(t.Job.ID), + "run_number": fmt.Sprint(t.Job.Run.Index), + "run_attempt": fmt.Sprint(t.Job.Attempt), + "actor": fmt.Sprint(t.Job.Run.TriggerUser.Name), + "repository": fmt.Sprint(t.Job.Run.Repo.Name), + "event_name": fmt.Sprint(t.Job.Run.Event.Event()), + "sha": fmt.Sprint(t.Job.Run.CommitSHA), + "ref": fmt.Sprint(t.Job.Run.Ref), "ref_name": "", "ref_type": "", "head_ref": "", "base_ref": "", "token": "", - "repository_owner": fmt.Sprint(run.Repo.OwnerName), + "repository_owner": fmt.Sprint(t.Job.Run.Repo.OwnerName), "retention_days": "", }) task := &runnerv1.Task{ Id: t.ID, - WorkflowPayload: job.WorkflowPayload, + WorkflowPayload: t.Job.WorkflowPayload, Context: taskContext, Secrets: nil, // TODO: query secrets } diff --git a/routers/api/bots/runner/unary.go b/routers/api/bots/runner/unary.go index 005db2103c..df6324e144 100644 --- a/routers/api/bots/runner/unary.go +++ b/routers/api/bots/runner/unary.go @@ -8,9 +8,11 @@ import ( "context" "strings" - "code.gitea.io/gitea/models/bots" + bots_model "code.gitea.io/gitea/models/bots" "github.com/bufbuild/connect-go" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" ) var WithRunner = connect.WithInterceptors(connect.UnaryInterceptorFunc(func(unaryFunc connect.UnaryFunc) connect.UnaryFunc { @@ -18,12 +20,15 @@ var WithRunner = connect.WithInterceptors(connect.UnaryInterceptorFunc(func(unar if methodName(request) == "Register" { return unaryFunc(ctx, request) } - uuid := request.Header().Get("X-Runner-Token") // TODO: shouldn't be X-Runner-Token, maybe X-Runner-UUID - // TODO: get runner from db, refuse request if it doesn't exist - r := &bots.Runner{ - UUID: uuid, + token := request.Header().Get("X-Runner-Token") // TODO: shouldn't be X-Runner-Token, maybe X-Runner-UUID + runner, err := bots_model.GetRunnerByToken(token) + if err != nil { + if _, ok := err.(*bots_model.ErrRunnerNotExist); ok { + return nil, status.Error(codes.Unauthenticated, "unregistered runner") + } + return nil, status.Error(codes.Internal, err.Error()) } - ctx = context.WithValue(ctx, runnerCtxKey{}, r) + ctx = context.WithValue(ctx, runnerCtxKey{}, runner) return unaryFunc(ctx, request) } })) @@ -38,9 +43,9 @@ func methodName(req connect.AnyRequest) string { type runnerCtxKey struct{} -func GetRunner(ctx context.Context) *bots.Runner { +func GetRunner(ctx context.Context) *bots_model.Runner { if v := ctx.Value(runnerCtxKey{}); v != nil { - if r, ok := v.(*bots.Runner); ok { + if r, ok := v.(*bots_model.Runner); ok { return r } }