| 1 | package extension |
| 2 | |
| 3 | import ( |
| 4 | "errors" |
| 5 | "io" |
| 6 | "sync" |
| 7 | ) |
| 8 | |
| 9 | // RuntimeSet tracks the closable resources bound to one RuntimeSnapshot — |
| 10 | // MCP sidecar processes, watchers, temp files. It exists separately from the |
| 11 | // snapshot because a snapshot is immutable value state while its resources |
| 12 | // have a lifecycle: when a newer snapshot activates, the old RuntimeSet is |
| 13 | // closed exactly once, and CloseIfGeneration makes sure a stale cleanup path |
| 14 | // can never close resources that now belong to a newer runtime. |
| 15 | type RuntimeSet struct { |
| 16 | mu sync.Mutex |
| 17 | generation uint64 |
| 18 | closers []io.Closer |
| 19 | closed bool |
| 20 | } |
| 21 | |
| 22 | // NewRuntimeSet returns an empty set bound to a snapshot generation. |
| 23 | func NewRuntimeSet(generation uint64) *RuntimeSet { |
| 24 | return &RuntimeSet{generation: generation} |
| 25 | } |
| 26 | |
| 27 | // Generation returns the snapshot generation this set is bound to. |
| 28 | func (s *RuntimeSet) Generation() uint64 { return s.generation } |
| 29 | |
| 30 | // Add registers closers. A closer added to an already-closed set is closed |
| 31 | // immediately — the alternative (silently leaking it because the owner went |
| 32 | // away between activation and registration) is how sidecar processes outlive |
| 33 | // their session. |
| 34 | func (s *RuntimeSet) Add(closers ...io.Closer) { |
| 35 | s.mu.Lock() |
| 36 | if s.closed { |
| 37 | s.mu.Unlock() |
| 38 | for _, c := range closers { |
| 39 | if c != nil { |
| 40 | _ = c.Close() |
| 41 | } |
| 42 | } |
| 43 | return |
| 44 | } |
| 45 | for _, c := range closers { |
| 46 | if c != nil { |
| 47 | s.closers = append(s.closers, c) |
| 48 | } |
| 49 | } |
| 50 | s.mu.Unlock() |
| 51 | } |
| 52 | |
| 53 | // Len returns the number of registered closers. |
| 54 | func (s *RuntimeSet) Len() int { |
| 55 | s.mu.Lock() |
| 56 | defer s.mu.Unlock() |
| 57 | return len(s.closers) |
| 58 | } |
| 59 | |
| 60 | // Close releases every registered resource, in reverse registration order |
| 61 | // (later resources may depend on earlier ones). It is idempotent: concurrent |
| 62 | // or repeated calls after the first return nil without re-closing, because |
| 63 | // most closers are not safe to invoke twice. |
| 64 | func (s *RuntimeSet) Close() error { |
| 65 | s.mu.Lock() |
| 66 | if s.closed { |
| 67 | s.mu.Unlock() |
| 68 | return nil |
| 69 | } |
| 70 | s.closed = true |
| 71 | closers := s.closers |
| 72 | s.closers = nil |
| 73 | s.mu.Unlock() |
| 74 | |
| 75 | var errs []error |
| 76 | for i := len(closers) - 1; i >= 0; i-- { |
| 77 | if err := closers[i].Close(); err != nil { |
| 78 | errs = append(errs, err) |
| 79 | } |
| 80 | } |
| 81 | return errors.Join(errs...) |
| 82 | } |
| 83 | |
| 84 | // CloseIfGeneration closes the set only when gen matches the generation it |
| 85 | // was built for, and reports whether the close ran. A wrong generation means |
| 86 | // the caller's snapshot is stale and these resources already belong to a |
| 87 | // newer runtime, so they are left untouched. Close errors are intentionally |
| 88 | // collapsed into the bool: this is the fire-and-forget cleanup path — callers |
| 89 | // that need error detail call Close themselves. |
| 90 | func (s *RuntimeSet) CloseIfGeneration(gen uint64) bool { |
| 91 | s.mu.Lock() |
| 92 | matched := gen == s.generation |
| 93 | s.mu.Unlock() |
| 94 | if !matched { |
| 95 | return false |
| 96 | } |
| 97 | _ = s.Close() |
| 98 | return true |
| 99 | } |
| 100 | |
| 101 | // Closed reports whether Close has run. |
| 102 | func (s *RuntimeSet) Closed() bool { |
| 103 | s.mu.Lock() |
| 104 | defer s.mu.Unlock() |
| 105 | return s.closed |
| 106 | } |
| 107 |