352 lines
12 KiB
Go
352 lines
12 KiB
Go
package main
|
||
|
||
import (
|
||
"bytes"
|
||
"context"
|
||
"encoding/json"
|
||
"errors"
|
||
"fmt"
|
||
"reflect"
|
||
"sync"
|
||
"time"
|
||
|
||
orbit "git.yuntuo.hk/k2safe/orbit/sdk/go"
|
||
)
|
||
|
||
// schedulePayload 冻结运行和计划序号；精确字段集合用于检测错配输入。
|
||
type schedulePayload struct {
|
||
RunID string `json:"run_id"` // RunID 本次测量身份，不是业务幂等保证。
|
||
Schedule int `json:"schedule"` // Schedule 从零开始的计划序号。
|
||
}
|
||
|
||
func decodePayload(raw json.RawMessage) (schedulePayload, bool) {
|
||
// map中的RawMessage保留整数文本，不接受缺字段、null和非精确字段名。
|
||
var fields map[string]json.RawMessage
|
||
var out schedulePayload
|
||
if json.Unmarshal(raw, &fields) != nil || len(fields) != 2 {
|
||
return out, false
|
||
}
|
||
run, hasRun := fields["run_id"]
|
||
index, hasIndex := fields["schedule"]
|
||
if !hasRun || !hasIndex || string(run) == "null" || string(index) == "null" || json.Unmarshal(run, &out.RunID) != nil || json.Unmarshal(index, &out.Schedule) != nil {
|
||
return schedulePayload{}, false
|
||
}
|
||
// 按原文计数，拒绝重复键；字段次序可以不同。
|
||
var keys []string
|
||
decoder := json.NewDecoder(bytes.NewReader(raw))
|
||
if _, err := decoder.Token(); err != nil {
|
||
return schedulePayload{}, false
|
||
}
|
||
for decoder.More() {
|
||
k, err := decoder.Token()
|
||
if err != nil {
|
||
return schedulePayload{}, false
|
||
}
|
||
keys = append(keys, k.(string))
|
||
var value json.RawMessage
|
||
if decoder.Decode(&value) != nil {
|
||
return schedulePayload{}, false
|
||
}
|
||
}
|
||
return out, len(keys) == 2
|
||
}
|
||
|
||
func handlerName(runID string) string { return "load.schedule." + runID }
|
||
func operationKey(runID string, index int, op string) string {
|
||
return fmt.Sprintf("schedule-%s-%02d-%s", runID, index, op)
|
||
}
|
||
func businessKey(runID string, index int) string { return operationKey(runID, index, "task") }
|
||
|
||
type scheduleAPI interface {
|
||
GetQueue(context.Context, string) (orbit.Queue, error)
|
||
CreateSchedule(context.Context, string, orbit.CreateScheduleRequest) (orbit.ScheduleResult, error)
|
||
GetSchedule(context.Context, string) (orbit.Schedule, error)
|
||
ResumeSchedule(context.Context, string, string, orbit.ControlScheduleRequest) (orbit.ScheduleResult, error)
|
||
PauseSchedule(context.Context, string, string, orbit.ControlScheduleRequest) (orbit.ScheduleResult, error)
|
||
ListScheduleOccurrences(context.Context, string, orbit.SchedulePageRequest) (orbit.ScheduleOccurrencePage, error)
|
||
GetTask(context.Context, string) (orbit.Task, error)
|
||
ListTaskAttempts(context.Context, string, orbit.HistoryPageRequest) (orbit.AttemptPage, error)
|
||
GetTaskSource(context.Context, string) (orbit.TaskSource, error)
|
||
}
|
||
|
||
func scheduleDefinition(cfg config, runID string, index int, start time.Time) orbit.ScheduleDefinition {
|
||
end := start.Add(time.Duration(cfg.Requests) * cfg.Interval)
|
||
payload, _ := json.Marshal(schedulePayload{RunID: runID, Schedule: index})
|
||
return orbit.ScheduleDefinition{
|
||
Rule: orbit.ScheduleRule{Kind: "FIXED_RATE", Timezone: "UTC", StartAt: start, EndAt: &end, IntervalMS: uint64(cfg.Interval.Milliseconds())},
|
||
Policy: orbit.SchedulePolicy{Misfire: "CATCH_UP", Overlap: "PARALLEL", CatchupLimit: 100, CatchupWindowMS: uint64((cfg.RunTimeout + time.Duration(cfg.Requests)*cfg.Interval + cfg.StartDelay).Milliseconds()), GraceMS: 1000},
|
||
Task: orbit.SubmitTaskRequest{QueueID: cfg.QueueID, Handler: handlerName(runID), HandlerVersion: "v1", PayloadVersion: "v1", BusinessKey: businessKey(runID, index), Payload: payload, TaskTimeoutMS: uint64(cfg.RunTimeout.Milliseconds()), RecoverySafety: orbit.RecoveryReconcileOnly, Retry: orbit.RetryPolicy{MaxAttempts: 1, AttemptTimeoutMS: uint64(cfg.RunTimeout.Milliseconds())}},
|
||
}
|
||
}
|
||
|
||
func measure(parent context.Context, cfg config, api *orbit.Client, runID string) (r scheduleReport) {
|
||
r = newReport(cfg, runID)
|
||
ctx, cancel := context.WithTimeout(parent, cfg.RunTimeout)
|
||
defer cancel()
|
||
var workers []*orbit.Worker
|
||
var stopWorkers context.CancelFunc
|
||
var done chan error
|
||
log := newInvocationLog(cfg.Requests * cfg.Schedules)
|
||
observations := newObservations(cfg, runID)
|
||
// 所有返回路径都执行有界收尾；超时后的暂停使用独立上下文，结果不明不换键重试。
|
||
defer func() {
|
||
if stopWorkers != nil {
|
||
stopWorkers()
|
||
for range workers {
|
||
if err := <-done; err != nil {
|
||
r.Errors["WORKER_STOP"]++
|
||
}
|
||
}
|
||
}
|
||
for _, w := range workers {
|
||
s := w.Stats()
|
||
r.WorkerStarted += s.Started
|
||
r.Reconnects += s.Reconnects
|
||
r.Renewals += s.RenewalsConfirmed
|
||
}
|
||
cleanup, stop := context.WithTimeout(context.Background(), 5*time.Second)
|
||
var wg sync.WaitGroup
|
||
codes := make(chan string, len(r.Plans))
|
||
for i := range r.Plans {
|
||
wg.Add(1)
|
||
go func(p *planReport) {
|
||
defer wg.Done()
|
||
if validID(p.ID) {
|
||
codes <- pauseSchedule(cleanup, cfg, api, runID, p)
|
||
}
|
||
}(&r.Plans[i])
|
||
}
|
||
wg.Wait()
|
||
stop()
|
||
close(codes)
|
||
for code := range codes {
|
||
if code != "" {
|
||
r.Errors[code]++
|
||
}
|
||
}
|
||
observations.summarize(&r)
|
||
log.summarize(&r, observations)
|
||
if ctx.Err() != nil {
|
||
r.Errors["RUN_"+classify(ctx.Err())]++
|
||
}
|
||
r = r.finish()
|
||
}()
|
||
call, stop := context.WithTimeout(ctx, cfg.Timeout)
|
||
queue, err := api.GetQueue(call, cfg.QueueID)
|
||
stop()
|
||
if err != nil {
|
||
r.Errors["PREFLIGHT_"+classify(err)]++
|
||
return
|
||
}
|
||
r.Queue = &queue
|
||
if queue.ID != cfg.QueueID || !validID(queue.NamespaceID) || queue.Paused || queue.MaxConcurrency == 0 || queue.RatePerSecond == 0 {
|
||
r.Errors["QUEUE_UNAVAILABLE"]++
|
||
return
|
||
}
|
||
r.ScheduleStart = time.Now().UTC().Truncate(time.Microsecond).Add(cfg.StartDelay)
|
||
r.ScheduleEnd = r.ScheduleStart.Add(time.Duration(cfg.Requests) * cfg.Interval)
|
||
// 在恢复任何计划之前完成全部创建；计划有限终止，即使进程被强杀也不会无限触发。
|
||
seen := map[string]bool{}
|
||
for i := 0; i < cfg.Schedules; i++ {
|
||
if ctx.Err() != nil {
|
||
return
|
||
}
|
||
definition := scheduleDefinition(cfg, runID, i, r.ScheduleStart)
|
||
key := operationKey(runID, i, "create")
|
||
p := planReport{Index: i, CreateKey: key}
|
||
r.Plans = append(r.Plans, p)
|
||
call, stop = context.WithTimeout(ctx, cfg.Timeout)
|
||
res, e := api.CreateSchedule(call, key, orbit.CreateScheduleRequest{Name: key, Paused: true, Definition: definition})
|
||
stop()
|
||
if e != nil {
|
||
r.Errors["CREATE_"+classify(e)]++
|
||
return
|
||
}
|
||
p.ID = res.Schedule.ID
|
||
p.CalendarVersion = res.Schedule.Current.Definition.Rule.CalendarVersion
|
||
p.EffectiveAt = res.Schedule.Current.EffectiveAt
|
||
r.Plans[i] = p
|
||
if res.Deduplicated || !validID(p.ID) || seen[p.ID] || !matchesSchedule(res.Schedule, cfg, runID, i, queue.NamespaceID, r.ScheduleStart) || !res.Schedule.Paused || res.Schedule.Current.EffectiveAt.After(r.ScheduleStart) {
|
||
r.Errors["CREATE_MISMATCH"]++
|
||
return
|
||
}
|
||
seen[p.ID] = true
|
||
observations.addPlan(p.ID, i, r.ScheduleStart, queue.NamespaceID)
|
||
}
|
||
for i := 0; i < cfg.Workers; i++ {
|
||
w, e := orbit.NewWorker(api, orbit.WorkerOptions{QueueID: cfg.QueueID, Capacity: uint32(cfg.Capacity), ShutdownTimeout: 5 * time.Second, Handlers: []orbit.WorkerHandler{log.handler(ctx, cfg, runID)}})
|
||
if e != nil {
|
||
r.Errors["WORKER_CONFIG"]++
|
||
return
|
||
}
|
||
workers = append(workers, w)
|
||
}
|
||
workerCtx, workerCancel := context.WithCancel(ctx)
|
||
stopWorkers = workerCancel
|
||
done = make(chan error, len(workers))
|
||
for _, w := range workers {
|
||
go func(w *orbit.Worker) {
|
||
err := w.Run(workerCtx)
|
||
if workerCtx.Err() == nil {
|
||
cancel()
|
||
}
|
||
done <- err
|
||
}(w)
|
||
}
|
||
for i := range r.Plans {
|
||
p := &r.Plans[i]
|
||
call, stop = context.WithTimeout(ctx, cfg.Timeout)
|
||
current, e := api.GetSchedule(call, p.ID)
|
||
stop()
|
||
if e != nil {
|
||
r.Errors["RESUME_READ_"+classify(e)]++
|
||
return
|
||
}
|
||
if !time.Now().Before(r.ScheduleStart) || current.ID != p.ID || !matchesSchedule(current, cfg, runID, i, queue.NamespaceID, r.ScheduleStart) || !current.Paused || !sameVersion(current, *p) {
|
||
r.Errors["SETUP_LATE_OR_CHANGED"]++
|
||
return
|
||
}
|
||
p.ResumeKey = operationKey(runID, i, "resume")
|
||
p.ResumeRevision = current.Revision
|
||
call, stop = context.WithTimeout(ctx, cfg.Timeout)
|
||
res, e := api.ResumeSchedule(call, p.ID, p.ResumeKey, orbit.ControlScheduleRequest{ExpectedRevision: current.Revision, Reason: "start bounded schedule verification"})
|
||
stop()
|
||
if e != nil {
|
||
r.Errors["RESUME_"+classify(e)]++
|
||
return
|
||
}
|
||
if res.Deduplicated || res.Schedule.ID != p.ID || res.Schedule.Paused || !matchesSchedule(res.Schedule, cfg, runID, i, queue.NamespaceID, r.ScheduleStart) || !res.Schedule.UpdatedAt.Before(r.ScheduleStart) || !sameVersion(res.Schedule, *p) {
|
||
r.Errors["RESUME_MISMATCH_OR_LATE"]++
|
||
return
|
||
}
|
||
p.Resumed = true
|
||
}
|
||
for {
|
||
if ctx.Err() != nil {
|
||
return
|
||
}
|
||
all := true
|
||
for i := range r.Plans {
|
||
p := &r.Plans[i]
|
||
call, stop = context.WithTimeout(ctx, cfg.Timeout)
|
||
page, e := api.ListScheduleOccurrences(call, p.ID, orbit.SchedulePageRequest{Limit: 100})
|
||
stop()
|
||
if e != nil {
|
||
r.Errors["OCCURRENCES_"+classify(e)]++
|
||
return
|
||
}
|
||
if code := observations.ingest(p.ID, page); code != "" {
|
||
r.Errors[code]++
|
||
return
|
||
}
|
||
complete, code := observations.verify(ctx, api, p.ID)
|
||
if code != "" {
|
||
r.Errors[code]++
|
||
return
|
||
}
|
||
if complete {
|
||
call, stop = context.WithTimeout(ctx, cfg.Timeout)
|
||
current, e := api.GetSchedule(call, p.ID)
|
||
stop()
|
||
if e != nil {
|
||
r.Errors["FINAL_READ_"+classify(e)]++
|
||
return
|
||
}
|
||
if current.ID != p.ID || !matchesSchedule(current, cfg, runID, i, queue.NamespaceID, r.ScheduleStart) || current.Paused || !sameVersion(current, *p) {
|
||
r.Errors["FINAL_SCHEDULE_MISMATCH"]++
|
||
return
|
||
}
|
||
p.Exhausted = current.CheckAt == nil && current.NextFireAt == nil && current.ProcessingVersion == 1
|
||
}
|
||
all = all && complete && p.Exhausted
|
||
}
|
||
if all {
|
||
return
|
||
}
|
||
timer := time.NewTimer(cfg.PollInterval)
|
||
select {
|
||
case <-ctx.Done():
|
||
timer.Stop()
|
||
return
|
||
case <-timer.C:
|
||
}
|
||
}
|
||
}
|
||
|
||
func matchesSchedule(s orbit.Schedule, cfg config, runID string, index int, namespace string, start time.Time) bool {
|
||
want := scheduleDefinition(cfg, runID, index, start)
|
||
got := s.Current.Definition
|
||
// 服务端冻结具体时区快照标识；不把API版本v1误当作时区数据版本。
|
||
if !validCalendarVersion(got.Rule.CalendarVersion) {
|
||
return false
|
||
}
|
||
got.Rule.CalendarVersion = ""
|
||
return s.NamespaceID == namespace && s.QueueID == cfg.QueueID && s.Name == operationKey(runID, index, "create") && s.CurrentVersion == 1 && s.Current.Version == 1 && s.Revision > 0 && !s.Current.EffectiveAt.IsZero() && sameDefinition(got, want)
|
||
}
|
||
|
||
func sameDefinition(a, b orbit.ScheduleDefinition) bool {
|
||
// 原始JSON可有空白差异；模板和结果保持精确整数，不经float64。
|
||
pa, ok := decodePayload(a.Task.Payload)
|
||
pb, ok2 := decodePayload(b.Task.Payload)
|
||
a.Task.Payload = nil
|
||
b.Task.Payload = nil
|
||
return ok && ok2 && pa == pb && reflect.DeepEqual(a, b)
|
||
}
|
||
|
||
func pauseSchedule(ctx context.Context, cfg config, api scheduleAPI, runID string, p *planReport) string {
|
||
for attempt := 0; attempt < 3; attempt++ {
|
||
call, stop := context.WithTimeout(ctx, cfg.Timeout)
|
||
current, err := api.GetSchedule(call, p.ID)
|
||
stop()
|
||
if err != nil {
|
||
return "PAUSE_READ_" + classify(err)
|
||
}
|
||
if current.ID != p.ID || current.Name != operationKey(runID, p.Index, "create") || current.QueueID != cfg.QueueID {
|
||
return "PAUSE_IDENTITY_MISMATCH"
|
||
}
|
||
if current.Paused {
|
||
p.PauseConfirmed = true
|
||
return ""
|
||
}
|
||
p.PauseKey = operationKey(runID, p.Index, fmt.Sprintf("pause-%d", attempt))
|
||
p.PauseRevision = current.Revision
|
||
call, stop = context.WithTimeout(ctx, cfg.Timeout)
|
||
res, err := api.PauseSchedule(call, p.ID, p.PauseKey, orbit.ControlScheduleRequest{ExpectedRevision: p.PauseRevision, Reason: "stop bounded schedule verification"})
|
||
stop()
|
||
if err == nil {
|
||
if res.Schedule.ID != p.ID || !res.Schedule.Paused {
|
||
return "PAUSE_RECEIPT_MISMATCH"
|
||
}
|
||
p.PauseConfirmed = true
|
||
return ""
|
||
}
|
||
var apiErr *orbit.APIError
|
||
if !errors.As(err, &apiErr) || apiErr.Code != "STATE_CONFLICT" || apiErr.StatusCode != 409 {
|
||
return "PAUSE_" + classify(err)
|
||
}
|
||
timer := time.NewTimer(25 * time.Millisecond)
|
||
select {
|
||
case <-ctx.Done():
|
||
timer.Stop()
|
||
return "PAUSE_" + classify(ctx.Err())
|
||
case <-timer.C:
|
||
}
|
||
}
|
||
return "PAUSE_CONFLICT_LIMIT"
|
||
}
|
||
|
||
func validCalendarVersion(v string) bool {
|
||
if len(v) == 0 || len(v) > 64 {
|
||
return false
|
||
}
|
||
for _, c := range v {
|
||
if !(c >= 'a' && c <= 'z' || c >= 'A' && c <= 'Z' || c >= '0' && c <= '9' || c == '.' || c == '-' || c == '_') {
|
||
return false
|
||
}
|
||
}
|
||
return true
|
||
}
|
||
func sameVersion(s orbit.Schedule, p planReport) bool {
|
||
return s.Current.Definition.Rule.CalendarVersion == p.CalendarVersion && s.Current.EffectiveAt.Equal(p.EffectiveAt)
|
||
}
|