|
| 1 | +// Package checkpoint lets state clients order their checkpoints after the resources they summarize. |
| 2 | +package checkpoint |
| 3 | + |
| 4 | +import ( |
| 5 | + "context" |
| 6 | + "errors" |
| 7 | + "slices" |
| 8 | + "sync" |
| 9 | + "sync/atomic" |
| 10 | + |
| 11 | + "github.com/cloudquery/plugin-sdk/v4/message" |
| 12 | + "github.com/cloudquery/plugin-sdk/v4/schema" |
| 13 | +) |
| 14 | + |
| 15 | +// ErrSyncAborted is returned for checkpoints that were dropped because the sync was cancelled. |
| 16 | +var ErrSyncAborted = errors.New("sync aborted before the checkpoint could be sent") |
| 17 | + |
| 18 | +// Scope is the part of a sync that a table resolver drives: the resources it produces and their relations. |
| 19 | +type Scope interface { |
| 20 | + // Checkpoint passes msg to Deliver after every resource the scope produced before the call is sent, or to Drop if |
| 21 | + // the sync is cancelled first. It may return before that. |
| 22 | + Checkpoint(ctx context.Context, msg message.SyncMessage) error |
| 23 | +} |
| 24 | + |
| 25 | +// Ordering is implemented by scopes that never send a checkpoint before an unsent one they already hold. |
| 26 | +type Ordering interface { |
| 27 | + // SendsAfter reports whether a checkpoint passed to the scope now will be sent no earlier than delivery. |
| 28 | + SendsAfter(delivery *Delivery) bool |
| 29 | +} |
| 30 | + |
| 31 | +// Collapsing is implemented by scopes that can become equivalent to a scope they share with others, such as a table |
| 32 | +// resolver's scope once the resolver ends. |
| 33 | +type Collapsing interface { |
| 34 | + // Collapsed returns the shared scope, or the scope itself. |
| 35 | + Collapsed() Scope |
| 36 | +} |
| 37 | + |
| 38 | +type scopeKey struct{} |
| 39 | + |
| 40 | +func WithScope(ctx context.Context, scope Scope) context.Context { |
| 41 | + return context.WithValue(ctx, scopeKey{}, scope) |
| 42 | +} |
| 43 | + |
| 44 | +// FromContext returns the scope ctx belongs to, or nil outside table resolvers. |
| 45 | +func FromContext(ctx context.Context) Scope { |
| 46 | + scope, _ := ctx.Value(scopeKey{}).(Scope) |
| 47 | + return scope |
| 48 | +} |
| 49 | + |
| 50 | +// Deliver sends msg on res. A *Barrier is marked as reached instead, and a *Tracked message is sent and marked as sent. |
| 51 | +func Deliver(ctx context.Context, res chan<- message.SyncMessage, msg message.SyncMessage) error { |
| 52 | + switch m := msg.(type) { |
| 53 | + case *Barrier: |
| 54 | + return m.reached(ctx, res) |
| 55 | + case *Tracked: |
| 56 | + if m.Delivery.Aborted() { |
| 57 | + return nil |
| 58 | + } |
| 59 | + if err := Deliver(ctx, res, m.Message); err != nil { |
| 60 | + m.Delivery.Abort() |
| 61 | + return err |
| 62 | + } |
| 63 | + return m.Delivery.Finish(ctx, res) |
| 64 | + default: |
| 65 | + select { |
| 66 | + case res <- msg: |
| 67 | + return nil |
| 68 | + case <-ctx.Done(): |
| 69 | + return context.Cause(ctx) |
| 70 | + } |
| 71 | + } |
| 72 | +} |
| 73 | + |
| 74 | +// Drop marks msg as never to be sent: a dropped *Tracked message, or the message a dropped *Barrier leads to, is |
| 75 | +// aborted, and so are the checkpoints waiting for it. |
| 76 | +func Drop(msg message.SyncMessage) { |
| 77 | + switch m := msg.(type) { |
| 78 | + case *Barrier: |
| 79 | + m.abort() |
| 80 | + case *Tracked: |
| 81 | + m.Delivery.Abort() |
| 82 | + } |
| 83 | +} |
| 84 | + |
| 85 | +// Barrier travels through scopes and sent checkpoints, and calls reached each time one of them gets to it. |
| 86 | +type Barrier struct { |
| 87 | + reached func(ctx context.Context, res chan<- message.SyncMessage) error |
| 88 | + abort func() |
| 89 | +} |
| 90 | + |
| 91 | +func (*Barrier) GetTable() *schema.Table { return nil } |
| 92 | +func (*Barrier) IsSyncMessage() bool { return true } |
| 93 | + |
| 94 | +// Tracked is a message that marks Delivery once Deliver sends it. |
| 95 | +type Tracked struct { |
| 96 | + Message message.SyncMessage |
| 97 | + Delivery *Delivery |
| 98 | +} |
| 99 | + |
| 100 | +func (*Tracked) GetTable() *schema.Table { return nil } |
| 101 | +func (*Tracked) IsSyncMessage() bool { return true } |
| 102 | + |
| 103 | +// Delivery tracks whether a checkpoint was sent or aborted, so that later checkpoints of the same keys can wait for it. |
| 104 | +type Delivery struct { |
| 105 | + mu sync.Mutex |
| 106 | + done bool |
| 107 | + aborted bool |
| 108 | + waiting []*Barrier |
| 109 | + // set when the checkpoint was merged into another one, which is then sent in its place |
| 110 | + mergedInto *Delivery |
| 111 | + // for a checkpoint that others are merged into: the scope holding it |
| 112 | + heldBy Scope |
| 113 | + // set when the checkpoint is aborted |
| 114 | + abortFlags map[*atomic.Bool]struct{} |
| 115 | +} |
| 116 | + |
| 117 | +// NewHeld returns a Delivery for a checkpoint that holder merges other checkpoints into. |
| 118 | +func NewHeld(holder Scope) *Delivery { |
| 119 | + return &Delivery{heldBy: holder} |
| 120 | +} |
| 121 | + |
| 122 | +// root returns the checkpoint s is sent with. |
| 123 | +func (s *Delivery) root() *Delivery { |
| 124 | + for { |
| 125 | + s.mu.Lock() |
| 126 | + into := s.mergedInto |
| 127 | + s.mu.Unlock() |
| 128 | + if into == nil { |
| 129 | + return s |
| 130 | + } |
| 131 | + s = into |
| 132 | + } |
| 133 | +} |
| 134 | + |
| 135 | +// lockRoot returns the checkpoint s is sent with, locked. It is still the root once locked, as MergeInto runs under the |
| 136 | +// lock of the checkpoint it merges. |
| 137 | +func (s *Delivery) lockRoot() *Delivery { |
| 138 | + for { |
| 139 | + r := s.root() |
| 140 | + r.mu.Lock() |
| 141 | + if r.mergedInto == nil { |
| 142 | + return r |
| 143 | + } |
| 144 | + r.mu.Unlock() |
| 145 | + } |
| 146 | +} |
| 147 | + |
| 148 | +// Delivered reports whether the checkpoint was sent. |
| 149 | +func (s *Delivery) Delivered() bool { |
| 150 | + r := s.lockRoot() |
| 151 | + defer r.mu.Unlock() |
| 152 | + return r.done |
| 153 | +} |
| 154 | + |
| 155 | +// Aborted reports whether the checkpoint was dropped. |
| 156 | +func (s *Delivery) Aborted() bool { |
| 157 | + r := s.lockRoot() |
| 158 | + defer r.mu.Unlock() |
| 159 | + return r.aborted |
| 160 | +} |
| 161 | + |
| 162 | +func (s *Delivery) holder() Scope { |
| 163 | + r := s.lockRoot() |
| 164 | + defer r.mu.Unlock() |
| 165 | + return r.heldBy |
| 166 | +} |
| 167 | + |
| 168 | +// NotifyAbort sets flag if the checkpoint is aborted, or now if it already was. |
| 169 | +func (s *Delivery) NotifyAbort(flag *atomic.Bool) { |
| 170 | + r := s.lockRoot() |
| 171 | + defer r.mu.Unlock() |
| 172 | + if r.aborted { |
| 173 | + flag.Store(true) |
| 174 | + return |
| 175 | + } |
| 176 | + if r.abortFlags == nil { |
| 177 | + r.abortFlags = make(map[*atomic.Bool]struct{}) |
| 178 | + } |
| 179 | + r.abortFlags[flag] = struct{}{} |
| 180 | +} |
| 181 | + |
| 182 | +// MergedInto reports whether s is sent with into. |
| 183 | +func (s *Delivery) MergedInto(into *Delivery) bool { |
| 184 | + return s != into && s.root() == into |
| 185 | +} |
| 186 | + |
| 187 | +// MergeInto marks s as sent once into is sent, and aborted if into is aborted. into must not be sent yet. |
| 188 | +func (s *Delivery) MergeInto(into *Delivery) { |
| 189 | + s.mu.Lock() |
| 190 | + waiting, flags := s.waiting, s.abortFlags |
| 191 | + s.waiting, s.abortFlags = nil, nil |
| 192 | + s.mergedInto = into |
| 193 | + s.mu.Unlock() |
| 194 | + for _, barrier := range waiting { |
| 195 | + into.wait(barrier) |
| 196 | + } |
| 197 | + for flag := range flags { |
| 198 | + into.NotifyAbort(flag) |
| 199 | + } |
| 200 | +} |
| 201 | + |
| 202 | +// wait makes barrier wait for s. It reports true, without waiting, if s was sent. |
| 203 | +func (s *Delivery) wait(barrier *Barrier) bool { |
| 204 | + r := s.lockRoot() |
| 205 | + switch { |
| 206 | + case r.done: |
| 207 | + r.mu.Unlock() |
| 208 | + return true |
| 209 | + case r.aborted: |
| 210 | + r.mu.Unlock() |
| 211 | + barrier.abort() |
| 212 | + return false |
| 213 | + default: |
| 214 | + r.waiting = append(r.waiting, barrier) |
| 215 | + r.mu.Unlock() |
| 216 | + return false |
| 217 | + } |
| 218 | +} |
| 219 | + |
| 220 | +// Finish marks s as sent, and releases the checkpoints waiting for it. |
| 221 | +func (s *Delivery) Finish(ctx context.Context, res chan<- message.SyncMessage) error { |
| 222 | + s.mu.Lock() |
| 223 | + if s.done || s.aborted { |
| 224 | + s.mu.Unlock() |
| 225 | + return nil |
| 226 | + } |
| 227 | + s.done = true |
| 228 | + waiting := s.waiting |
| 229 | + s.waiting, s.abortFlags = nil, nil |
| 230 | + s.mu.Unlock() |
| 231 | + var errs []error |
| 232 | + for _, barrier := range waiting { |
| 233 | + if err := barrier.reached(ctx, res); err != nil { |
| 234 | + errs = append(errs, err) |
| 235 | + } |
| 236 | + } |
| 237 | + return errors.Join(errs...) |
| 238 | +} |
| 239 | + |
| 240 | +// Abort marks s as dropped, and aborts the checkpoints waiting for it. |
| 241 | +func (s *Delivery) Abort() { |
| 242 | + r := s.lockRoot() |
| 243 | + if r.done || r.aborted { |
| 244 | + r.mu.Unlock() |
| 245 | + return |
| 246 | + } |
| 247 | + r.aborted = true |
| 248 | + waiting, flags := r.waiting, r.abortFlags |
| 249 | + r.waiting, r.abortFlags = nil, nil |
| 250 | + r.mu.Unlock() |
| 251 | + for flag := range flags { |
| 252 | + flag.Store(true) |
| 253 | + } |
| 254 | + for _, barrier := range waiting { |
| 255 | + barrier.abort() |
| 256 | + } |
| 257 | +} |
| 258 | + |
| 259 | +// Schedule passes msg to Deliver once every one of scope's resources produced before the call is sent, and every |
| 260 | +// checkpoint in after is sent, which may be after it returns. Without a scope or pending checkpoints to wait for, it |
| 261 | +// sends msg on res. The returned Delivery tracks msg. If a checkpoint in after was aborted, Schedule returns |
| 262 | +// ErrSyncAborted. |
| 263 | +func Schedule(ctx context.Context, res chan<- message.SyncMessage, scope Scope, after []*Delivery, msg message.SyncMessage) (*Delivery, error) { |
| 264 | + pending := make([]*Delivery, 0, len(after)) |
| 265 | + for _, previous := range after { |
| 266 | + switch { |
| 267 | + case previous.Aborted(): |
| 268 | + return nil, ErrSyncAborted |
| 269 | + case !previous.Delivered(): |
| 270 | + pending = append(pending, previous) |
| 271 | + } |
| 272 | + } |
| 273 | + if scope == nil { |
| 274 | + // waiting only for checkpoints one scope holds: going through that scope sends no earlier, and lets it merge msg |
| 275 | + scope = commonHolder(pending) |
| 276 | + } |
| 277 | + if ordering, ok := scope.(Ordering); ok { |
| 278 | + pending = slices.DeleteFunc(pending, ordering.SendsAfter) |
| 279 | + } |
| 280 | + |
| 281 | + sent := &Delivery{} |
| 282 | + tracked := &Tracked{Message: msg, Delivery: sent} |
| 283 | + scopes := flatten(scope) |
| 284 | + if len(pending) == 0 { |
| 285 | + switch len(scopes) { |
| 286 | + case 0: |
| 287 | + return sent, Deliver(ctx, res, tracked) |
| 288 | + case 1: |
| 289 | + return sent, scopes[0].Checkpoint(ctx, tracked) |
| 290 | + } |
| 291 | + } |
| 292 | + |
| 293 | + var remaining atomic.Int64 |
| 294 | + remaining.Store(int64(len(scopes) + len(pending))) |
| 295 | + barrier := &Barrier{ |
| 296 | + reached: func(ctx context.Context, res chan<- message.SyncMessage) error { |
| 297 | + if remaining.Add(-1) > 0 { |
| 298 | + return nil |
| 299 | + } |
| 300 | + return Deliver(ctx, res, tracked) |
| 301 | + }, |
| 302 | + abort: sent.Abort, |
| 303 | + } |
| 304 | + var errs []error |
| 305 | + for _, previous := range pending { |
| 306 | + if previous.wait(barrier) { |
| 307 | + if err := barrier.reached(ctx, res); err != nil { |
| 308 | + errs = append(errs, err) |
| 309 | + } |
| 310 | + } |
| 311 | + } |
| 312 | + for _, s := range scopes { |
| 313 | + if err := s.Checkpoint(ctx, barrier); err != nil { |
| 314 | + errs = append(errs, err) |
| 315 | + } |
| 316 | + } |
| 317 | + if err := errors.Join(errs...); err != nil { |
| 318 | + sent.Abort() |
| 319 | + return nil, err |
| 320 | + } |
| 321 | + return sent, nil |
| 322 | +} |
| 323 | + |
| 324 | +func commonHolder(pending []*Delivery) Scope { |
| 325 | + var holder Scope |
| 326 | + for _, sent := range pending { |
| 327 | + h := sent.holder() |
| 328 | + if h == nil || (holder != nil && h != holder) { |
| 329 | + return nil |
| 330 | + } |
| 331 | + holder = h |
| 332 | + } |
| 333 | + return holder |
| 334 | +} |
| 335 | + |
| 336 | +// Collapse returns the shared scope a Collapsing scope collapsed into, or scope itself. |
| 337 | +func Collapse(scope Scope) Scope { |
| 338 | + if c, ok := scope.(Collapsing); ok { |
| 339 | + return c.Collapsed() |
| 340 | + } |
| 341 | + return scope |
| 342 | +} |
| 343 | + |
| 344 | +// Join returns a scope whose checkpoints follow the resources of all the given scopes. It skips nil scopes, merges |
| 345 | +// scopes that collapsed into the same one, and returns nil if all are nil. |
| 346 | +func Join(scopes ...Scope) Scope { |
| 347 | + seen := make(map[Scope]struct{}, len(scopes)) |
| 348 | + var unique []Scope |
| 349 | + for _, scope := range scopes { |
| 350 | + for _, s := range flatten(scope) { |
| 351 | + s = Collapse(s) |
| 352 | + if _, ok := seen[s]; ok { |
| 353 | + continue |
| 354 | + } |
| 355 | + seen[s] = struct{}{} |
| 356 | + unique = append(unique, s) |
| 357 | + } |
| 358 | + } |
| 359 | + switch len(unique) { |
| 360 | + case 0: |
| 361 | + return nil |
| 362 | + case 1: |
| 363 | + return unique[0] |
| 364 | + default: |
| 365 | + return &joined{scopes: unique} |
| 366 | + } |
| 367 | +} |
| 368 | + |
| 369 | +type joined struct { |
| 370 | + scopes []Scope |
| 371 | +} |
| 372 | + |
| 373 | +func (j *joined) Checkpoint(ctx context.Context, msg message.SyncMessage) error { |
| 374 | + _, err := Schedule(ctx, nil, j, nil, msg) |
| 375 | + return err |
| 376 | +} |
| 377 | + |
| 378 | +func flatten(scope Scope) []Scope { |
| 379 | + switch s := scope.(type) { |
| 380 | + case nil: |
| 381 | + return nil |
| 382 | + case *joined: |
| 383 | + return s.scopes |
| 384 | + default: |
| 385 | + return []Scope{s} |
| 386 | + } |
| 387 | +} |
0 commit comments