Skip to content

Commit a4bba72

Browse files
authored
feat: Send incremental state as checkpoints after the resources it covers (#2632)
#### Summary Adds source checkpoints and destination flushes, so the CLI can persist state only after its resources commit. Sources opt in with `state.ClientOptions{Checkpoints: res}` and `NewConnectedClientWithOptions`, destinations with `plugin.WithFlushSupport()`. `state.NewClientWithOptions` returns an error when `Checkpoints` is set, as it cannot tell whether the CLI persists checkpoints. With `StrategyShuffleQueue`, checkpointed state is held until the scheduler's `Sync` completes. New write API: `message.WriteFlush`, `message.FailedTablesError` and `message.FailedTables`, `writers.WithFailedTables` for custom writers, and `SkipFlush` in the writer test suite. Behaviour change: with checkpoints on, a failed batch write is reported on the next flush, with the tables it failed for, so the CLI stops saving state for that sync instead of only logging the failure. After a cancelled sync, including a quota stop, `Flush` returns `state.ErrSyncAborted` for keys set in table resolvers and does not save them, as the scheduler may have dropped their rows. `scheduler.Sync` now returns the cancel cause, not nil, when the context is cancelled while the batcher closes. Also fixes `NewConnectedClientWithOptions` with a zero `MaxMsgSizeInBytes`, which set a 0-byte gRPC message limit. Uses the protocol from plugin-pb-go v1.28.2.
1 parent 05070de commit a4bba72

35 files changed

Lines changed: 4940 additions & 144 deletions

‎internal/checkpoint/checkpoint.go‎

Lines changed: 387 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,387 @@
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

Comments
 (0)