| 1 | package extension |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "fmt" |
| 6 | "sync" |
| 7 | "time" |
| 8 | ) |
| 9 | |
| 10 | // PublishGate enforces the activate → publish → drain order for runtime |
| 11 | // generations. Only one generation is Published at a time; older generations |
| 12 | // are Draining and their late traffic must be dropped. |
| 13 | type PublishGate struct { |
| 14 | mu sync.RWMutex |
| 15 | published uint64 |
| 16 | draining map[uint64]time.Time // gen → drain start |
| 17 | drainTTL time.Duration |
| 18 | onStale func(gen uint64, kind string) |
| 19 | receipts *ReceiptStore |
| 20 | drainCancels map[uint64]map[uint64]func() |
| 21 | drainCancelID uint64 |
| 22 | expired map[uint64]struct{} |
| 23 | expiredOrder []uint64 |
| 24 | expiredLimit int |
| 25 | drainWatching bool |
| 26 | } |
| 27 | |
| 28 | const defaultExpiredGenerationLimit = 256 |
| 29 | |
| 30 | // NewPublishGate returns a gate with a default drain timeout of 30s. |
| 31 | func NewPublishGate() *PublishGate { |
| 32 | return newPublishGate(NewReceiptStore()) |
| 33 | } |
| 34 | |
| 35 | func newPublishGate(receipts *ReceiptStore) *PublishGate { |
| 36 | if receipts == nil { |
| 37 | receipts = NewReceiptStore() |
| 38 | } |
| 39 | return &PublishGate{ |
| 40 | draining: make(map[uint64]time.Time), |
| 41 | drainTTL: 30 * time.Second, |
| 42 | receipts: receipts, |
| 43 | drainCancels: make(map[uint64]map[uint64]func()), |
| 44 | expired: make(map[uint64]struct{}), |
| 45 | expiredLimit: defaultExpiredGenerationLimit, |
| 46 | } |
| 47 | } |
| 48 | |
| 49 | // WithDrainTTL sets how long a draining generation is tracked before force-forget. |
| 50 | func (g *PublishGate) WithDrainTTL(d time.Duration) *PublishGate { |
| 51 | if g != nil && d > 0 { |
| 52 | g.mu.Lock() |
| 53 | g.drainTTL = d |
| 54 | g.mu.Unlock() |
| 55 | } |
| 56 | return g |
| 57 | } |
| 58 | |
| 59 | // Published returns the currently published generation (0 if none). |
| 60 | func (g *PublishGate) Published() uint64 { |
| 61 | if g == nil { |
| 62 | return 0 |
| 63 | } |
| 64 | g.mu.RLock() |
| 65 | defer g.mu.RUnlock() |
| 66 | return g.published |
| 67 | } |
| 68 | |
| 69 | // Publish atomically switches the published generation. The previous generation |
| 70 | // enters Draining. Publishing the same generation is a no-op. |
| 71 | func (g *PublishGate) Publish(gen uint64) { |
| 72 | if g == nil || gen == 0 { |
| 73 | return |
| 74 | } |
| 75 | g.mu.Lock() |
| 76 | defer g.mu.Unlock() |
| 77 | if g.published == gen { |
| 78 | return |
| 79 | } |
| 80 | if g.published != 0 { |
| 81 | if g.draining == nil { |
| 82 | g.draining = make(map[uint64]time.Time) |
| 83 | } |
| 84 | g.draining[g.published] = time.Now() |
| 85 | } |
| 86 | delete(g.draining, gen) |
| 87 | g.clearExpiredLocked(gen) |
| 88 | g.published = gen |
| 89 | DefaultLifecycleMetrics.Publishes.Add(1) |
| 90 | } |
| 91 | |
| 92 | // BeginDrain marks gen as draining without changing the published pointer |
| 93 | // (used when an un-published activation fails and its resources are disposed). |
| 94 | func (g *PublishGate) BeginDrain(gen uint64) { |
| 95 | if g == nil || gen == 0 { |
| 96 | return |
| 97 | } |
| 98 | g.mu.Lock() |
| 99 | defer g.mu.Unlock() |
| 100 | if g.draining == nil { |
| 101 | g.draining = make(map[uint64]time.Time) |
| 102 | } |
| 103 | g.clearExpiredLocked(gen) |
| 104 | g.draining[gen] = time.Now() |
| 105 | DefaultLifecycleMetrics.Drains.Add(1) |
| 106 | } |
| 107 | |
| 108 | // IsStale reports whether messageGen must be dropped: it is non-zero and does |
| 109 | // not match the published generation. |
| 110 | func (g *PublishGate) IsStale(messageGen uint64) bool { |
| 111 | if g == nil || messageGen == 0 { |
| 112 | return false |
| 113 | } |
| 114 | g.mu.RLock() |
| 115 | pub := g.published |
| 116 | g.mu.RUnlock() |
| 117 | return StaleGeneration(messageGen, pub) |
| 118 | } |
| 119 | |
| 120 | // IsDraining reports whether gen is currently in the drain set. |
| 121 | func (g *PublishGate) IsDraining(gen uint64) bool { |
| 122 | if g == nil || gen == 0 { |
| 123 | return false |
| 124 | } |
| 125 | g.mu.RLock() |
| 126 | defer g.mu.RUnlock() |
| 127 | _, ok := g.draining[gen] |
| 128 | return ok |
| 129 | } |
| 130 | |
| 131 | // AdmitNewWork reports whether a new turn may be admitted for gen. Only the |
| 132 | // published generation admits new work; draining generations refuse. |
| 133 | func (g *PublishGate) AdmitNewWork(gen uint64) bool { |
| 134 | if g == nil { |
| 135 | return true |
| 136 | } |
| 137 | g.mu.RLock() |
| 138 | defer g.mu.RUnlock() |
| 139 | if g.published == 0 { |
| 140 | return true |
| 141 | } |
| 142 | return gen == g.published |
| 143 | } |
| 144 | |
| 145 | // DropStale logs (via onStale) and returns true when messageGen is stale. |
| 146 | func (g *PublishGate) DropStale(messageGen uint64, kind string) bool { |
| 147 | if !g.IsStale(messageGen) { |
| 148 | return false |
| 149 | } |
| 150 | DefaultLifecycleMetrics.StaleDrops.Add(1) |
| 151 | if g != nil && g.onStale != nil { |
| 152 | g.onStale(messageGen, kind) |
| 153 | } |
| 154 | return true |
| 155 | } |
| 156 | |
| 157 | // SweepExpiredDrains removes drain entries older than drainTTL. |
| 158 | func (g *PublishGate) SweepExpiredDrains() []uint64 { |
| 159 | if g == nil { |
| 160 | return nil |
| 161 | } |
| 162 | g.mu.Lock() |
| 163 | defer g.mu.Unlock() |
| 164 | var expired []uint64 |
| 165 | now := time.Now() |
| 166 | for gen, started := range g.draining { |
| 167 | if now.Sub(started) >= g.drainTTL { |
| 168 | expired = append(expired, gen) |
| 169 | delete(g.draining, gen) |
| 170 | g.markExpiredLocked(gen) |
| 171 | } |
| 172 | } |
| 173 | return expired |
| 174 | } |
| 175 | |
| 176 | // DrainTimeoutError is returned when remaining in-flight work is cancelled |
| 177 | // after the drain TTL. |
| 178 | type DrainTimeoutError struct { |
| 179 | Generation uint64 |
| 180 | } |
| 181 | |
| 182 | func (e *DrainTimeoutError) Error() string { |
| 183 | return fmt.Sprintf("extension: generation %d drain timed out", e.Generation) |
| 184 | } |
| 185 | |
| 186 | // RegisterDrainCancel registers a cancel func for gen and returns an idempotent |
| 187 | // unregister function. Remaining callbacks fire once when the generation is |
| 188 | // force-expired after drain TTL (or explicit ForceExpireDrain). |
| 189 | func (g *PublishGate) RegisterDrainCancel(gen uint64, cancel func()) func() { |
| 190 | if g == nil || gen == 0 || cancel == nil { |
| 191 | return func() {} |
| 192 | } |
| 193 | g.mu.Lock() |
| 194 | _, expired := g.expired[gen] |
| 195 | _, draining := g.draining[gen] |
| 196 | // Generation IDs increase monotonically. Once an old expiry marker leaves |
| 197 | // bounded retention, a generation below the published one is still stale |
| 198 | // and its late registration must be cancelled instead of retained forever. |
| 199 | forgottenExpired := !expired && !draining && g.published != 0 && gen < g.published |
| 200 | if expired || forgottenExpired { |
| 201 | g.mu.Unlock() |
| 202 | cancel() |
| 203 | return func() {} |
| 204 | } |
| 205 | if g.drainCancels == nil { |
| 206 | g.drainCancels = make(map[uint64]map[uint64]func()) |
| 207 | } |
| 208 | if g.drainCancels[gen] == nil { |
| 209 | g.drainCancels[gen] = make(map[uint64]func()) |
| 210 | } |
| 211 | g.drainCancelID++ |
| 212 | id := g.drainCancelID |
| 213 | g.drainCancels[gen][id] = cancel |
| 214 | g.mu.Unlock() |
| 215 | |
| 216 | var once sync.Once |
| 217 | return func() { |
| 218 | once.Do(func() { |
| 219 | g.mu.Lock() |
| 220 | if callbacks := g.drainCancels[gen]; callbacks != nil { |
| 221 | delete(callbacks, id) |
| 222 | if len(callbacks) == 0 { |
| 223 | delete(g.drainCancels, gen) |
| 224 | } |
| 225 | } |
| 226 | g.mu.Unlock() |
| 227 | }) |
| 228 | } |
| 229 | } |
| 230 | |
| 231 | // FireDrainCancels runs and clears all cancel callbacks for gen. |
| 232 | func (g *PublishGate) FireDrainCancels(gen uint64) { |
| 233 | if g == nil || gen == 0 { |
| 234 | return |
| 235 | } |
| 236 | g.mu.Lock() |
| 237 | fns := g.drainCancels[gen] |
| 238 | delete(g.drainCancels, gen) |
| 239 | g.markExpiredLocked(gen) |
| 240 | g.mu.Unlock() |
| 241 | for _, fn := range fns { |
| 242 | if fn != nil { |
| 243 | fn() |
| 244 | } |
| 245 | } |
| 246 | } |
| 247 | |
| 248 | // ForceExpireDrain cancels remaining in-flight work for gen, then records a |
| 249 | // cleanup receipt and forgets the drain entry. |
| 250 | func (g *PublishGate) ForceExpireDrain(gen uint64) { |
| 251 | if g == nil || gen == 0 { |
| 252 | return |
| 253 | } |
| 254 | g.mu.Lock() |
| 255 | fns := g.drainCancels[gen] |
| 256 | delete(g.drainCancels, gen) |
| 257 | delete(g.draining, gen) |
| 258 | g.markExpiredLocked(gen) |
| 259 | g.mu.Unlock() |
| 260 | for _, fn := range fns { |
| 261 | if fn != nil { |
| 262 | fn() |
| 263 | } |
| 264 | } |
| 265 | g.receipts.Record(EffectReceipt{ |
| 266 | ID: fmt.Sprintf("drain-timeout-%d", gen), |
| 267 | Generation: gen, |
| 268 | Class: Irreversible, |
| 269 | CompensationStatus: "not_applicable", |
| 270 | Error: (&DrainTimeoutError{Generation: gen}).Error(), |
| 271 | }) |
| 272 | DefaultLifecycleMetrics.Drains.Add(1) |
| 273 | } |
| 274 | |
| 275 | // SweepAndForceExpire removes expired drain entries, cancels in-flight work, |
| 276 | // and records a timeout receipt for each. |
| 277 | func (g *PublishGate) SweepAndForceExpire() []uint64 { |
| 278 | expired := g.SweepExpiredDrains() |
| 279 | for _, gen := range expired { |
| 280 | g.FireDrainCancels(gen) |
| 281 | g.receipts.Record(EffectReceipt{ |
| 282 | ID: fmt.Sprintf("drain-timeout-%d", gen), |
| 283 | Generation: gen, |
| 284 | Class: Irreversible, |
| 285 | CompensationStatus: "not_applicable", |
| 286 | Error: (&DrainTimeoutError{Generation: gen}).Error(), |
| 287 | }) |
| 288 | DefaultLifecycleMetrics.Drains.Add(1) |
| 289 | } |
| 290 | return expired |
| 291 | } |
| 292 | |
| 293 | // ScheduleDrainWatch starts one background timer while a gate has active |
| 294 | // draining generations. Calls are coalesced so cold publishes do not allocate |
| 295 | // a timer goroutine and rapid publishes do not create one watcher per publish. |
| 296 | func (g *PublishGate) ScheduleDrainWatch() { |
| 297 | if g == nil { |
| 298 | return |
| 299 | } |
| 300 | g.mu.Lock() |
| 301 | if len(g.draining) == 0 || g.drainWatching { |
| 302 | g.mu.Unlock() |
| 303 | return |
| 304 | } |
| 305 | g.drainWatching = true |
| 306 | ttl := g.drainTTL |
| 307 | if ttl <= 0 { |
| 308 | ttl = 30 * time.Second |
| 309 | } |
| 310 | wakeAfter := ttl |
| 311 | now := time.Now() |
| 312 | for _, started := range g.draining { |
| 313 | remaining := ttl - now.Sub(started) |
| 314 | if remaining < wakeAfter { |
| 315 | wakeAfter = remaining |
| 316 | } |
| 317 | } |
| 318 | if wakeAfter < 0 { |
| 319 | wakeAfter = 0 |
| 320 | } |
| 321 | g.mu.Unlock() |
| 322 | go func() { |
| 323 | timer := time.NewTimer(wakeAfter) |
| 324 | defer timer.Stop() |
| 325 | <-timer.C |
| 326 | g.SweepAndForceExpire() |
| 327 | g.mu.Lock() |
| 328 | g.drainWatching = false |
| 329 | watchAgain := len(g.draining) > 0 |
| 330 | g.mu.Unlock() |
| 331 | if watchAgain { |
| 332 | g.ScheduleDrainWatch() |
| 333 | } |
| 334 | }() |
| 335 | } |
| 336 | |
| 337 | func (g *PublishGate) markExpiredLocked(gen uint64) { |
| 338 | if gen == 0 { |
| 339 | return |
| 340 | } |
| 341 | if _, ok := g.expired[gen]; ok { |
| 342 | return |
| 343 | } |
| 344 | if g.expired == nil { |
| 345 | g.expired = make(map[uint64]struct{}) |
| 346 | } |
| 347 | g.expired[gen] = struct{}{} |
| 348 | g.expiredOrder = append(g.expiredOrder, gen) |
| 349 | limit := g.expiredLimit |
| 350 | if limit < 1 { |
| 351 | limit = defaultExpiredGenerationLimit |
| 352 | } |
| 353 | for len(g.expiredOrder) > limit { |
| 354 | old := g.expiredOrder[0] |
| 355 | g.expiredOrder = g.expiredOrder[1:] |
| 356 | delete(g.expired, old) |
| 357 | } |
| 358 | } |
| 359 | |
| 360 | func (g *PublishGate) clearExpiredLocked(gen uint64) { |
| 361 | if _, ok := g.expired[gen]; !ok { |
| 362 | return |
| 363 | } |
| 364 | delete(g.expired, gen) |
| 365 | for i, item := range g.expiredOrder { |
| 366 | if item != gen { |
| 367 | continue |
| 368 | } |
| 369 | copy(g.expiredOrder[i:], g.expiredOrder[i+1:]) |
| 370 | g.expiredOrder = g.expiredOrder[:len(g.expiredOrder)-1] |
| 371 | break |
| 372 | } |
| 373 | } |
| 374 | |
| 375 | // DrainingGenerations returns generations currently in drain. |
| 376 | func (g *PublishGate) DrainingGenerations() []uint64 { |
| 377 | if g == nil { |
| 378 | return nil |
| 379 | } |
| 380 | g.mu.RLock() |
| 381 | defer g.mu.RUnlock() |
| 382 | out := make([]uint64, 0, len(g.draining)) |
| 383 | for gen := range g.draining { |
| 384 | out = append(out, gen) |
| 385 | } |
| 386 | return out |
| 387 | } |
| 388 | |
| 389 | // DefaultPublishGate returns the compatibility owner gate. Product boot paths |
| 390 | // bind an isolated RuntimeOwner instead. |
| 391 | func DefaultPublishGate() *PublishGate { |
| 392 | return RuntimeOwnerOrDefault(nil).Gate |
| 393 | } |
| 394 | |
| 395 | // RegisterDrainCancel preserves the package-level compatibility API. |
| 396 | func RegisterDrainCancel(gen uint64, cancel func()) func() { |
| 397 | return DefaultPublishGate().RegisterDrainCancel(gen, cancel) |
| 398 | } |
| 399 | |
| 400 | // FireDrainCancels preserves the package-level compatibility API. |
| 401 | func FireDrainCancels(gen uint64) { |
| 402 | DefaultPublishGate().FireDrainCancels(gen) |
| 403 | } |
| 404 | |
| 405 | // AwaitReady waits until ctx is done or ready is closed. Used by activation |
| 406 | // transactions to gate publish on component readiness. |
| 407 | func AwaitReady(ctx context.Context, ready <-chan struct{}) error { |
| 408 | if ready == nil { |
| 409 | return nil |
| 410 | } |
| 411 | select { |
| 412 | case <-ctx.Done(): |
| 413 | return ctx.Err() |
| 414 | case <-ready: |
| 415 | return nil |
| 416 | } |
| 417 | } |
| 418 |