| 1 | package extension |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "errors" |
| 6 | "fmt" |
| 7 | "io" |
| 8 | "slices" |
| 9 | "sync" |
| 10 | "time" |
| 11 | ) |
| 12 | |
| 13 | // EffectClass classifies how an effect may be undone during scope disposal. |
| 14 | type EffectClass uint8 |
| 15 | |
| 16 | const ( |
| 17 | // Reversible effects can be fully undone by Dispose. |
| 18 | Reversible EffectClass = iota |
| 19 | // Cancelable effects stop further work; Dispose waits for completion. |
| 20 | Cancelable |
| 21 | // Compensatable effects may run Compensate when Dispose cannot reverse them. |
| 22 | Compensatable |
| 23 | // Irreversible effects record a receipt only; never claim rollback success. |
| 24 | Irreversible |
| 25 | ) |
| 26 | |
| 27 | // Effect is one live resource owned by an EffectScope generation. |
| 28 | type Effect struct { |
| 29 | ID string |
| 30 | Owner string |
| 31 | Component string |
| 32 | Class EffectClass |
| 33 | Dispose func(context.Context) error |
| 34 | Compensate func(context.Context) error |
| 35 | } |
| 36 | |
| 37 | // EffectReceipt records irreversible or compensatable work for recovery. |
| 38 | type EffectReceipt struct { |
| 39 | ID string `json:"id"` |
| 40 | Owner string `json:"owner"` |
| 41 | Generation uint64 `json:"generation"` |
| 42 | Component string `json:"component,omitempty"` |
| 43 | Class EffectClass `json:"class"` |
| 44 | StartedAt time.Time `json:"startedAt"` |
| 45 | CompletedAt time.Time `json:"completedAt,omitempty"` |
| 46 | CompensationStatus string `json:"compensationStatus,omitempty"` |
| 47 | Error string `json:"error,omitempty"` |
| 48 | } |
| 49 | |
| 50 | // EffectScope tracks live resources for one runtime generation. |
| 51 | type EffectScope interface { |
| 52 | Track(Effect) error |
| 53 | TrackCloser(id string, c io.Closer) error |
| 54 | Dispose(context.Context) error |
| 55 | Generation() uint64 |
| 56 | Receipts() []EffectReceipt |
| 57 | Closed() bool |
| 58 | } |
| 59 | |
| 60 | // LiveScope is the default EffectScope implementation: reverse-order, |
| 61 | // once-only dispose with generation identity and receipt aggregation. |
| 62 | type LiveScope struct { |
| 63 | mu sync.Mutex |
| 64 | generation uint64 |
| 65 | effects []trackedEffect |
| 66 | receipts []EffectReceipt |
| 67 | closed bool |
| 68 | seen map[string]struct{} |
| 69 | } |
| 70 | |
| 71 | type trackedEffect struct { |
| 72 | effect Effect |
| 73 | disposed bool |
| 74 | started time.Time |
| 75 | } |
| 76 | |
| 77 | // NewEffectScope returns an empty scope bound to generation. |
| 78 | func NewEffectScope(generation uint64) *LiveScope { |
| 79 | return &LiveScope{ |
| 80 | generation: generation, |
| 81 | seen: make(map[string]struct{}), |
| 82 | } |
| 83 | } |
| 84 | |
| 85 | // Generation returns the bound runtime generation. |
| 86 | func (s *LiveScope) Generation() uint64 { return s.generation } |
| 87 | |
| 88 | // Closed reports whether Dispose has completed. |
| 89 | func (s *LiveScope) Closed() bool { |
| 90 | s.mu.Lock() |
| 91 | defer s.mu.Unlock() |
| 92 | return s.closed |
| 93 | } |
| 94 | |
| 95 | // Receipts returns a copy of recorded effect receipts. |
| 96 | func (s *LiveScope) Receipts() []EffectReceipt { |
| 97 | s.mu.Lock() |
| 98 | defer s.mu.Unlock() |
| 99 | out := make([]EffectReceipt, len(s.receipts)) |
| 100 | copy(out, s.receipts) |
| 101 | return out |
| 102 | } |
| 103 | |
| 104 | // Track registers an effect. If the scope is already closed the effect is |
| 105 | // disposed immediately so activation races cannot leak resources. |
| 106 | func (s *LiveScope) Track(e Effect) error { |
| 107 | if e.ID == "" { |
| 108 | return fmt.Errorf("extension: effect id is required") |
| 109 | } |
| 110 | if e.Dispose == nil && e.Class != Irreversible { |
| 111 | return fmt.Errorf("extension: effect %q requires Dispose", e.ID) |
| 112 | } |
| 113 | s.mu.Lock() |
| 114 | if _, dup := s.seen[e.ID]; dup { |
| 115 | s.mu.Unlock() |
| 116 | return fmt.Errorf("extension: duplicate effect id %q", e.ID) |
| 117 | } |
| 118 | if s.closed { |
| 119 | s.mu.Unlock() |
| 120 | return disposeNow(context.Background(), e) |
| 121 | } |
| 122 | s.seen[e.ID] = struct{}{} |
| 123 | s.effects = append(s.effects, trackedEffect{effect: e, started: time.Now().UTC()}) |
| 124 | if e.Class == Irreversible || e.Class == Compensatable { |
| 125 | s.receipts = append(s.receipts, EffectReceipt{ |
| 126 | ID: e.ID, |
| 127 | Owner: e.Owner, |
| 128 | Generation: s.generation, |
| 129 | Component: e.Component, |
| 130 | Class: e.Class, |
| 131 | StartedAt: time.Now().UTC(), |
| 132 | }) |
| 133 | } |
| 134 | s.mu.Unlock() |
| 135 | return nil |
| 136 | } |
| 137 | |
| 138 | // TrackCloser registers an io.Closer as a reversible effect. |
| 139 | func (s *LiveScope) TrackCloser(id string, c io.Closer) error { |
| 140 | if c == nil { |
| 141 | return nil |
| 142 | } |
| 143 | if id == "" { |
| 144 | id = fmt.Sprintf("closer-%p", c) |
| 145 | } |
| 146 | return s.Track(Effect{ |
| 147 | ID: id, |
| 148 | Class: Reversible, |
| 149 | Dispose: func(context.Context) error { |
| 150 | return c.Close() |
| 151 | }, |
| 152 | }) |
| 153 | } |
| 154 | |
| 155 | // Dispose releases every tracked effect in reverse registration order. It is |
| 156 | // idempotent. Individual dispose errors are joined and do not skip remaining |
| 157 | // effects. Cancelable dispose functions receive the caller's context so they |
| 158 | // can wait for background work. |
| 159 | func (s *LiveScope) Dispose(ctx context.Context) error { |
| 160 | if ctx == nil { |
| 161 | ctx = context.Background() |
| 162 | } |
| 163 | s.mu.Lock() |
| 164 | if s.closed { |
| 165 | s.mu.Unlock() |
| 166 | return nil |
| 167 | } |
| 168 | s.closed = true |
| 169 | effects := s.effects |
| 170 | s.effects = nil |
| 171 | s.mu.Unlock() |
| 172 | |
| 173 | var errs []error |
| 174 | for _, te := range slices.Backward(effects) { |
| 175 | if te.disposed { |
| 176 | continue |
| 177 | } |
| 178 | if err := disposeTracked(ctx, s, te); err != nil { |
| 179 | errs = append(errs, err) |
| 180 | } |
| 181 | } |
| 182 | return errors.Join(errs...) |
| 183 | } |
| 184 | |
| 185 | func disposeTracked(ctx context.Context, s *LiveScope, te trackedEffect) error { |
| 186 | e := te.effect |
| 187 | var disposeErr error |
| 188 | if e.Dispose != nil { |
| 189 | disposeErr = e.Dispose(ctx) |
| 190 | } |
| 191 | compStatus := "" |
| 192 | if e.Class == Compensatable && e.Compensate != nil { |
| 193 | if cerr := e.Compensate(ctx); cerr != nil { |
| 194 | compStatus = "failed" |
| 195 | disposeErr = errors.Join(disposeErr, fmt.Errorf("compensate %s: %w", e.ID, cerr)) |
| 196 | } else { |
| 197 | compStatus = "applied" |
| 198 | } |
| 199 | } |
| 200 | if e.Class == Irreversible { |
| 201 | // Cancellation/dispose never means the external action was undone. |
| 202 | compStatus = "not_applicable" |
| 203 | } |
| 204 | s.mu.Lock() |
| 205 | for i := range s.receipts { |
| 206 | if s.receipts[i].ID == e.ID && s.receipts[i].CompletedAt.IsZero() { |
| 207 | s.receipts[i].CompletedAt = time.Now().UTC() |
| 208 | s.receipts[i].CompensationStatus = compStatus |
| 209 | if disposeErr != nil { |
| 210 | s.receipts[i].Error = disposeErr.Error() |
| 211 | } |
| 212 | break |
| 213 | } |
| 214 | } |
| 215 | s.mu.Unlock() |
| 216 | if disposeErr != nil { |
| 217 | return fmt.Errorf("dispose %s: %w", e.ID, disposeErr) |
| 218 | } |
| 219 | return nil |
| 220 | } |
| 221 | |
| 222 | func disposeNow(ctx context.Context, e Effect) error { |
| 223 | var errs []error |
| 224 | if e.Dispose != nil { |
| 225 | if err := e.Dispose(ctx); err != nil { |
| 226 | errs = append(errs, err) |
| 227 | } |
| 228 | } |
| 229 | if e.Class == Compensatable && e.Compensate != nil { |
| 230 | if err := e.Compensate(ctx); err != nil { |
| 231 | errs = append(errs, err) |
| 232 | } |
| 233 | } |
| 234 | return errors.Join(errs...) |
| 235 | } |
| 236 | |
| 237 | var _ EffectScope = (*LiveScope)(nil) |
| 238 |