《后端架构设计·高性能与高可用》1.3 性能预算：把端到端延迟拆到每个调用

// file: internal/platform/budget/budget.go
// tier: L1
// verify: go test ./internal/platform/budget/ -count=1
package budget

import (
	"context"
	"sort"
	"time"
)

type Mode string

const (
	Sequential Mode = "sequential"
	Parallel   Mode = "parallel"
	Retry      Mode = "retry"
)

type Step struct {
	Name       string
	Mode       Mode
	Group      string
	Estimate   time.Duration
	Attempts   int
	Reserve    time.Duration
	Degradable bool
	Priority   int
}

type Plan struct{ Steps []Step }

type StepAllocation struct {
	Name           string
	Deadline       time.Duration
	AttemptTimeout time.Duration
	Attempts       int
}

type StageAllocation struct {
	Deadline time.Duration
	Budget   time.Duration
	Steps    []StepAllocation
}

type Allocation struct {
	Total        time.Duration
	Need         time.Duration
	Feasible     bool
	Shortfall    time.Duration
	Degradations []string
	Stages       []StageAllocation
}

type stage struct {
	steps []Step
	need  time.Duration
}

func buildStages(steps []Step) []stage {
	var out []stage
	for i := 0; i < len(steps); i++ {
		s := steps[i]
		if s.Mode == Parallel && s.Group != "" {
			j := i + 1
			for j < len(steps) && steps[j].Mode == Parallel && steps[j].Group == s.Group {
				j++
			}
			group := steps[i:j]
			need := group[0].Estimate
			for _, g := range group[1:] {
				if g.Estimate > need {
					need = g.Estimate
				}
			}
			out = append(out, stage{steps: group, need: need})
			i = j - 1
			continue
		}
		need := s.Estimate
		if s.Mode == Retry {
			attempts := s.Attempts
			if attempts < 1 {
				attempts = 1
			}
			need = time.Duration(attempts)*s.Estimate + s.Reserve
		}
		out = append(out, stage{steps: []Step{s}, need: need})
	}
	return out
}

func totalNeed(stages []stage) time.Duration {
	var total time.Duration
	for _, st := range stages {
		total += st.need
	}
	return total
}

// Allocate converts a plan into cumulative deadlines. Over budget, it drops
// degradable steps by descending Priority until the plan fits.
func (p Plan) Allocate(total time.Duration) Allocation {
	steps := append([]Step(nil), p.Steps...)
	alloc := Allocation{Total: total}
	stages := buildStages(steps)
	alloc.Need = totalNeed(stages)

	if total > 0 && alloc.Need > total {
		candidates := make([]Step, 0)
		for _, s := range steps {
			if s.Degradable {
				candidates = append(candidates, s)
			}
		}
		sort.SliceStable(candidates, func(i, j int) bool { return candidates[i].Priority > candidates[j].Priority })
		for _, cand := range candidates {
			filtered := steps[:0]
			for _, s := range steps {
				if s.Name != cand.Name {
					filtered = append(filtered, s)
				}
			}
			steps = filtered
			alloc.Degradations = append(alloc.Degradations, cand.Name)
			stages = buildStages(steps)
			alloc.Need = totalNeed(stages)
			if alloc.Need <= total {
				break
			}
		}
	}
	alloc.Feasible = alloc.Need <= total
	if !alloc.Feasible {
		alloc.Shortfall = alloc.Need - total
	}

	var cursor time.Duration
	for _, st := range stages {
		stage := StageAllocation{Deadline: cursor + st.need, Budget: st.need}
		for _, s := range st.steps {
			attempts := s.Attempts
			if attempts < 1 {
				attempts = 1
			}
			stage.Steps = append(stage.Steps, StepAllocation{
				Name: s.Name, Deadline: stage.Deadline, AttemptTimeout: s.Estimate, Attempts: attempts,
			})
		}
		alloc.Stages = append(alloc.Stages, stage)
		cursor = stage.Deadline
	}
	return alloc
}

// WithRemaining takes min(parent deadline, now+budget); it never extends the
// parent, so downstream cannot spend budget the caller no longer has.
func WithRemaining(parent context.Context, budget time.Duration) (context.Context, context.CancelFunc) {
	deadline := time.Now().Add(budget)
	if d, ok := parent.Deadline(); ok && d.Before(deadline) {
		deadline = d
	}
	return context.WithDeadline(parent, deadline)
}

type AttemptFunc func(ctx context.Context, attempt int) error

// RunRetries gives each attempt its own timeout and spends Reserve only on
// gaps. The stage deadline bounds the whole sequence, so retries cannot
// multiply the caller's timeout.
func RunRetries(parent context.Context, stageDeadline time.Time, attemptTimeout time.Duration, attempts int, reserve time.Duration, fn AttemptFunc) error {
	if attempts < 1 {
		attempts = 1
	}
	stageCtx, cancel := context.WithDeadline(parent, stageDeadline)
	defer cancel()
	gap := time.Duration(0)
	if attempts > 1 {
		gap = reserve / time.Duration(attempts-1)
	}
	var last error
	for i := 1; i <= attempts; i++ {
		attemptCtx, attemptCancel := context.WithTimeout(stageCtx, attemptTimeout)
		err := fn(attemptCtx, i)
		attemptCancel()
		if err == nil {
			return nil
		}
		last = err
		if i == attempts {
			break
		}
		if gap > 0 {
			t := time.NewTimer(gap)
			select {
			case <-stageCtx.Done():
				t.Stop()
				return last
			case <-t.C:
			}
		}
	}
	return last
}

$ go test ./internal/platform/budget/ -count=5 -v
attempts=3 elapsed=95.707625ms stage_deadline=120ms attempt_timeout=20ms reserve=30ms
attempts=3 elapsed=95.171ms    stage_deadline=120ms attempt_timeout=20ms reserve=30ms
attempts=3 elapsed=96.411125ms stage_deadline=120ms attempt_timeout=20ms reserve=30ms
attempts=3 elapsed=96.294792ms stage_deadline=120ms attempt_timeout=20ms reserve=30ms
attempts=3 elapsed=96.909333ms stage_deadline=120ms attempt_timeout=20ms reserve=30ms