| 1 | package activator |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "sync" |
| 6 | "time" |
| 7 | ) |
| 8 | |
| 9 | type cleanupBudgetKey struct{} |
| 10 | type cleanupContextKey struct{} |
| 11 | type cleanupBudget struct { |
| 12 | mu sync.Mutex |
| 13 | base context.Context |
| 14 | source context.Context |
| 15 | parent *cleanupBudget |
| 16 | limit time.Duration |
| 17 | ctx context.Context |
| 18 | cancel context.CancelFunc |
| 19 | } |
| 20 | |
| 21 | func (b *cleanupBudget) start() context.Context { |
| 22 | b.mu.Lock() |
| 23 | defer b.mu.Unlock() |
| 24 | if b.ctx == nil { |
| 25 | base := b.base |
| 26 | if b.parent != nil && b.parent.source.Err() != nil { |
| 27 | base = b.parent.start() |
| 28 | } |
| 29 | b.ctx, b.cancel = context.WithTimeout(base, b.limit) |
| 30 | } |
| 31 | return b.ctx |
| 32 | } |
| 33 | |
| 34 | // WithCleanupBudget shares the cancellation grace period across a request, |
| 35 | // its nested failure paths, and the owner's final join. Normal work is untimed. |
| 36 | func WithCleanupBudget(ctx context.Context, limit time.Duration) (context.Context, context.CancelFunc) { |
| 37 | parent, _ := ctx.Value(cleanupBudgetKey{}).(*cleanupBudget) |
| 38 | b := &cleanupBudget{base: context.WithoutCancel(ctx), source: ctx, parent: parent, limit: limit} |
| 39 | stop := context.AfterFunc(ctx, func() { b.start() }) |
| 40 | return context.WithValue(ctx, cleanupBudgetKey{}, b), func() { stop(); b.start(); b.cancel() } |
| 41 | } |
| 42 | |
| 43 | // CleanupContext keeps an existing cleanup deadline. Cancellation starts the |
| 44 | // owner's single grace period; it never grants each nested call a fresh one. |
| 45 | func CleanupContext(ctx context.Context, limit time.Duration) (context.Context, context.CancelFunc) { |
| 46 | if ctx.Value(cleanupContextKey{}) != nil { |
| 47 | return context.WithTimeout(ctx, limit) |
| 48 | } |
| 49 | base := context.WithoutCancel(ctx) |
| 50 | var stop func() bool |
| 51 | if b, ok := ctx.Value(cleanupBudgetKey{}).(*cleanupBudget); ok { |
| 52 | if ctx.Err() != nil { |
| 53 | base = b.start() |
| 54 | } |
| 55 | cleanup, cancel := context.WithTimeout(base, limit) |
| 56 | stop = context.AfterFunc(ctx, func() { budget := b.start(); context.AfterFunc(budget, cancel) }) |
| 57 | return context.WithValue(cleanup, cleanupContextKey{}, true), func() { stop(); cancel() } |
| 58 | } |
| 59 | // Callers without an owner budget still cannot extend an explicit deadline. |
| 60 | if deadline, ok := ctx.Deadline(); ok { |
| 61 | bounded, cancel := context.WithDeadline(base, deadline) |
| 62 | cleanup, done := context.WithTimeout(bounded, limit) |
| 63 | return context.WithValue(cleanup, cleanupContextKey{}, true), func() { done(); cancel() } |
| 64 | } |
| 65 | cleanup, cancel := context.WithTimeout(base, limit) |
| 66 | return context.WithValue(cleanup, cleanupContextKey{}, true), cancel |
| 67 | } |
| 68 | |
amadan