diff --git a/internal/checkpoint/checkpoint.go b/internal/checkpoint/checkpoint.go new file mode 100644 index 0000000000..6cfaae3dd8 --- /dev/null +++ b/internal/checkpoint/checkpoint.go @@ -0,0 +1,386 @@ +// Package checkpoint lets state clients order their checkpoints after the resources they summarize. +package checkpoint + +import ( + "context" + "errors" + "slices" + "sync" + "sync/atomic" + + "github.com/cloudquery/plugin-sdk/v4/message" + "github.com/cloudquery/plugin-sdk/v4/schema" +) + +// ErrSyncAborted is returned for checkpoints that were dropped because the sync was cancelled. +var ErrSyncAborted = errors.New("sync aborted before the checkpoint could be sent") + +// Scope is the part of a sync that a table resolver drives: the resources it produces and their relations. +type Scope interface { + // Checkpoint passes msg to Deliver after every resource the scope produced before the call is sent, or to Drop if + // the sync is cancelled first. It may return before that. + Checkpoint(ctx context.Context, msg message.SyncMessage) error +} + +// Covering is implemented by scopes that never send a checkpoint before an unsent one they already hold. +type Covering interface { + // Covers reports whether a checkpoint passed to the scope now will be sent no earlier than sent. + Covers(sent *Sent) bool +} + +// Collapsing is implemented by scopes that can become equivalent to a scope they share with others, such as a table +// resolver's scope once the resolver ends. +type Collapsing interface { + // Collapsed returns the shared scope, or the scope itself. + Collapsed() Scope +} + +type scopeKey struct{} + +func WithScope(ctx context.Context, scope Scope) context.Context { + return context.WithValue(ctx, scopeKey{}, scope) +} + +// FromContext returns the scope ctx belongs to, or nil outside table resolvers. +func FromContext(ctx context.Context) Scope { + scope, _ := ctx.Value(scopeKey{}).(Scope) + return scope +} + +// Deliver sends msg on res. A *Barrier is marked as reached instead, and a *Tracked message is sent and marked as sent. +func Deliver(ctx context.Context, res chan<- message.SyncMessage, msg message.SyncMessage) error { + switch m := msg.(type) { + case *Barrier: + return m.reached(ctx, res) + case *Tracked: + if m.Sent.Aborted() { + return nil + } + if err := Deliver(ctx, res, m.Message); err != nil { + m.Sent.Abort() + return err + } + return m.Sent.Finish(ctx, res) + default: + select { + case res <- msg: + return nil + case <-ctx.Done(): + return context.Cause(ctx) + } + } +} + +// Drop marks msg as never to be sent: a dropped *Tracked message, or the message a dropped *Barrier leads to, is +// aborted, and so are the checkpoints waiting for it. +func Drop(msg message.SyncMessage) { + switch m := msg.(type) { + case *Barrier: + m.abort() + case *Tracked: + m.Sent.Abort() + } +} + +// Barrier travels through scopes and sent checkpoints, and calls reached each time one of them gets to it. +type Barrier struct { + reached func(ctx context.Context, res chan<- message.SyncMessage) error + abort func() +} + +func (*Barrier) GetTable() *schema.Table { return nil } +func (*Barrier) IsSyncMessage() bool { return true } + +// Tracked is a message that marks Sent once Deliver sends it. +type Tracked struct { + Message message.SyncMessage + Sent *Sent +} + +func (*Tracked) GetTable() *schema.Table { return nil } +func (*Tracked) IsSyncMessage() bool { return true } + +// Sent tracks whether a checkpoint was sent or aborted, so that later checkpoints of the same keys can wait for it. +type Sent struct { + mu sync.Mutex + done bool + aborted bool + waiting []*Barrier + // set when the checkpoint was merged into another one, which is then sent in its place + mergedInto *Sent + // for a checkpoint that others are merged into: the scope holding it + heldBy Scope + // set when the checkpoint is aborted + abortFlags map[*atomic.Bool]struct{} +} + +// NewHeld returns a Sent for a checkpoint that holder merges other checkpoints into. +func NewHeld(holder Scope) *Sent { + return &Sent{heldBy: holder} +} + +// root returns the checkpoint s is sent with. +func (s *Sent) root() *Sent { + for { + s.mu.Lock() + into := s.mergedInto + s.mu.Unlock() + if into == nil { + return s + } + s = into + } +} + +// lockRoot returns the checkpoint s is sent with, locked. It is still the root once locked, as MergeInto runs under the +// lock of the checkpoint it merges. +func (s *Sent) lockRoot() *Sent { + for { + r := s.root() + r.mu.Lock() + if r.mergedInto == nil { + return r + } + r.mu.Unlock() + } +} + +// Done reports whether the checkpoint was sent. +func (s *Sent) Done() bool { + r := s.lockRoot() + defer r.mu.Unlock() + return r.done +} + +// Aborted reports whether the checkpoint was dropped. +func (s *Sent) Aborted() bool { + r := s.lockRoot() + defer r.mu.Unlock() + return r.aborted +} + +func (s *Sent) holder() Scope { + r := s.lockRoot() + defer r.mu.Unlock() + return r.heldBy +} + +// NotifyAbort sets flag if the checkpoint is aborted, or now if it already was. +func (s *Sent) NotifyAbort(flag *atomic.Bool) { + r := s.lockRoot() + defer r.mu.Unlock() + if r.aborted { + flag.Store(true) + return + } + if r.abortFlags == nil { + r.abortFlags = make(map[*atomic.Bool]struct{}) + } + r.abortFlags[flag] = struct{}{} +} + +// MergedInto reports whether s is sent with into. +func (s *Sent) MergedInto(into *Sent) bool { + return s != into && s.root() == into +} + +// MergeInto marks s as sent once into is sent, and aborted if into is aborted. into must not be sent yet. +func (s *Sent) MergeInto(into *Sent) { + s.mu.Lock() + waiting, flags := s.waiting, s.abortFlags + s.waiting, s.abortFlags = nil, nil + s.mergedInto = into + s.mu.Unlock() + for _, barrier := range waiting { + into.wait(barrier) + } + for flag := range flags { + into.NotifyAbort(flag) + } +} + +// wait makes barrier wait for s. It reports true, without waiting, if s was sent. +func (s *Sent) wait(barrier *Barrier) bool { + r := s.lockRoot() + switch { + case r.done: + r.mu.Unlock() + return true + case r.aborted: + r.mu.Unlock() + barrier.abort() + return false + default: + r.waiting = append(r.waiting, barrier) + r.mu.Unlock() + return false + } +} + +// Finish marks s as sent, and releases the checkpoints waiting for it. +func (s *Sent) Finish(ctx context.Context, res chan<- message.SyncMessage) error { + s.mu.Lock() + if s.done || s.aborted { + s.mu.Unlock() + return nil + } + s.done = true + waiting := s.waiting + s.waiting, s.abortFlags = nil, nil + s.mu.Unlock() + var errs []error + for _, barrier := range waiting { + if err := barrier.reached(ctx, res); err != nil { + errs = append(errs, err) + } + } + return errors.Join(errs...) +} + +// Abort marks s as dropped, and aborts the checkpoints waiting for it. +func (s *Sent) Abort() { + r := s.lockRoot() + if r.done || r.aborted { + r.mu.Unlock() + return + } + r.aborted = true + waiting, flags := r.waiting, r.abortFlags + r.waiting, r.abortFlags = nil, nil + r.mu.Unlock() + for flag := range flags { + flag.Store(true) + } + for _, barrier := range waiting { + barrier.abort() + } +} + +// Send passes msg to Deliver once every one of scope's resources produced before the call is sent, and every +// checkpoint in after is sent. Without a scope or pending checkpoints to wait for, it sends msg on res. +// The returned Sent tracks msg. If a checkpoint in after was aborted, Send returns ErrSyncAborted. +func Send(ctx context.Context, res chan<- message.SyncMessage, scope Scope, after []*Sent, msg message.SyncMessage) (*Sent, error) { + pending := make([]*Sent, 0, len(after)) + for _, previous := range after { + switch { + case previous.Aborted(): + return nil, ErrSyncAborted + case !previous.Done(): + pending = append(pending, previous) + } + } + if scope == nil { + // waiting only for checkpoints one scope holds: going through that scope sends no earlier, and lets it merge msg + scope = commonHolder(pending) + } + if covering, ok := scope.(Covering); ok { + pending = slices.DeleteFunc(pending, covering.Covers) + } + + sent := &Sent{} + tracked := &Tracked{Message: msg, Sent: sent} + scopes := flatten(scope) + if len(pending) == 0 { + switch len(scopes) { + case 0: + return sent, Deliver(ctx, res, tracked) + case 1: + return sent, scopes[0].Checkpoint(ctx, tracked) + } + } + + var remaining atomic.Int64 + remaining.Store(int64(len(scopes) + len(pending))) + barrier := &Barrier{ + reached: func(ctx context.Context, res chan<- message.SyncMessage) error { + if remaining.Add(-1) > 0 { + return nil + } + return Deliver(ctx, res, tracked) + }, + abort: sent.Abort, + } + var errs []error + for _, previous := range pending { + if previous.wait(barrier) { + if err := barrier.reached(ctx, res); err != nil { + errs = append(errs, err) + } + } + } + for _, s := range scopes { + if err := s.Checkpoint(ctx, barrier); err != nil { + errs = append(errs, err) + } + } + if err := errors.Join(errs...); err != nil { + sent.Abort() + return nil, err + } + return sent, nil +} + +func commonHolder(pending []*Sent) Scope { + var holder Scope + for _, sent := range pending { + h := sent.holder() + if h == nil || (holder != nil && h != holder) { + return nil + } + holder = h + } + return holder +} + +// Collapse returns the shared scope a Collapsing scope collapsed into, or scope itself. +func Collapse(scope Scope) Scope { + if c, ok := scope.(Collapsing); ok { + return c.Collapsed() + } + return scope +} + +// Join returns a scope whose checkpoints follow the resources of all the given scopes. It skips nil scopes, merges +// scopes that collapsed into the same one, and returns nil if all are nil. +func Join(scopes ...Scope) Scope { + seen := make(map[Scope]struct{}, len(scopes)) + var unique []Scope + for _, scope := range scopes { + for _, s := range flatten(scope) { + s = Collapse(s) + if _, ok := seen[s]; ok { + continue + } + seen[s] = struct{}{} + unique = append(unique, s) + } + } + switch len(unique) { + case 0: + return nil + case 1: + return unique[0] + default: + return &joined{scopes: unique} + } +} + +type joined struct { + scopes []Scope +} + +func (j *joined) Checkpoint(ctx context.Context, msg message.SyncMessage) error { + _, err := Send(ctx, nil, j, nil, msg) + return err +} + +func flatten(scope Scope) []Scope { + switch s := scope.(type) { + case nil: + return nil + case *joined: + return s.scopes + default: + return []Scope{s} + } +} diff --git a/internal/checkpoint/checkpoint_test.go b/internal/checkpoint/checkpoint_test.go new file mode 100644 index 0000000000..535633f2b8 --- /dev/null +++ b/internal/checkpoint/checkpoint_test.go @@ -0,0 +1,196 @@ +package checkpoint + +import ( + "context" + "sync" + "sync/atomic" + "testing" + + "github.com/cloudquery/plugin-sdk/v4/message" + "github.com/stretchr/testify/require" +) + +// heldScope holds messages until the test delivers them, as a scope does until its resources are sent. +type heldScope struct { + msgs []message.SyncMessage +} + +func (s *heldScope) Checkpoint(_ context.Context, msg message.SyncMessage) error { + s.msgs = append(s.msgs, msg) + return nil +} + +func (s *heldScope) deliver(t *testing.T, res chan<- message.SyncMessage) { + t.Helper() + for _, msg := range s.msgs { + require.NoError(t, Deliver(context.Background(), res, msg)) + } + s.msgs = nil +} + +func marker(name string) message.SyncMessage { + return &message.SyncError{TableName: name} +} + +func received(res chan message.SyncMessage) []string { + var names []string + for len(res) > 0 { + names = append(names, (<-res).GetTable().Name) + } + return names +} + +func TestSendWaitsForScopesAndPreviousCheckpoints(t *testing.T) { + ctx := context.Background() + res := make(chan message.SyncMessage, 10) + a, b := &heldScope{}, &heldScope{} + + first, err := Send(ctx, res, a, nil, marker("first")) + require.NoError(t, err) + second, err := Send(ctx, res, b, []*Sent{first}, marker("second")) + require.NoError(t, err) + + b.deliver(t, res) + require.Empty(t, received(res), "second must wait for first") + require.False(t, second.Done()) + a.deliver(t, res) + require.Equal(t, []string{"first", "second"}, received(res)) + require.True(t, first.Done()) + require.True(t, second.Done()) + + // without a scope, a checkpoint waits only for previous ones + third, err := Send(ctx, res, nil, []*Sent{second}, marker("third")) + require.NoError(t, err) + require.True(t, third.Done()) + require.Equal(t, []string{"third"}, received(res)) +} + +func TestJoin(t *testing.T) { + ctx := context.Background() + res := make(chan message.SyncMessage, 10) + a, b := &heldScope{}, &heldScope{} + + require.Nil(t, Join(nil, nil)) + require.Same(t, a, Join(nil, a, a)) + joined := Join(Join(a, nil), b, a) + require.NoError(t, joined.Checkpoint(ctx, marker("joined"))) + + a.deliver(t, res) + require.Empty(t, received(res)) + b.deliver(t, res) + require.Equal(t, []string{"joined"}, received(res)) +} + +func TestMergeInto(t *testing.T) { + ctx := context.Background() + res := make(chan message.SyncMessage, 10) + a := &heldScope{} + + merged := &Sent{} + into := &Sent{} + merged.MergeInto(into) + _, err := Send(ctx, res, a, []*Sent{merged}, marker("after merged")) + require.NoError(t, err) + + a.deliver(t, res) + require.Empty(t, received(res)) + require.False(t, merged.Done()) + require.NoError(t, into.Finish(ctx, res)) + require.True(t, merged.Done()) + require.Equal(t, []string{"after merged"}, received(res)) +} + +// collapsingScope collapses into into once ended. +type collapsingScope struct { + heldScope + ended bool + into Scope +} + +func (s *collapsingScope) Collapsed() Scope { + if s.ended { + return s.into + } + return s +} + +func TestJoinCollapsesEndedScopes(t *testing.T) { + shared := &heldScope{} + a := &collapsingScope{ended: true, into: shared} + b := &collapsingScope{ended: true, into: shared} + running := &collapsingScope{into: shared} + + require.Same(t, shared, Join(a, b)) + joined, ok := Join(a, b, running).(*joined) + require.True(t, ok) + require.Len(t, joined.scopes, 2) +} + +func TestAbortReleasesLaterCheckpoints(t *testing.T) { + ctx := context.Background() + res := make(chan message.SyncMessage, 10) + a, b := &heldScope{}, &heldScope{} + + first, err := Send(ctx, res, a, nil, marker("first")) + require.NoError(t, err) + second, err := Send(ctx, res, b, []*Sent{first}, marker("second")) + require.NoError(t, err) + + for _, msg := range a.msgs { + Drop(msg) + } + require.True(t, first.Aborted()) + require.True(t, second.Aborted(), "a checkpoint waiting for an aborted one is aborted too") + b.deliver(t, res) + require.Empty(t, received(res)) + + _, err = Send(ctx, res, b, []*Sent{second}, marker("third")) + require.ErrorIs(t, err, ErrSyncAborted) +} + +func TestNotifyAbort(t *testing.T) { + var flag atomic.Bool + sent := &Sent{} + sent.NotifyAbort(&flag) + require.False(t, flag.Load()) + sent.Abort() + require.True(t, flag.Load()) + + var late atomic.Bool + sent.NotifyAbort(&late) + require.True(t, late.Load(), "a checkpoint aborted already sets the flag at once") + + var finished atomic.Bool + done := &Sent{} + done.NotifyAbort(&finished) + require.NoError(t, done.Finish(context.Background(), make(chan message.SyncMessage))) + done.Abort() + require.False(t, done.Aborted(), "a sent checkpoint is never aborted") + require.False(t, finished.Load()) +} + +func TestNotifyAbortFollowsMerges(t *testing.T) { + var flag atomic.Bool + held := NewHeld(&heldScope{}) + sent := &Sent{} + sent.NotifyAbort(&flag) + sent.MergeInto(held) + held.Abort() + require.True(t, flag.Load()) + require.True(t, sent.Aborted()) +} + +func TestConcurrentMergeKeepsAborts(t *testing.T) { + // with the root read before its lock is taken, a flag added while the checkpoint merges can be lost + for range 20_000 { + var flag atomic.Bool + held := NewHeld(&heldScope{}) + sent := &Sent{} + var wg sync.WaitGroup + wg.Go(func() { sent.MergeInto(held) }) + wg.Go(func() { sent.NotifyAbort(&flag) }) + wg.Wait() + held.Abort() + require.True(t, flag.Load()) + } +} diff --git a/internal/clients/state/v3/scope_set_test.go b/internal/clients/state/v3/scope_set_test.go new file mode 100644 index 0000000000..7273225b1b --- /dev/null +++ b/internal/clients/state/v3/scope_set_test.go @@ -0,0 +1,54 @@ +package state + +import ( + "context" + "testing" + + "github.com/cloudquery/plugin-sdk/v4/internal/checkpoint" + "github.com/cloudquery/plugin-sdk/v4/message" + "github.com/stretchr/testify/require" +) + +type liveScope struct{ id int } + +func (*liveScope) Checkpoint(context.Context, message.SyncMessage) error { return nil } + +// countingScope counts how often a compaction visits it. +type countingScope struct { + liveScope + visits *int +} + +func (s *countingScope) Collapsed() checkpoint.Scope { + *s.visits++ + return s +} + +type endedScope struct { + liveScope + into checkpoint.Scope +} + +func (s *endedScope) Collapsed() checkpoint.Scope { return s.into } + +func TestScopeSetCompactsOnlyOnceItDoubles(t *testing.T) { + const scopes = 100_000 + s := newScopeSet() + visits := 0 + for i := range scopes { + s.add(&countingScope{liveScope: liveScope{id: i}, visits: &visits}) + } + require.Len(t, s.scopes, scopes, "live scopes are kept") + // compactions at 64, 128, ..., 65536 visit 131008 scopes in all: the cost per add stays constant + require.Equal(t, 131008, visits) +} + +func TestScopeSetCollapsesEndedScopes(t *testing.T) { + shared := &liveScope{} + s := newScopeSet() + for i := range 100_000 { + s.add(&endedScope{liveScope: liveScope{id: i}, into: shared}) + } + require.Less(t, len(s.scopes), minCompactSize, "ended scopes collapse into the one they share") + require.Contains(t, s.scopes, checkpoint.Scope(shared)) +} diff --git a/internal/clients/state/v3/state.go b/internal/clients/state/v3/state.go index cdd9aae524..549e3e318e 100644 --- a/internal/clients/state/v3/state.go +++ b/internal/clients/state/v3/state.go @@ -3,14 +3,23 @@ package state import ( "bytes" "context" + "errors" + "fmt" "io" + "iter" + "maps" + "slices" + "strings" "sync" + "sync/atomic" "github.com/apache/arrow-go/v18/arrow" "github.com/apache/arrow-go/v18/arrow/array" "github.com/apache/arrow-go/v18/arrow/ipc" "github.com/apache/arrow-go/v18/arrow/memory" pb "github.com/cloudquery/plugin-pb-go/pb/plugin/v3" + "github.com/cloudquery/plugin-sdk/v4/internal/checkpoint" + "github.com/cloudquery/plugin-sdk/v4/message" "github.com/cloudquery/plugin-sdk/v4/schema" "google.golang.org/grpc" ) @@ -22,13 +31,31 @@ const ( ) type Client struct { - client pb.PluginClient - mem map[string]versionedValue - changes map[string]struct{} // changed keys + client pb.PluginClient + mem map[string]versionedValue + // changed keys. With checkpoints, each key has the scopes of its changes since the last flush, nil for changes made + // outside table resolvers. + changes map[string]*scopeSet + // the last checkpoint of each key, which the next checkpoint of the key must not overtake + sent map[string]*checkpoint.Sent + // set when one of the checkpoints in sent is aborted + abortSeen atomic.Bool mutex *sync.RWMutex schema *arrow.Schema versionedMode bool conn *grpc.ClientConn + // when set, Flush sends checkpoints instead of writing to the backend + checkpoints chan<- message.SyncMessage +} + +type Option func(*Client) + +// WithCheckpoints makes Flush send the changes as message.SyncCheckpoint on res instead of writing them to the +// backend. Changes made within a table resolver are sent after the resources it produced before the Flush. +func WithCheckpoints(res chan<- message.SyncMessage) Option { + return func(c *Client) { + c.checkpoints = res + } } type versionedValue struct { @@ -36,16 +63,20 @@ type versionedValue struct { version uint64 } -func NewClient(ctx context.Context, conn *grpc.ClientConn, tableName string) (*Client, error) { +func NewClient(ctx context.Context, conn *grpc.ClientConn, tableName string, opts ...Option) (*Client, error) { table := Table(tableName) c := &Client{ conn: conn, client: pb.NewPluginClient(conn), mem: make(map[string]versionedValue), - changes: make(map[string]struct{}), + changes: make(map[string]*scopeSet), + sent: make(map[string]*checkpoint.Sent), mutex: &sync.RWMutex{}, versionedMode: table.Column(versionColumn) != nil, } + for _, opt := range opts { + opt(c) + } sc := table.ToArrowSchema() c.schema = sc tableBytes, err := pb.SchemaToBytes(sc) @@ -131,39 +162,72 @@ func NewClient(ctx context.Context, conn *grpc.ClientConn, tableName string) (*C return c, nil } -func (c *Client) SetKey(_ context.Context, key string, value string) error { +func (c *Client) SetKey(ctx context.Context, key string, value string) error { c.mutex.Lock() defer c.mutex.Unlock() + scopes, pending := c.changes[key] if c.mem[key].value == value { + // the pending change summarizes this caller's resources too + if pending && c.checkpoints != nil { + scopes.add(checkpoint.FromContext(ctx)) + } return nil // don't update if the value is the same } c.mem[key] = versionedValue{ value: value, version: c.mem[key].version + 1, } - c.changes[key] = struct{}{} + if c.checkpoints == nil { + c.changes[key] = nil + return nil + } + if !pending { + scopes = newScopeSet() + c.changes[key] = scopes + } + scopes.add(checkpoint.FromContext(ctx)) return nil } +// minCompactSize is the size at which a scope set is first compacted. +const minCompactSize = 64 + +// scopeSet holds the scopes of a key's changes since the last flush. +type scopeSet struct { + scopes map[checkpoint.Scope]struct{} + // the size at which the set is compacted next + compactAt int +} + +func newScopeSet() *scopeSet { + return &scopeSet{scopes: make(map[checkpoint.Scope]struct{}), compactAt: minCompactSize} +} + +// add adds scope. Once the set reaches compactAt, it collapses the scopes of ended resolvers into the one they share, +// so that a key that many resolvers change keeps a small set. It compacts again only once the set doubles, so that a +// set of scopes that do not collapse costs constant time per add. +func (s *scopeSet) add(scope checkpoint.Scope) { + s.scopes[scope] = struct{}{} + if len(s.scopes) < s.compactAt { + return + } + for scope := range s.scopes { + if collapsed := checkpoint.Collapse(scope); collapsed != scope { + delete(s.scopes, scope) + s.scopes[collapsed] = struct{}{} + } + } + s.compactAt = max(minCompactSize, 2*len(s.scopes)) +} + func (c *Client) Flush(ctx context.Context) error { c.mutex.Lock() defer c.mutex.Unlock() - bldr := array.NewRecordBuilder(memory.DefaultAllocator, c.schema) - keys, values := bldr.Field(0).(*array.StringBuilder), bldr.Field(1).(*array.StringBuilder) - var version *array.Uint64Builder - if c.versionedMode { - version = bldr.Field(2).(*array.Uint64Builder) - } - for k := range c.changes { - val := c.mem[k] - keys.Append(k) - values.Append(val.value) - if version != nil { - version.Append(val.version) - } + if c.checkpoints != nil { + return c.flushCheckpoints(ctx) } - rec := bldr.NewRecordBatch() + rec := c.record(maps.Keys(c.changes)) recordBytes, err := pb.RecordToBytes(rec) if err != nil { return err @@ -185,10 +249,130 @@ func (c *Client) Flush(ctx context.Context) error { return err } - c.changes = make(map[string]struct{}) + c.changes = make(map[string]*scopeSet) return nil } +// flushCheckpoints sends one checkpoint per scope, so each waits only for the resources of the scopes that changed its +// keys. Keys whose checkpoint fails stay pending. +func (c *Client) flushCheckpoints(ctx context.Context) error { + aborted := c.abortedKeys() + groups := make(map[string]*checkpointGroup) + for key, scopes := range c.changes { + if _, ok := aborted[key]; ok { + continue + } + scope, id := joinScopes(scopes.scopes) + group, ok := groups[id] + if !ok { + group = &checkpointGroup{scope: scope} + groups[id] = group + } + group.keys = append(group.keys, key) + } + + var errs []error + if len(aborted) > 0 { + errs = append(errs, fmt.Errorf("%d keys: %w", len(aborted), checkpoint.ErrSyncAborted)) + } + // keys changed outside table resolvers do not wait for any scope, and their group id is empty, so they go first + for _, id := range slices.Sorted(maps.Keys(groups)) { + group := groups[id] + msg := &message.SyncCheckpoint{Record: c.record(slices.Values(group.keys))} + sent, err := checkpoint.Send(ctx, c.checkpoints, group.scope, c.unsent(group.keys), msg) + if err != nil { + errs = append(errs, err) + continue + } + sent.NotifyAbort(&c.abortSeen) + for _, key := range group.keys { + delete(c.changes, key) + c.sent[key] = sent + } + } + return errors.Join(errs...) +} + +// checkpointGroup is the keys whose changes wait for the same scopes, so one checkpoint covers them. +type checkpointGroup struct { + scope checkpoint.Scope + keys []string +} + +// joinScopes joins scopes, and returns an id that is the same for sets that join the same scopes, and empty for none. +func joinScopes(scopes map[checkpoint.Scope]struct{}) (checkpoint.Scope, string) { + unique := make(map[checkpoint.Scope]struct{}, len(scopes)) + for s := range scopes { + if s = checkpoint.Collapse(s); s != nil { + unique[s] = struct{}{} + } + } + ids := make([]string, 0, len(unique)) + for s := range unique { + ids = append(ids, fmt.Sprintf("%p", s)) + } + slices.Sort(ids) + return checkpoint.Join(slices.Collect(maps.Keys(unique))...), strings.Join(ids, ",") +} + +// abortedKeys returns the keys whose last checkpoint was dropped by a cancelled sync. They stay pending, as the +// resources they summarize may have been dropped too. It only looks once a checkpoint was aborted. +func (c *Client) abortedKeys() map[string]struct{} { + if !c.abortSeen.Load() { + return nil + } + aborted := make(map[string]struct{}) + for key, sent := range c.sent { + switch { + case sent.Done(): + delete(c.sent, key) + case sent.Aborted(): + aborted[key] = struct{}{} + if _, ok := c.changes[key]; !ok { + c.changes[key] = newScopeSet() + } + } + } + return aborted +} + +// unsent returns the previous checkpoints of keys that are not sent yet, as the scopes they wait for may still be +// sending the resources they cover. +func (c *Client) unsent(keys []string) []*checkpoint.Sent { + var unsent []*checkpoint.Sent + for _, key := range keys { + sent, ok := c.sent[key] + if !ok { + continue + } + if sent.Done() { + delete(c.sent, key) + } else if !slices.Contains(unsent, sent) { + unsent = append(unsent, sent) + } + } + return unsent +} + +func (c *Client) record(keys iter.Seq[string]) arrow.RecordBatch { + bldr := array.NewRecordBuilder(memory.DefaultAllocator, c.schema) + defer bldr.Release() + keyBldr, valueBldr := bldr.Field(0).(*array.StringBuilder), bldr.Field(1).(*array.StringBuilder) + var versionBldr *array.Uint64Builder + if c.versionedMode { + versionBldr = bldr.Field(2).(*array.Uint64Builder) + } + for k := range keys { + val := c.mem[k] + keyBldr.Append(k) + valueBldr.Append(val.value) + if versionBldr != nil { + versionBldr.Append(val.version) + } + } + return bldr.NewRecordBatch() +} + func (c *Client) GetKey(_ context.Context, key string) (string, error) { c.mutex.RLock() defer c.mutex.RUnlock() diff --git a/internal/memdb/memdb.go b/internal/memdb/memdb.go index ef2eb20231..233ea8ad91 100644 --- a/internal/memdb/memdb.go +++ b/internal/memdb/memdb.go @@ -240,6 +240,8 @@ func (c *client) Write(ctx context.Context, msgs <-chan message.WriteMessage) er } table := c.tables[tableName] c.overwrite(table, msg.Record) + case *message.WriteFlush: + msg.Done(nil) } c.memoryDBLock.Unlock() diff --git a/internal/memdb/memdb_test.go b/internal/memdb/memdb_test.go index d17ecac73e..30c3eeb4a0 100644 --- a/internal/memdb/memdb_test.go +++ b/internal/memdb/memdb_test.go @@ -9,7 +9,7 @@ import ( func TestPlugin(t *testing.T) { ctx := context.Background() - p := plugin.NewPlugin("test", "development", NewMemDBClient) + p := plugin.NewPlugin("test", "development", NewMemDBClient, plugin.WithFlushSupport()) if err := p.Init(ctx, nil, plugin.NewClientOptions{}); err != nil { t.Fatal(err) } diff --git a/internal/servers/plugin/v3/plugin.go b/internal/servers/plugin/v3/plugin.go index a5cf5c1ffd..520336f0f4 100644 --- a/internal/servers/plugin/v3/plugin.go +++ b/internal/servers/plugin/v3/plugin.go @@ -5,6 +5,7 @@ import ( "errors" "fmt" "io" + "strings" "github.com/apache/arrow-go/v18/arrow" pb "github.com/cloudquery/plugin-pb-go/pb/plugin/v3" @@ -120,7 +121,8 @@ func (s *Server) Init(ctx context.Context, req *pb.Init_Request) (*pb.Init_Respo if err := s.Plugin.Init(ctx, req.Spec, plugin.NewClientOptions{NoConnection: req.NoConnection, InvocationID: req.InvocationId}); err != nil { return nil, status.Errorf(codes.Internal, "failed to init plugin: %v", err) } - return &pb.Init_Response{}, nil + // a client without a connection does not write, and its spec may not be set + return &pb.Init_Response{SupportsFlush: !req.NoConnection && s.Plugin.SupportsFlush()}, nil } func (s *Server) Read(req *pb.Read_Request, stream pb.Plugin_ReadServer) error { @@ -184,8 +186,9 @@ func (s *Server) Sync(req *pb.Sync_Request, stream pb.Plugin_SyncServer) error { } if req.Backend != nil { syncOptions.BackendOptions = &plugin.BackendOptions{ - TableName: req.Backend.TableName, - Connection: req.Backend.Connection, + TableName: req.Backend.TableName, + Connection: req.Backend.Connection, + SupportsCheckpoints: req.SupportsCheckpoints, } } if req.Shard != nil { @@ -274,6 +277,15 @@ func (s *Server) Sync(req *pb.Sync_Request, stream pb.Plugin_SyncServer) error { Error: m.Error, }, } + case *message.SyncCheckpoint: + if !req.SupportsCheckpoints { + return status.Errorf(codes.Internal, "plugin sent a checkpoint, but the client does not support checkpoints") + } + // each part of a split checkpoint holds whole rows, and the client persists them all after the same point + if err := s.sendRecord(stream, m.Record, checkpointResponse); err != nil { + return err + } + continue default: return status.Errorf(codes.Internal, "unknown message type: %T", msg) } @@ -293,7 +305,32 @@ func (s *Server) Sync(req *pb.Sync_Request, stream pb.Plugin_SyncServer) error { return syncErr } +func insertResponse(recordBytes []byte) *pb.Sync_Response { + return &pb.Sync_Response{ + Message: &pb.Sync_Response_Insert{ + Insert: &pb.Sync_MessageInsert{ + Record: recordBytes, + }, + }, + } +} + +func checkpointResponse(recordBytes []byte) *pb.Sync_Response { + return &pb.Sync_Response{ + Message: &pb.Sync_Response_Checkpoint{ + Checkpoint: &pb.Sync_MessageCheckpoint{ + Record: recordBytes, + }, + }, + } +} + func (s *Server) sendInsert(stream pb.Plugin_SyncServer, record arrow.RecordBatch) error { + return s.sendRecord(stream, record, insertResponse) +} + +// sendRecord sends record in responses made by toResponse, split to stay within the max message size. +func (s *Server) sendRecord(stream pb.Plugin_SyncServer, record arrow.RecordBatch, toResponse func([]byte) *pb.Sync_Response) error { maxMsgSize := s.maxMsgSize() pending := []arrow.RecordBatch{record} for len(pending) > 0 { @@ -304,13 +341,7 @@ func (s *Server) sendInsert(stream pb.Plugin_SyncServer, record arrow.RecordBatc if err != nil { return status.Errorf(codes.Internal, "failed to encode record for table %q: %v", recordTableName(current), err) } - pbMsg := &pb.Sync_Response{ - Message: &pb.Sync_Response_Insert{ - Insert: &pb.Sync_MessageInsert{ - Record: recordBytes, - }, - }, - } + pbMsg := toResponse(recordBytes) size := proto.Size(pbMsg) if size <= maxMsgSize { @@ -375,76 +406,14 @@ func (s *Server) Write(stream pb.Plugin_WriteServer) error { } return status.Errorf(codes.Internal, "failed to receive msg: %v", err) } - var pluginMessage message.WriteMessage - var pbMsgConvertErr error - switch pbMsg := r.Message.(type) { - case *pb.Write_Request_MigrateTable: - sc, err := pb.NewSchemaFromBytes(pbMsg.MigrateTable.Table) - if err != nil { - pbMsgConvertErr = status.Errorf(codes.InvalidArgument, "failed to create schema from bytes: %v", err) - break - } - table, err := schema.NewTableFromArrowSchema(sc) - if err != nil { - pbMsgConvertErr = status.Errorf(codes.InvalidArgument, "failed to create table from schema: %v", err) - break - } - pluginMessage = &message.WriteMigrateTable{ - Table: table, - MigrateForce: pbMsg.MigrateTable.MigrateForce, - } - case *pb.Write_Request_Insert: - record, err := pb.NewRecordFromBytes(pbMsg.Insert.Record) - if err != nil { - pbMsgConvertErr = status.Errorf(codes.InvalidArgument, "failed to create record: %v", err) - break - } - pluginMessage = &message.WriteInsert{ - Record: record, - } - case *pb.Write_Request_Delete: - pluginMessage = &message.WriteDeleteStale{ - TableName: pbMsg.Delete.TableName, - SourceName: pbMsg.Delete.SourceName, - SyncTime: pbMsg.Delete.SyncTime.AsTime(), - } - - case *pb.Write_Request_DeleteRecord: - whereClause := make(message.PredicateGroups, len(pbMsg.DeleteRecord.WhereClause)) - - for j, predicateGroup := range pbMsg.DeleteRecord.WhereClause { - whereClause[j].GroupingType = predicateGroup.GroupingType.String() - whereClause[j].Predicates = make(message.Predicates, len(predicateGroup.Predicates)) - for i, predicate := range predicateGroup.Predicates { - record, err := pb.NewRecordFromBytes(predicate.Record) - if err != nil { - pbMsgConvertErr = status.Errorf(codes.InvalidArgument, "failed to create record: %v", err) - break - } - whereClause[j].Predicates[i] = message.Predicate{ - Record: record, - Column: predicate.Column, - Operator: predicate.Operator.String(), - } - } - } - - tableRelations := make([]message.TableRelation, len(pbMsg.DeleteRecord.TableRelations)) - for i, tr := range pbMsg.DeleteRecord.TableRelations { - tableRelations[i] = message.TableRelation{ - TableName: tr.TableName, - ParentTable: tr.ParentTable, - } - } - pluginMessage = &message.WriteDeleteRecord{ - DeleteRecord: message.DeleteRecord{ - TableName: pbMsg.DeleteRecord.TableName, - TableRelations: tableRelations, - WhereClause: whereClause, - }, + if _, ok := r.Message.(*pb.Write_Request_Flush); ok { + close(msgs) + if wgErr := eg.Wait(); wgErr != nil { + return status.Errorf(codes.InvalidArgument, "flush is only supported on WriteStream and write failed: %v", wgErr) } + return status.Errorf(codes.InvalidArgument, "flush is only supported on WriteStream") } - + pluginMessage, pbMsgConvertErr := writeRequestToMessage(r) if pbMsgConvertErr != nil { close(msgs) if wgErr := eg.Wait(); wgErr != nil { @@ -471,6 +440,183 @@ func (s *Server) Write(stream pb.Plugin_WriteServer) error { } } +// WriteStream is Write that also accepts flushes. It answers each flush once the plugin's Write has tried to commit +// every message sent before it, with an error if some of them may not have been committed. A failed flush does not end +// the stream. +func (s *Server) WriteStream(stream pb.Plugin_WriteStreamServer) error { + if !s.Plugin.SupportsFlush() { + return status.Errorf(codes.Unimplemented, "plugin does not support flush") + } + + msgs := make(chan message.WriteMessage) + ctx := stream.Context() + eg, gctx := errgroup.WithContext(ctx) + writeDone := make(chan struct{}) + eg.Go(func() error { + defer close(writeDone) + return s.Plugin.Write(gctx, msgs) + }) + stop := func(code codes.Code, format string, args ...any) error { + close(msgs) + if err := eg.Wait(); err != nil { + return status.Errorf(code, format+" and write failed: %v", append(args, err)...) + } + return status.Errorf(code, format, args...) + } + + for { + r, err := stream.Recv() + if err == io.EOF { + close(msgs) + if err := eg.Wait(); err != nil { + return status.Errorf(codes.Internal, "write failed: %v", err) + } + return nil + } + if err != nil { + return stop(codes.Internal, "failed to receive msg: %v", err) + } + + var ( + pluginMessage message.WriteMessage + flush *message.WriteFlush + ) + if _, ok := r.Message.(*pb.Write_Request_Flush); ok { + flush = message.NewWriteFlush() + pluginMessage = flush + } else { + pluginMessage, err = writeRequestToMessage(r) + if err != nil { + return stop(codes.InvalidArgument, "failed to convert message: %v", err) + } + } + + select { + case msgs <- pluginMessage: + case <-writeDone: + return stop(codes.Internal, "plugin stopped writing") + case <-gctx.Done(): + return stop(codes.Canceled, "plugin stopped writing") + case <-ctx.Done(): + return stop(codes.Canceled, "context done: %v", ctx.Err()) + } + if flush == nil { + continue + } + + flushErr, completed, err := waitForFlush(ctx, flush, writeDone) + switch { + case err != nil: + return stop(codes.Canceled, "context done: %v", err) + case !completed: + return stop(codes.Internal, "plugin stopped writing before the flush completed") + } + if err := stream.Send(flushResponse(r.GetFlush().GetId(), flushErr)); err != nil { + return stop(codes.Internal, "failed to send flush response: %v", err) + } + } +} + +// waitForFlush returns the error flush completed with, or completed=false if Write returned without completing it. +func waitForFlush(ctx context.Context, flush *message.WriteFlush, writeDone <-chan struct{}) (flushErr error, completed bool, err error) { + select { + case result := <-flush.Result(): + return result, true, nil + case <-writeDone: + // Write may have completed the flush right before returning + select { + case result := <-flush.Result(): + return result, true, nil + default: + return nil, false, nil + } + case <-ctx.Done(): + return nil, false, ctx.Err() + } +} + +func flushResponse(id uint64, flushErr error) *pb.Write_Response { + res := &pb.Write_Response{FlushId: id} + if flushErr == nil { + return res + } + // protobuf strings must be valid UTF-8, and an empty error would report success + res.FlushError = strings.ToValidUTF8(flushErr.Error(), "\uFFFD") + if res.FlushError == "" { + res.FlushError = "flush failed" + } + for _, table := range message.FailedTables(flushErr) { + res.FailedTables = append(res.FailedTables, strings.ToValidUTF8(table, "\uFFFD")) + } + return res +} + +func writeRequestToMessage(r *pb.Write_Request) (message.WriteMessage, error) { + switch pbMsg := r.Message.(type) { + case *pb.Write_Request_MigrateTable: + sc, err := pb.NewSchemaFromBytes(pbMsg.MigrateTable.Table) + if err != nil { + return nil, status.Errorf(codes.InvalidArgument, "failed to create schema from bytes: %v", err) + } + table, err := schema.NewTableFromArrowSchema(sc) + if err != nil { + return nil, status.Errorf(codes.InvalidArgument, "failed to create table from schema: %v", err) + } + return &message.WriteMigrateTable{ + Table: table, + MigrateForce: pbMsg.MigrateTable.MigrateForce, + }, nil + case *pb.Write_Request_Insert: + record, err := pb.NewRecordFromBytes(pbMsg.Insert.Record) + if err != nil { + return nil, status.Errorf(codes.InvalidArgument, "failed to create record: %v", err) + } + return &message.WriteInsert{ + Record: record, + }, nil + case *pb.Write_Request_Delete: + return &message.WriteDeleteStale{ + TableName: pbMsg.Delete.TableName, + SourceName: pbMsg.Delete.SourceName, + SyncTime: pbMsg.Delete.SyncTime.AsTime(), + }, nil + case *pb.Write_Request_DeleteRecord: + whereClause := make(message.PredicateGroups, len(pbMsg.DeleteRecord.WhereClause)) + for j, predicateGroup := range pbMsg.DeleteRecord.WhereClause { + whereClause[j].GroupingType = predicateGroup.GroupingType.String() + whereClause[j].Predicates = make(message.Predicates, len(predicateGroup.Predicates)) + for i, predicate := range predicateGroup.Predicates { + record, err := pb.NewRecordFromBytes(predicate.Record) + if err != nil { + return nil, status.Errorf(codes.InvalidArgument, "failed to create record: %v", err) + } + whereClause[j].Predicates[i] = message.Predicate{ + Record: record, + Column: predicate.Column, + Operator: predicate.Operator.String(), + } + } + } + + tableRelations := make([]message.TableRelation, len(pbMsg.DeleteRecord.TableRelations)) + for i, tr := range pbMsg.DeleteRecord.TableRelations { + tableRelations[i] = message.TableRelation{ + TableName: tr.TableName, + ParentTable: tr.ParentTable, + } + } + return &message.WriteDeleteRecord{ + DeleteRecord: message.DeleteRecord{ + TableName: pbMsg.DeleteRecord.TableName, + TableRelations: tableRelations, + WhereClause: whereClause, + }, + }, nil + default: + return nil, status.Errorf(codes.InvalidArgument, "unknown message type: %T", pbMsg) + } +} + func (s *Server) Transform(stream pb.Plugin_TransformServer) error { var ( recvRecords = make(chan arrow.RecordBatch) diff --git a/internal/servers/plugin/v3/plugin_msgsize_test.go b/internal/servers/plugin/v3/plugin_msgsize_test.go index c0af39d191..511cfaac76 100644 --- a/internal/servers/plugin/v3/plugin_msgsize_test.go +++ b/internal/servers/plugin/v3/plugin_msgsize_test.go @@ -113,3 +113,26 @@ func TestMaxRowsPerMessage(t *testing.T) { }) } } + +func TestSendRecordSplitsOversizedCheckpoint(t *testing.T) { + t.Parallel() + + const maxMsgSize = 8 * 1024 + record := newTestRecord(t, "cq_state", 40, 1024) + s := &Server{Logger: zerolog.Nop(), MaxMsgSize: maxMsgSize} + stream := &recordingSyncServer{} + + require.NoError(t, s.sendRecord(stream, record, checkpointResponse)) + require.Greater(t, len(stream.sent), 1, "oversized checkpoint should be split into several messages") + + var gotRows int64 + for _, msg := range stream.sent { + require.LessOrEqual(t, proto.Size(msg), maxMsgSize) + checkpoint := msg.GetCheckpoint() + require.NotNil(t, checkpoint, "every part of a checkpoint is a checkpoint") + part, err := pb.NewRecordFromBytes(checkpoint.Record) + require.NoError(t, err) + gotRows += part.NumRows() + } + require.EqualValues(t, 40, gotRows) +} diff --git a/internal/servers/plugin/v3/plugin_test.go b/internal/servers/plugin/v3/plugin_test.go index f5daf4a817..85b1546688 100644 --- a/internal/servers/plugin/v3/plugin_test.go +++ b/internal/servers/plugin/v3/plugin_test.go @@ -14,6 +14,7 @@ import ( "github.com/apache/arrow-go/v18/arrow/memory" pb "github.com/cloudquery/plugin-pb-go/pb/plugin/v3" "github.com/cloudquery/plugin-sdk/v4/internal/memdb" + "github.com/cloudquery/plugin-sdk/v4/message" "github.com/cloudquery/plugin-sdk/v4/plugin" "github.com/cloudquery/plugin-sdk/v4/schema" "github.com/rs/zerolog" @@ -392,3 +393,76 @@ func makeRecordFromString(s string) arrow.RecordBatch { return array.NewRecordBatch(sch, []arrow.Array{arr}, 1) } + +type checkpointingSourceClient struct { + plugin.UnimplementedDestination +} + +func (*checkpointingSourceClient) Close(context.Context) error { return nil } + +func (*checkpointingSourceClient) Tables(context.Context, plugin.TableOptions) (schema.Tables, error) { + return nil, nil +} + +func (*checkpointingSourceClient) Sync(_ context.Context, _ plugin.SyncOptions, res chan<- message.SyncMessage) error { + bldr := array.NewRecordBuilder(memory.DefaultAllocator, (&schema.Table{Name: "state", Columns: schema.ColumnList{{Name: "key", Type: arrow.BinaryTypes.String}}}).ToArrowSchema()) + defer bldr.Release() + bldr.Field(0).(*array.StringBuilder).Append("cursor") + res <- &message.SyncCheckpoint{Record: bldr.NewRecordBatch()} + return nil +} + +func TestPluginSyncCheckpoints(t *testing.T) { + ctx := context.Background() + s := Server{ + Plugin: plugin.NewPlugin("test", "development", func(context.Context, zerolog.Logger, []byte, plugin.NewClientOptions) (plugin.Client, error) { + return &checkpointingSourceClient{}, nil + }), + } + _, err := s.Init(ctx, &pb.Init_Request{}) + require.NoError(t, err) + + stream := &mockSyncServer{} + require.NoError(t, s.Sync(&pb.Sync_Request{SupportsCheckpoints: true}, stream)) + require.Len(t, stream.messages, 1) + + stream = &mockSyncServer{} + require.ErrorContains(t, s.Sync(&pb.Sync_Request{}, stream), "client does not support checkpoints") + require.Empty(t, stream.messages) +} + +func TestWaitForFlush(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + writeDone := make(chan struct{}) + close(writeDone) + // Write completed the flush and returned: whichever the select picks, the result is not lost + for range 100 { + flush := message.NewWriteFlush() + flush.Done(errors.New("disk full")) + flushErr, completed, err := waitForFlush(ctx, flush, writeDone) + require.NoError(t, err) + require.True(t, completed) + require.EqualError(t, flushErr, "disk full") + } + + _, completed, err := waitForFlush(ctx, message.NewWriteFlush(), writeDone) + require.NoError(t, err) + require.False(t, completed, "Write returned without completing the flush") + + cancelled, cancel := context.WithCancel(ctx) + cancel() + _, _, err = waitForFlush(cancelled, message.NewWriteFlush(), make(chan struct{})) + require.ErrorIs(t, err, context.Canceled) +} + +func TestFlushResponse(t *testing.T) { + require.Equal(t, &pb.Write_Response{FlushId: 1}, flushResponse(1, nil)) + require.Equal(t, &pb.Write_Response{FlushId: 2, FlushError: "flush failed"}, flushResponse(2, errors.New(""))) + require.Equal(t, &pb.Write_Response{FlushId: 3, FlushError: "bad � byte"}, flushResponse(3, errors.New("bad \xff byte"))) + require.Equal(t, &pb.Write_Response{ + FlushId: 4, + FlushError: "table b: disk full", + FailedTables: []string{"a", "b�"}, + }, flushResponse(4, &message.FailedTablesError{Tables: []string{"b\xff", "a"}, Err: errors.New("table b: disk full")})) +} diff --git a/message/sync_message.go b/message/sync_message.go index c92e449abb..a9852c7349 100644 --- a/message/sync_message.go +++ b/message/sync_message.go @@ -131,3 +131,18 @@ type SyncError struct { func (e SyncError) GetTable() *schema.Table { return &schema.Table{Name: e.TableName} } + +// SyncCheckpoint carries state backend rows (key, value, version). The client persists them once every message sent +// before the checkpoint is committed to all destinations. +type SyncCheckpoint struct { + syncBaseMessage + Record arrow.RecordBatch +} + +func (m *SyncCheckpoint) GetTable() *schema.Table { + table, err := schema.NewTableFromArrowSchema(m.Record.Schema()) + if err != nil { + panic(err) + } + return table +} diff --git a/message/write_message.go b/message/write_message.go index 7ba369d4aa..6f432e292d 100644 --- a/message/write_message.go +++ b/message/write_message.go @@ -2,6 +2,8 @@ package message import ( "slices" + "strings" + "sync" "time" "github.com/apache/arrow-go/v18/arrow" @@ -169,3 +171,80 @@ func (m WriteDeleteRecord) GetTable() *schema.Table { } type WriteDeleteRecords []*WriteDeleteRecord + +// WriteFlush asks the writer to commit every message it received before this one. +// The writer calls Done once it has tried to commit them, with an error if some messages it received since the previous +// flush may not have been committed. A *FailedTablesError names the tables of those messages. +type WriteFlush struct { + writeBaseMessage + result chan error + once sync.Once +} + +func NewWriteFlush() *WriteFlush { + return &WriteFlush{result: make(chan error, 1)} +} + +func (*WriteFlush) GetTable() *schema.Table { + return &schema.Table{} +} + +// Done reports the outcome of the flush. Only the first call has an effect. +func (m *WriteFlush) Done(err error) { + m.once.Do(func() { m.result <- err }) +} + +// Result receives the error passed to Done. +func (m *WriteFlush) Result() <-chan error { + return m.result +} + +// FailedTablesError is a flush error that names the tables whose messages may not have been committed. +type FailedTablesError struct { + Tables []string + Err error +} + +func (e *FailedTablesError) Error() string { + if e.Err == nil { + return "failed to write tables " + strings.Join(e.Tables, ", ") + } + return e.Err.Error() +} + +func (e *FailedTablesError) Unwrap() error { return e.Err } + +// FailedTables returns the tables a flush error names, sorted and without duplicates. It returns nil if some error +// joined in err names no tables, as messages of any table may then have been lost. +func FailedTables(err error) []string { + tables, ok := collectFailedTables(err) + if !ok { + return nil + } + slices.Sort(tables) + return slices.Compact(tables) +} + +func collectFailedTables(err error) ([]string, bool) { + switch e := err.(type) { + case *FailedTablesError: + if len(e.Tables) == 0 || slices.Contains(e.Tables, "") { + return nil, false + } + return slices.Clone(e.Tables), true + case interface{ Unwrap() []error }: + var tables []string + for _, err := range e.Unwrap() { + t, ok := collectFailedTables(err) + if !ok { + return nil, false + } + tables = append(tables, t...) + } + return tables, len(tables) > 0 + case interface{ Unwrap() error }: + return collectFailedTables(e.Unwrap()) + default: + return nil, false + } +} diff --git a/message/write_message_test.go b/message/write_message_test.go new file mode 100644 index 0000000000..7cbd8c635c --- /dev/null +++ b/message/write_message_test.go @@ -0,0 +1,40 @@ +package message + +import ( + "errors" + "fmt" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestFailedTables(t *testing.T) { + failed := func(tables ...string) error { + return &FailedTablesError{Tables: tables, Err: errors.New("write failed")} + } + for _, tc := range []struct { + name string + err error + want []string + }{ + {name: "nil", err: nil}, + {name: "unattributed", err: errors.New("write failed")}, + {name: "tables", err: failed("b", "a", "b"), want: []string{"a", "b"}}, + {name: "wrapped", err: fmt.Errorf("flush: %w", failed("a")), want: []string{"a"}}, + {name: "joined", err: errors.Join(failed("a"), fmt.Errorf("flush: %w", failed("b"))), want: []string{"a", "b"}}, + {name: "joined with unattributed", err: errors.Join(failed("a"), errors.New("migrate failed"))}, + {name: "no tables", err: failed()}, + {name: "empty table name", err: failed("a", "")}, + } { + t.Run(tc.name, func(t *testing.T) { + require.Equal(t, tc.want, FailedTables(tc.err)) + }) + } +} + +func TestFailedTablesErrorWithoutCause(t *testing.T) { + err := &FailedTablesError{Tables: []string{"a", "b"}} + require.EqualError(t, err, "failed to write tables a, b") + require.NoError(t, errors.Unwrap(err)) + require.Equal(t, []string{"a", "b"}, FailedTables(err)) +} diff --git a/plugin/options.go b/plugin/options.go index 44c44744d8..a89cf75023 100644 --- a/plugin/options.go +++ b/plugin/options.go @@ -53,6 +53,27 @@ func WithConnectionTester(tester ConnectionTester) Option { } } +// WithFlushSupport declares that the client's Write handles message.WriteFlush: it tries to commit every message +// received before the flush, then calls Done, with an error if some of them may not have been committed. The writers in +// the writers package do this, so a client whose Write passes the message channel to one of them can set this option. A Write that wraps or filters the channel must forward every +// message.WriteFlush, or call its Done itself, or the sync waits for it forever. +// +// A flush ends every open batch, like reaching the batch size does. Do not set this option if ending a batch early +// loses or corrupts data, for example a file destination that writes the whole sync to a single object. A client whose +// support depends on its spec implements FlushSupporter. +// +// A flush fails if a message since the previous flush may not have been written, and names the tables of those messages +// if the writer can tell. The CLI then stops saving state for the source, so that its state never passes the lost +// rows, skips delete-stale for the failed tables, and reports the failure once the sync finishes. The next sync then +// syncs again every row after the last saved state, so in append write mode, the destination gets those rows again. +// Writers that stop on a failed write, such as the mixed batch writer, end the sync, unless they reported the failure +// on a flush first. Without flush support, the batch writer only logs failed batches. +func WithFlushSupport() Option { + return func(p *Plugin) { + p.supportsFlush = true + } +} + type TableOptions struct { Tables []string SkipTables []string diff --git a/plugin/plugin.go b/plugin/plugin.go index 3871a830c4..cfc86c5fa2 100644 --- a/plugin/plugin.go +++ b/plugin/plugin.go @@ -101,6 +101,8 @@ type Plugin struct { invocationID string // Method to test connection given a spec testConnFn ConnectionTester + // the client's Write commits everything received before a message.WriteFlush, then calls its Done + supportsFlush bool } // NewPlugin returns a new CloudQuery Plugin with the given name, version and implementation. @@ -133,6 +135,23 @@ func (p *Plugin) InvocationID() string { return p.invocationID } +// FlushSupporter is implemented by clients whose flush support depends on their spec. +// It is only consulted for plugins created with WithFlushSupport. +type FlushSupporter interface { + SupportsFlush() bool +} + +// SupportsFlush reports whether the plugin's Write handles message.WriteFlush +func (p *Plugin) SupportsFlush() bool { + if !p.supportsFlush { + return false + } + if client, ok := p.client.(FlushSupporter); ok { + return client.SupportsFlush() + } + return true +} + // Name returns the name of this plugin func (p *Plugin) Name() string { return p.name diff --git a/plugin/plugin_source.go b/plugin/plugin_source.go index f626244c10..1c74387ec4 100644 --- a/plugin/plugin_source.go +++ b/plugin/plugin_source.go @@ -14,6 +14,9 @@ import ( type BackendOptions struct { TableName string Connection string + // SupportsCheckpoints is set when the client persists state from message.SyncCheckpoint, + // so state clients send their writes on the sync stream instead of writing to the backend. + SupportsCheckpoints bool } type Shard struct { diff --git a/plugin/testing_write.go b/plugin/testing_write.go index dca0e9cca0..7feb10abbb 100644 --- a/plugin/testing_write.go +++ b/plugin/testing_write.go @@ -196,6 +196,18 @@ func TestWriterSuiteRunner(t *testing.T, p *Plugin, tests WriterTestSuiteTests, }) }) + t.Run("TestFlush", func(t *testing.T) { + if !p.SupportsFlush() { + t.Skip("plugin does not support flush") + } + if suite.tests.SkipInsert { + t.Skip("skipping " + t.Name()) + } + if err := suite.testFlush(ctx); err != nil { + t.Fatal(err) + } + }) + t.Run("TestMigrate", func(t *testing.T) { if suite.tests.SkipMigrate { t.Skip("skipping " + t.Name()) diff --git a/plugin/testing_write_flush.go b/plugin/testing_write_flush.go new file mode 100644 index 0000000000..eb9802f88f --- /dev/null +++ b/plugin/testing_write_flush.go @@ -0,0 +1,157 @@ +package plugin + +import ( + "context" + "fmt" + "time" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/cloudquery/plugin-sdk/v4/message" + "github.com/cloudquery/plugin-sdk/v4/schema" +) + +// flushTestTimeout bounds each wait in testFlush, so that a plugin that never completes a flush fails the test +// instead of hanging it. +const flushTestTimeout = time.Minute + +// testFlush checks that the changes before each flush are readable once the flush completes, while Write is still +// running. +func (s *WriterTestSuite) testFlush(ctx context.Context) error { + table := &schema.Table{ + Name: s.tableNameForTest("flush"), + Columns: schema.ColumnList{ + {Name: "id", Type: arrow.PrimitiveTypes.Int64, PrimaryKey: true, NotNull: true}, + schema.CqSourceNameColumn, + schema.CqSyncTimeColumn, + }, + } + const sourceName = "source-test" + syncTime := time.Now().UTC().Truncate(s.genDatOptions.TimePrecision).Truncate(time.Microsecond) + + ctx, cancel := context.WithCancel(ctx) + defer cancel() + msgs := make(chan message.WriteMessage) + writeErr := make(chan error, 1) + writeDone := make(chan struct{}) + go func() { + defer close(writeDone) + writeErr <- s.plugin.Write(ctx, msgs) + }() + msgsOpen := true + defer func() { + if msgsOpen { + close(msgs) + } + // a failed step must not leave Write running into the next test + cancel() + select { + case <-writeDone: + case <-time.After(flushTestTimeout): + } + }() + + send := func(msg message.WriteMessage) error { + select { + case msgs <- msg: + return nil + case err := <-writeErr: + return fmt.Errorf("write returned early: %w", err) + case <-time.After(flushTestTimeout): + return fmt.Errorf("write did not take a %T message within %s", msg, flushTestTimeout) + } + } + flushAndCount := func(step string, wantRows int64) error { + msg := message.NewWriteFlush() + if err := send(msg); err != nil { + return fmt.Errorf("%s: %w", step, err) + } + select { + case err := <-msg.Result(): + if err != nil { + return fmt.Errorf("%s: flush failed: %w", step, err) + } + case err := <-writeErr: + return fmt.Errorf("%s: write returned before the flush completed: %w", step, err) + case <-time.After(flushTestTimeout): + return fmt.Errorf("%s: flush did not complete within %s", step, flushTestTimeout) + } + records, err := s.plugin.readAll(ctx, table) + if err != nil { + return fmt.Errorf("%s: failed to read: %w", step, err) + } + if rows := TotalRows(records); rows != wantRows { + return fmt.Errorf("%s: expected %d rows after the flush, got %d", step, wantRows, rows) + } + return nil + } + insert := func(syncTime time.Time, ids ...int64) error { + bldr := array.NewRecordBuilder(memory.DefaultAllocator, table.ToArrowSchema()) + defer bldr.Release() + for _, id := range ids { + bldr.Field(0).(*array.Int64Builder).Append(id) + bldr.Field(1).(*array.StringBuilder).Append(sourceName) + bldr.Field(2).(*array.TimestampBuilder).AppendTime(syncTime) + } + return send(&message.WriteInsert{Record: bldr.NewRecordBatch()}) + } + + if err := send(&message.WriteMigrateTable{Table: table}); err != nil { + return err + } + if err := insert(syncTime, 1, 2, 3); err != nil { + return err + } + if err := flushAndCount("first insert", 3); err != nil { + return err + } + syncTime = syncTime.Add(time.Second) + if err := insert(syncTime, 4, 5); err != nil { + return err + } + if err := flushAndCount("second insert", 5); err != nil { + return err + } + + rows := int64(5) + if !s.tests.SkipDeleteStale { + if err := send(&message.WriteDeleteStale{TableName: table.Name, SourceName: sourceName, SyncTime: syncTime}); err != nil { + return err + } + rows = 2 + if err := flushAndCount("delete stale", rows); err != nil { + return err + } + } + if !s.tests.SkipDeleteRecord { + bldr := array.NewRecordBuilder(memory.DefaultAllocator, (&schema.Table{ + Name: table.Name, + Columns: schema.ColumnList{{Name: "id", Type: arrow.PrimitiveTypes.Int64}}, + }).ToArrowSchema()) + defer bldr.Release() + bldr.Field(0).(*array.Int64Builder).Append(4) + if err := send(&message.WriteDeleteRecord{DeleteRecord: message.DeleteRecord{ + TableName: table.Name, + WhereClause: message.PredicateGroups{{ + GroupingType: "AND", + Predicates: message.Predicates{{Operator: "eq", Column: "id", Record: bldr.NewRecordBatch()}}, + }}, + }}); err != nil { + return err + } + rows-- + if err := flushAndCount("delete record", rows); err != nil { + return err + } + } + + close(msgs) + msgsOpen = false + select { + case err := <-writeErr: + return err + case <-time.After(flushTestTimeout): + return fmt.Errorf("write did not return within %s after its messages ended", flushTestTimeout) + } +} diff --git a/scheduler/batch.go b/scheduler/batch.go index f869dac9a8..114c393459 100644 --- a/scheduler/batch.go +++ b/scheduler/batch.go @@ -84,6 +84,9 @@ func (s *BatchSettings) getBatcher(ctx context.Context, res chan<- message.SyncM type batcherInterface interface { process(res *schema.Resource) + // flush sends every resource passed to process so far. It reports false if the sync was cancelled first, as + // workers then drop the resources they hold. + flush() bool close() } @@ -95,6 +98,8 @@ func (n *nopBatcher) process(resource *schema.Resource) { n.res <- &message.SyncInsert{Record: resource.GetValues().ToArrowRecord(resource.Table.ToArrowSchema())} } +func (*nopBatcher) flush() bool { return true } + func (*nopBatcher) close() {} var _ batcherInterface = (*nopBatcher)(nil) @@ -165,6 +170,17 @@ func (w *worker) reachedSizeLimit() bool { return w.maxSizeBytes > 0 && w.bytesPerRow > 0 && int64(w.curRows)*w.bytesPerRow >= w.maxSizeBytes } +// add appends r to the batch and reports whether that sent the batch. +func (w *worker) add(r *schema.Resource) bool { + scalar.AppendToRecordBuilder(w.builder, r.GetValues()) + w.curRows++ + if (w.maxRows > 0 && w.curRows == w.maxRows) || w.reachedSizeLimit() { + w.send() + return true + } + return false +} + func (w *worker) work(done <-chan struct{}, timeout time.Duration) { ticker := writers.NewTicker(timeout) defer ticker.Stop() @@ -180,12 +196,7 @@ func (w *worker) work(done <-chan struct{}, timeout time.Duration) { return } - // append to builder - scalar.AppendToRecordBuilder(w.builder, r.GetValues()) - w.curRows++ - // check if we need to flush - if (w.maxRows > 0 && w.curRows == w.maxRows) || w.reachedSizeLimit() { - w.send() + if w.add(r) { ticker.Reset(timeout) } @@ -195,6 +206,10 @@ func (w *worker) work(done <-chan struct{}, timeout time.Duration) { } case ch := <-w.flush: + // resources queued before the flush request must go out with it + for queued := len(w.ch); queued > 0; queued-- { + w.add(<-w.ch) + } if w.curRows > 0 { w.send() ticker.Reset(timeout) @@ -217,9 +232,12 @@ func (b *batcher) process(res *schema.Resource) { return } - // we alloc only ch here, as it may be needed right away + // we alloc only the channels here, as they may be needed right away // for instance, if another goroutine will get the value allocated by us - wr := &worker{ch: make(chan *schema.Resource, 5)} // 5 is quite enough + wr := &worker{ + ch: make(chan *schema.Resource, 5), // 5 is quite enough + flush: make(chan chan struct{}), + } v, loaded = b.workers.LoadOrStore(table.Name, wr) if loaded { // means that the worker was already in tne sync.Map, so we just discard the wr value @@ -235,7 +253,6 @@ func (b *batcher) process(res *schema.Resource) { defer b.wg.Done() // fill in the worker fields - wr.flush = make(chan chan struct{}) wr.maxRows = b.maxRows wr.maxSizeBytes = b.maxSizeBytes wr.builder = array.NewRecordBuilder(memory.DefaultAllocator, table.ToArrowSchema()) @@ -251,6 +268,27 @@ func (b *batcher) process(res *schema.Resource) { wr.ch <- res } +func (b *batcher) flush() bool { + flushed := true + b.workers.Range(func(_, v any) bool { + ack := make(chan struct{}) + select { + case v.(*worker).flush <- ack: + case <-b.done: + flushed = false + return false + } + select { + case <-ack: + return true + case <-b.done: + flushed = false + return false + } + }) + return flushed +} + func (b *batcher) close() { b.workers.Range(func(_, v any) bool { close(v.(*worker).ch) diff --git a/scheduler/batchsender/batch_sender.go b/scheduler/batchsender/batch_sender.go index a851814a44..971a60d287 100644 --- a/scheduler/batchsender/batch_sender.go +++ b/scheduler/batchsender/batch_sender.go @@ -75,7 +75,8 @@ func (bs *BatchSender) flushLocked(items ...any) { bs.items = nil } -func (bs *BatchSender) Close() { +// Flush sends the buffered items and waits for sendFn to return. +func (bs *BatchSender) Flush() { bs.mu.Lock() defer bs.mu.Unlock() @@ -84,3 +85,7 @@ func (bs *BatchSender) Close() { } bs.flushLocked() } + +func (bs *BatchSender) Close() { + bs.Flush() +} diff --git a/scheduler/checkpoint.go b/scheduler/checkpoint.go new file mode 100644 index 0000000000..913605a9fa --- /dev/null +++ b/scheduler/checkpoint.go @@ -0,0 +1,383 @@ +package scheduler + +import ( + "context" + "errors" + "maps" + "slices" + "sync" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/cloudquery/plugin-sdk/v4/schema" + + "github.com/cloudquery/plugin-sdk/v4/internal/checkpoint" + "github.com/cloudquery/plugin-sdk/v4/message" +) + +// syncCheckpoints sends checkpoints on the sync message channel after the resources handed to the batcher before them. +// As a checkpoint.Scope, it holds checkpoints until the sync completes, for strategies that cannot track what a table +// resolver produced. +type syncCheckpoints struct { + res chan<- message.SyncMessage + requests chan []message.SyncMessage + // closed when Sync stops reading requests + loopDone chan struct{} + // closed once the batcher is closed + finished chan struct{} + // the scope of every table resolver that ended: they all send checkpoints after the resources Sync received + ended *endedScope + + mu sync.Mutex + deferred []message.SyncMessage + // state rows held until the sync completes, by table, so repeated checkpoints of a key keep only the newest row + deferredState map[string]*heldState + // sent once the held state is sent, to track the checkpoints merged into it + heldSent *checkpoint.Sent + // table scopes that did not end, whose pending checkpoints an aborted sync drops + scopes map[*tableScope]struct{} + done bool + completed bool + + // set when the caller closes res once Sync returns, as SyncAll does + closeRes bool + // held to send on res after Sync, so that Sync returns only once those sends are over + resMu sync.RWMutex + resClosed bool +} + +var ( + _ checkpoint.Scope = (*syncCheckpoints)(nil) + _ checkpoint.Covering = (*syncCheckpoints)(nil) +) + +// state table columns, see the state client +const ( + stateKeyColumn = "key" + stateValueColumn = "value" + stateVersionColumn = "version" +) + +func newSyncCheckpoints(res chan<- message.SyncMessage) *syncCheckpoints { + c := &syncCheckpoints{ + res: res, + requests: make(chan []message.SyncMessage), + loopDone: make(chan struct{}), + finished: make(chan struct{}), + } + c.ended = &endedScope{checkpoints: c} + return c +} + +// send sends msgs after every resource Sync received so far. It drops msgs if it fails. +func (c *syncCheckpoints) send(ctx context.Context, msgs []message.SyncMessage) error { + if len(msgs) == 0 { + return nil + } + select { + case c.requests <- msgs: + return nil + case <-c.loopDone: + return c.sendAfterSync(ctx, msgs) + case <-ctx.Done(): + dropMessages(msgs) + return context.Cause(ctx) + } +} + +func (c *syncCheckpoints) sendAfterSync(ctx context.Context, msgs []message.SyncMessage) error { + select { + case <-c.finished: + case <-ctx.Done(): + dropMessages(msgs) + return context.Cause(ctx) + } + c.resMu.RLock() + defer c.resMu.RUnlock() + switch { + case !c.completed: + dropMessages(msgs) + return checkpoint.ErrSyncAborted + case c.resClosed: + dropMessages(msgs) + return errResClosed + } + return sendMessages(ctx, c.res, msgs) +} + +var errResClosed = errors.New("the sync message channel is closed, as the sync returned") + +// returned is called when Sync returns. It waits for checkpoints being sent, if the caller then closes res. +func (c *syncCheckpoints) returned() { + if !c.closeRes { + return + } + c.resMu.Lock() + defer c.resMu.Unlock() + c.resClosed = true +} + +// endedScope is the scope of table resolvers that ended. +type endedScope struct { + checkpoints *syncCheckpoints +} + +func (e *endedScope) Checkpoint(ctx context.Context, msg message.SyncMessage) error { + return e.checkpoints.send(ctx, []message.SyncMessage{msg}) +} + +func (c *syncCheckpoints) Checkpoint(ctx context.Context, msg message.SyncMessage) error { + c.mu.Lock() + if !c.done { + if !c.holdState(msg) { + c.deferred = append(c.deferred, msg) + } + c.mu.Unlock() + return nil + } + c.mu.Unlock() + return c.sendAfterSync(ctx, []message.SyncMessage{msg}) +} + +// Covers reports whether sent is merged into the held state, which a checkpoint passed now is merged into too, or sent +// after. +func (c *syncCheckpoints) Covers(sent *checkpoint.Sent) bool { + c.mu.Lock() + defer c.mu.Unlock() + return !c.done && c.heldSent != nil && sent.MergedInto(c.heldSent) +} + +// complete is called once the batcher is closed after sending every resource. +func (c *syncCheckpoints) complete(ctx context.Context) error { + c.mu.Lock() + deferred := c.deferred + for _, state := range c.deferredState { + deferred = append(deferred, &message.SyncCheckpoint{Record: state.record()}) + } + heldSent := c.heldSent + c.deferred = nil + c.deferredState = nil + c.heldSent = nil + c.done = true + c.completed = true + c.mu.Unlock() + close(c.finished) + if err := sendMessages(ctx, c.res, deferred); err != nil { + if heldSent != nil { + heldSent.Abort() + } + return err + } + if heldSent == nil { + return nil + } + return heldSent.Finish(ctx, c.res) +} + +// abort is called once the batcher is closed when some resources may not have been sent. +func (c *syncCheckpoints) abort() { + c.mu.Lock() + deferred, heldSent, scopes := c.deferred, c.heldSent, c.scopes + c.deferred = nil + c.deferredState = nil + c.heldSent = nil + c.scopes = nil + c.done = true + c.mu.Unlock() + close(c.finished) + dropMessages(deferred) + // a resolver goroutine that is stuck after the cancellation never forwards its pending checkpoints + for scope := range scopes { + dropMessages(scope.end()) + } + if heldSent != nil { + heldSent.Abort() + } +} + +// tableScope sends checkpoints after the resources its table resolver produced before them, including their relations. +// The goroutine reading the resolver's results forwards the checkpoints once it has processed everything received +// before them. +type tableScope struct { + checkpoints *syncCheckpoints + notify chan struct{} + + mu sync.Mutex + pending []message.SyncMessage + ended bool +} + +var ( + _ checkpoint.Scope = (*tableScope)(nil) + _ checkpoint.Collapsing = (*tableScope)(nil) +) + +func newTableScope(checkpoints *syncCheckpoints) *tableScope { + t := &tableScope{checkpoints: checkpoints, notify: make(chan struct{}, 1)} + checkpoints.mu.Lock() + defer checkpoints.mu.Unlock() + if checkpoints.done { + // nothing drops the checkpoints of a scope that starts after the sync, so they go straight to the sync + t.ended = true + return t + } + if checkpoints.scopes == nil { + checkpoints.scopes = make(map[*tableScope]struct{}) + } + checkpoints.scopes[t] = struct{}{} + return t +} + +func (t *tableScope) Checkpoint(ctx context.Context, msg message.SyncMessage) error { + t.mu.Lock() + if t.ended { + t.mu.Unlock() + return t.checkpoints.send(ctx, []message.SyncMessage{msg}) + } + t.pending = append(t.pending, msg) + t.mu.Unlock() + select { + case t.notify <- struct{}{}: + default: + } + return nil +} + +// Collapsed returns the scope shared by all ended table resolvers once the resolver ended. +func (t *tableScope) Collapsed() checkpoint.Scope { + t.mu.Lock() + defer t.mu.Unlock() + if t.ended { + return t.checkpoints.ended + } + return t +} + +// take returns the pending checkpoints. +func (t *tableScope) take() []message.SyncMessage { + t.mu.Lock() + defer t.mu.Unlock() + pending := t.pending + t.pending = nil + return pending +} + +// end returns the pending checkpoints. Later checkpoints go straight to the sync. +func (t *tableScope) end() []message.SyncMessage { + t.mu.Lock() + pending := t.pending + t.pending = nil + t.ended = true + t.mu.Unlock() + t.checkpoints.mu.Lock() + delete(t.checkpoints.scopes, t) + t.checkpoints.mu.Unlock() + return pending +} + +// sendMessages delivers msgs on res, and drops the ones it cannot deliver. +func sendMessages(ctx context.Context, res chan<- message.SyncMessage, msgs []message.SyncMessage) error { + for i, msg := range msgs { + if err := checkpoint.Deliver(ctx, res, msg); err != nil { + dropMessages(msgs[i+1:]) + return err + } + } + return nil +} + +func dropMessages(msgs []message.SyncMessage) { + for _, msg := range msgs { + checkpoint.Drop(msg) + } +} + +type stateRow struct { + value string + version uint64 +} + +type heldState struct { + schema *arrow.Schema + rows map[string]stateRow +} + +// holdState merges a state checkpoint into the held state, and reports false for other messages. +// Must be called with c.mu held. +func (c *syncCheckpoints) holdState(msg message.SyncMessage) bool { + tracked, isTracked := msg.(*checkpoint.Tracked) + if isTracked { + msg = tracked.Message + } + m, ok := msg.(*message.SyncCheckpoint) + if !ok { + return false + } + tableName, _ := m.Record.Schema().Metadata().GetValue(schema.MetadataTableName) + keys, keysOK := column[*array.String](m.Record, stateKeyColumn) + values, valuesOK := column[*array.String](m.Record, stateValueColumn) + versions, versioned := column[*array.Uint64](m.Record, stateVersionColumn) + if !keysOK || !valuesOK || (!versioned && m.Record.Schema().HasField(stateVersionColumn)) { + return false + } + state := c.deferredState[tableName] + if state == nil || !state.schema.Equal(m.Record.Schema()) { + if state != nil { + c.deferred = append(c.deferred, &message.SyncCheckpoint{Record: state.record()}) + } + state = &heldState{schema: m.Record.Schema(), rows: make(map[string]stateRow)} + if c.deferredState == nil { + c.deferredState = make(map[string]*heldState) + } + c.deferredState[tableName] = state + } + for i := range keys.Len() { + row := stateRow{value: values.Value(i)} + if versioned { + row.version = versions.Value(i) + if held, ok := state.rows[keys.Value(i)]; ok && held.version > row.version { + continue + } + } + state.rows[keys.Value(i)] = row + } + if isTracked { + if c.heldSent == nil { + c.heldSent = checkpoint.NewHeld(c) + } + tracked.Sent.MergeInto(c.heldSent) + } + return true +} + +func (s *heldState) record() arrow.RecordBatch { + bldr := array.NewRecordBuilder(memory.DefaultAllocator, s.schema) + defer bldr.Release() + for _, key := range slices.Sorted(maps.Keys(s.rows)) { + row := s.rows[key] + for i, field := range s.schema.Fields() { + switch field.Name { + case stateKeyColumn: + bldr.Field(i).(*array.StringBuilder).Append(key) + case stateValueColumn: + bldr.Field(i).(*array.StringBuilder).Append(row.value) + case stateVersionColumn: + bldr.Field(i).(*array.Uint64Builder).Append(row.version) + default: + bldr.Field(i).AppendNull() + } + } + } + return bldr.NewRecordBatch() +} + +func column[T arrow.Array](record arrow.RecordBatch, name string) (T, bool) { + var zero T + indices := record.Schema().FieldIndices(name) + if len(indices) != 1 { + return zero, false + } + col, ok := record.Column(indices[0]).(T) + return col, ok +} diff --git a/scheduler/checkpoint_test.go b/scheduler/checkpoint_test.go new file mode 100644 index 0000000000..c037838777 --- /dev/null +++ b/scheduler/checkpoint_test.go @@ -0,0 +1,938 @@ +package scheduler + +import ( + "context" + "errors" + "fmt" + "net" + "slices" + "strconv" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/memory" + pb "github.com/cloudquery/plugin-pb-go/pb/plugin/v3" + "github.com/cloudquery/plugin-sdk/v4/internal/checkpoint" + stateV3 "github.com/cloudquery/plugin-sdk/v4/internal/clients/state/v3" + "github.com/cloudquery/plugin-sdk/v4/internal/memdb" + serversv3 "github.com/cloudquery/plugin-sdk/v4/internal/servers/plugin/v3" + "github.com/cloudquery/plugin-sdk/v4/message" + "github.com/cloudquery/plugin-sdk/v4/plugin" + "github.com/cloudquery/plugin-sdk/v4/schema" + "github.com/rs/zerolog" + "github.com/stretchr/testify/require" + "google.golang.org/grpc" + "google.golang.org/grpc/credentials/insecure" + "google.golang.org/grpc/test/bufconn" +) + +// testConcurrency keeps shuffle-queue syncs from starting a worker for each of the default 50,000 concurrent resolvers. +const testConcurrency = 10 + +type testMarker struct { + name string + // rows that must be sent before the marker + parents []int64 + children []int64 +} + +func (*testMarker) GetTable() *schema.Table { return nil } +func (*testMarker) IsSyncMessage() bool { return true } + +const ( + checkpointParentTable = "checkpoint_parent" + checkpointChildTable = "checkpoint_child" +) + +func idColumn() schema.Column { + return schema.Column{ + Name: "id", + Type: arrow.PrimitiveTypes.Int64, + Resolver: func(_ context.Context, _ schema.ClientMeta, resource *schema.Resource, c schema.Column) error { + return resource.Set(c.Name, resource.Item.(int64)) + }, + } +} + +func childIDs(parent int64) []int64 { + return []int64{parent * 100, parent*100 + 1} +} + +func sendSlice(res chan<- any, ids []int64) { + res <- ids +} + +func sendEach(res chan<- any, ids []int64) { + for _, id := range ids { + res <- id + } +} + +// checkpointTables returns a parent table that checkpoints after each page, and a child table that checkpoints after +// its rows. +func checkpointTables(t *testing.T, pages, pageSize int, sendPage func(chan<- any, []int64)) schema.Tables { + child := &schema.Table{ + Name: checkpointChildTable, + Columns: []schema.Column{idColumn()}, + Resolver: func(ctx context.Context, _ schema.ClientMeta, parent *schema.Resource, res chan<- any) error { + ids := childIDs(parent.Item.(int64)) + res <- ids + return checkpoint.FromContext(ctx).Checkpoint(ctx, &testMarker{ + name: fmt.Sprintf("child-%d", parent.Item.(int64)), + children: ids, + }) + }, + } + parent := &schema.Table{ + Name: checkpointParentTable, + Columns: []schema.Column{idColumn()}, + Relations: schema.Tables{child}, + Resolver: func(ctx context.Context, _ schema.ClientMeta, _ *schema.Resource, res chan<- any) error { + scope := checkpoint.FromContext(ctx) + require.NotNil(t, scope) + var parents, children []int64 + for page := range pages { + ids := make([]int64, pageSize) + for i := range ids { + ids[i] = int64(page*pageSize + i + 1) + parents = append(parents, ids[i]) + children = append(children, childIDs(ids[i])...) + } + sendPage(res, ids) + err := scope.Checkpoint(ctx, &testMarker{ + name: fmt.Sprintf("page-%d", page), + parents: append([]int64(nil), parents...), + children: append([]int64(nil), children...), + }) + if err != nil { + return err + } + } + return nil + }, + } + child.Parent = parent + return schema.Tables{parent} +} + +func collectSyncMessages(res <-chan message.SyncMessage) (wait func() message.SyncMessages) { + var msgs message.SyncMessages + done := make(chan struct{}) + go func() { + defer close(done) + for msg := range res { + msgs = append(msgs, msg) + } + }() + return func() message.SyncMessages { + <-done + return msgs + } +} + +func recordIDs(record arrow.RecordBatch) []int64 { + col := record.Column(record.Schema().FieldIndices("id")[0]).(*array.Int64) + return col.Int64Values() +} + +// requireMarkersAfterRows checks every marker follows the rows it lists, and returns the markers in order. +func requireMarkersAfterRows(t *testing.T, msgs message.SyncMessages) []*testMarker { + t.Helper() + sent := map[string]map[int64]bool{ + checkpointParentTable: {}, + checkpointChildTable: {}, + } + var markers []*testMarker + for _, msg := range msgs { + switch m := msg.(type) { + case *message.SyncInsert: + table := m.GetTable().Name + for _, id := range recordIDs(m.Record) { + sent[table][id] = true + } + case *testMarker: + for _, id := range m.parents { + require.Truef(t, sent[checkpointParentTable][id], "%s sent before parent %d", m.name, id) + } + for _, id := range m.children { + require.Truef(t, sent[checkpointChildTable][id], "%s sent before child %d", m.name, id) + } + markers = append(markers, m) + } + } + return markers +} + +func TestCheckpointsFollowTheirResources(t *testing.T) { + const pages, pageSize = 4, 3 + for _, strategy := range []Strategy{StrategyDFS, StrategyRoundRobin, StrategyShuffle} { + for _, batching := range []struct { + name string + opt Option + }{ + {name: "batching", opt: WithBatchOptions(WithBatchMaxRows(1000), WithBatchTimeout(time.Hour))}, + {name: "no_batching", opt: WithoutBatching()}, + } { + for _, send := range []struct { + name string + fn func(chan<- any, []int64) + }{ + {name: "slices", fn: sendSlice}, + {name: "single_items", fn: sendEach}, + } { + t.Run(fmt.Sprintf("%s/%s/%s", strategy.String(), batching.name, send.name), func(t *testing.T) { + for range 20 { + s := NewScheduler(WithConcurrency(testConcurrency), WithStrategy(strategy), batching.opt) + res := make(chan message.SyncMessage) + wait := collectSyncMessages(res) + err := s.Sync(context.Background(), &testExecutionClient{}, checkpointTables(t, pages, pageSize, send.fn), res) + close(res) + msgs := wait() + require.NoError(t, err) + + markers := requireMarkersAfterRows(t, msgs) + require.Len(t, markers, pages+pages*pageSize) + var pageMarkers []string + for _, m := range markers { + if m.parents != nil { + pageMarkers = append(pageMarkers, m.name) + } + } + require.Equal(t, []string{"page-0", "page-1", "page-2", "page-3"}, pageMarkers) + } + }) + } + } + } +} + +func TestCheckpointsWaitForSyncWithShuffleQueue(t *testing.T) { + const pages, pageSize = 3, 2 + s := NewScheduler(WithConcurrency(testConcurrency), WithStrategy(StrategyShuffleQueue), WithBatchOptions(WithBatchMaxRows(1000), WithBatchTimeout(time.Hour))) + res := make(chan message.SyncMessage) + wait := collectSyncMessages(res) + err := s.Sync(context.Background(), &testExecutionClient{}, checkpointTables(t, pages, pageSize, sendSlice), res) + close(res) + msgs := wait() + require.NoError(t, err) + + markers := requireMarkersAfterRows(t, msgs) + require.Len(t, markers, pages+pages*pageSize) + for _, msg := range msgs[len(msgs)-len(markers):] { + require.IsType(t, &testMarker{}, msg, "checkpoints must follow every resource") + } +} + +func TestCheckpointAfterSync(t *testing.T) { + var ( + mu sync.Mutex + scopes []checkpoint.Scope + ) + table := &schema.Table{ + Name: checkpointParentTable, + Columns: []schema.Column{idColumn()}, + Resolver: func(ctx context.Context, _ schema.ClientMeta, _ *schema.Resource, res chan<- any) error { + mu.Lock() + scopes = append(scopes, checkpoint.FromContext(ctx)) + mu.Unlock() + res <- []int64{1, 2} + return nil + }, + } + + for _, strategy := range []Strategy{StrategyDFS, StrategyShuffleQueue} { + t.Run(strategy.String(), func(t *testing.T) { + scopes = nil + s := NewScheduler(WithConcurrency(testConcurrency), WithStrategy(strategy), WithBatchOptions(WithBatchMaxRows(1000), WithBatchTimeout(time.Hour))) + res := make(chan message.SyncMessage) + wait := collectSyncMessages(res) + require.NoError(t, s.Sync(context.Background(), &testExecutionClient{}, schema.Tables{table}, res)) + require.Len(t, scopes, 1) + require.NoError(t, scopes[0].Checkpoint(context.Background(), &testMarker{name: "after", parents: []int64{1, 2}})) + close(res) + + markers := requireMarkersAfterRows(t, wait()) + require.Len(t, markers, 1) + }) + } +} + +func TestCheckpointAfterCancelledSync(t *testing.T) { + var scope checkpoint.Scope + ctx, cancel := context.WithCancel(context.Background()) + table := &schema.Table{ + Name: checkpointParentTable, + Columns: []schema.Column{idColumn()}, + Resolver: func(ctx context.Context, _ schema.ClientMeta, _ *schema.Resource, res chan<- any) error { + scope = checkpoint.FromContext(ctx) + res <- []int64{1, 2} + cancel() + return nil + }, + } + + // with a table scope, Checkpoint could still be queued while the resolver goroutine winds down + s := NewScheduler(WithConcurrency(testConcurrency), WithStrategy(StrategyShuffleQueue), WithBatchOptions(WithBatchMaxRows(1000), WithBatchTimeout(time.Hour))) + res := make(chan message.SyncMessage) + wait := collectSyncMessages(res) + _ = s.Sync(ctx, &testExecutionClient{}, schema.Tables{table}, res) + require.ErrorIs(t, scope.Checkpoint(context.Background(), &testMarker{name: "after"}), checkpoint.ErrSyncAborted) + close(res) + require.Empty(t, requireMarkersAfterRows(t, wait())) +} + +// cancellingBatcher cancels the sync while it flushes. +type cancellingBatcher struct { + cancel context.CancelFunc + flushed bool +} + +func (*cancellingBatcher) process(*schema.Resource) {} +func (b *cancellingBatcher) flush() bool { + b.cancel() + return b.flushed +} +func (*cancellingBatcher) close() {} + +func TestCheckpointNotSentWhenSyncIsCancelledDuringFlush(t *testing.T) { + for _, flushed := range []bool{false, true} { + t.Run(fmt.Sprintf("flushed=%t", flushed), func(t *testing.T) { + // sending is a random choice between res and ctx.Done, so repeat to catch a checkpoint going out + for range 50 { + ctx, cancel := context.WithCancel(context.Background()) + res := make(chan message.SyncMessage, 1) + s := &syncClient{checkpoints: newSyncCheckpoints(res), msgChan: res, logger: zerolog.Nop()} + resources := make(chan *schema.Resource) + errCh := make(chan error, 1) + go func() { + errCh <- s.processResources(ctx, &cancellingBatcher{cancel: cancel, flushed: flushed}, resources) + }() + sent := &checkpoint.Sent{} + require.NoError(t, s.checkpoints.send(context.Background(), []message.SyncMessage{&checkpoint.Tracked{Message: &testMarker{name: "checkpoint"}, Sent: sent}})) + close(resources) + require.ErrorIs(t, receiveWithin(t, errCh), context.Canceled) + require.Empty(t, res) + require.True(t, sent.Aborted(), "a state client waiting for the checkpoint learns it was dropped") + } + }) + } +} + +func cursorCheckpoint(value string, version uint64) *message.SyncCheckpoint { + sc := (&schema.Table{Name: "cq_state", Columns: schema.ColumnList{ + {Name: stateKeyColumn, Type: arrow.BinaryTypes.String, PrimaryKey: true}, + {Name: stateValueColumn, Type: arrow.BinaryTypes.String}, + {Name: stateVersionColumn, Type: arrow.PrimitiveTypes.Uint64}, + }}).ToArrowSchema() + bldr := array.NewRecordBuilder(memory.DefaultAllocator, sc) + defer bldr.Release() + bldr.Field(0).(*array.StringBuilder).Append("cursor") + bldr.Field(1).(*array.StringBuilder).Append(value) + bldr.Field(2).(*array.Uint64Builder).Append(version) + return &message.SyncCheckpoint{Record: bldr.NewRecordBatch()} +} + +func stateValues(t *testing.T, msgs message.SyncMessages) map[string]string { + t.Helper() + values := map[string]string{} + for _, msg := range msgs { + if m, ok := msg.(*message.SyncCheckpoint); ok { + keys, _ := column[*array.String](m.Record, stateKeyColumn) + vals, _ := column[*array.String](m.Record, stateValueColumn) + for i := range keys.Len() { + values[keys.Value(i)] = vals.Value(i) + } + } + } + return values +} + +func TestHeldCheckpointsAreCoalesced(t *testing.T) { + table := &schema.Table{ + Name: checkpointParentTable, + Columns: []schema.Column{idColumn()}, + Resolver: func(ctx context.Context, _ schema.ClientMeta, _ *schema.Resource, res chan<- any) error { + scope := checkpoint.FromContext(ctx) + for page := range 5 { + res <- []int64{int64(page)} + if err := scope.Checkpoint(ctx, cursorCheckpoint(strconv.Itoa(page), uint64(page+1))); err != nil { + return err + } + } + // an older version does not replace the newer one + return scope.Checkpoint(ctx, cursorCheckpoint("stale", 1)) + }, + } + s := NewScheduler(WithConcurrency(testConcurrency), WithStrategy(StrategyShuffleQueue)) + res := make(chan message.SyncMessage) + wait := collectSyncMessages(res) + require.NoError(t, s.Sync(context.Background(), &testExecutionClient{}, schema.Tables{table}, res)) + close(res) + msgs := wait() + + var checkpoints int + for _, msg := range msgs { + if _, ok := msg.(*message.SyncCheckpoint); ok { + checkpoints++ + } + } + require.Equal(t, 1, checkpoints) + require.Equal(t, map[string]string{"cursor": "4"}, stateValues(t, msgs)) + require.IsType(t, &message.SyncCheckpoint{}, msgs[len(msgs)-1]) +} + +func TestPendingCheckpointsDroppedWhenSyncIsCancelled(t *testing.T) { + for _, strategy := range []Strategy{StrategyDFS, StrategyShuffleQueue} { + t.Run(strategy.String(), func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + table := &schema.Table{ + Name: checkpointParentTable, + Columns: []schema.Column{idColumn()}, + Resolver: func(ctx context.Context, _ schema.ClientMeta, _ *schema.Resource, res chan<- any) error { + res <- []int64{1, 2} + err := checkpoint.FromContext(ctx).Checkpoint(ctx, &testMarker{name: "pending", parents: []int64{1, 2}}) + cancel() + return err + }, + } + // the batch never fills, so the rows are dropped when the sync is cancelled + s := NewScheduler(WithConcurrency(testConcurrency), WithStrategy(strategy), WithBatchOptions(WithBatchMaxRows(1000), WithBatchTimeout(time.Hour))) + res := make(chan message.SyncMessage) + wait := collectSyncMessages(res) + _ = s.Sync(ctx, &testExecutionClient{}, schema.Tables{table}, res) + close(res) + require.Empty(t, requireMarkersAfterRows(t, wait())) + }) + } +} + +func TestJoinedCheckpointFollowsEveryScope(t *testing.T) { + var ( + parentScope = make(chan checkpoint.Scope, 1) + joinedSent = make(chan struct{}) + ) + parent := &schema.Table{ + Name: checkpointParentTable, + Columns: []schema.Column{idColumn()}, + Resolver: func(ctx context.Context, _ schema.ClientMeta, _ *schema.Resource, res chan<- any) error { + res <- []int64{1, 2, 3} + parentScope <- checkpoint.FromContext(ctx) + <-joinedSent + return nil + }, + } + other := &schema.Table{ + Name: checkpointChildTable, + Columns: []schema.Column{idColumn()}, + Resolver: func(ctx context.Context, _ schema.ClientMeta, _ *schema.Resource, res chan<- any) error { + res <- []int64{10, 11} + scope := checkpoint.Join(<-parentScope, checkpoint.FromContext(ctx)) + defer close(joinedSent) + return scope.Checkpoint(ctx, &testMarker{name: "joined", parents: []int64{1, 2, 3}, children: []int64{10, 11}}) + }, + } + for range 20 { + joinedSent = make(chan struct{}) + s := NewScheduler(WithConcurrency(testConcurrency), WithBatchOptions(WithBatchMaxRows(1000), WithBatchTimeout(time.Hour))) + res := make(chan message.SyncMessage) + wait := collectSyncMessages(res) + require.NoError(t, s.Sync(context.Background(), &testExecutionClient{}, schema.Tables{parent, other}, res)) + close(res) + require.Len(t, requireMarkersAfterRows(t, wait()), 1) + } +} + +func TestHeldCheckpointsAreMarkedSentWhenTheSyncCompletes(t *testing.T) { + var sent []*checkpoint.Sent + table := &schema.Table{ + Name: checkpointParentTable, + Columns: []schema.Column{idColumn()}, + Resolver: func(ctx context.Context, _ schema.ClientMeta, _ *schema.Resource, res chan<- any) error { + for page := range 3 { + res <- []int64{int64(page)} + s, err := checkpoint.Send(ctx, nil, checkpoint.FromContext(ctx), nil, cursorCheckpoint(strconv.Itoa(page), uint64(page+1))) + if err != nil { + return err + } + sent = append(sent, s) + } + return nil + }, + } + s := NewScheduler(WithConcurrency(testConcurrency), WithStrategy(StrategyShuffleQueue)) + res := make(chan message.SyncMessage) + wait := collectSyncMessages(res) + require.NoError(t, s.Sync(context.Background(), &testExecutionClient{}, schema.Tables{table}, res)) + close(res) + msgs := wait() + + require.Equal(t, map[string]string{"cursor": "2"}, stateValues(t, msgs)) + require.Len(t, sent, 3) + for _, s := range sent { + require.True(t, s.Done()) + } +} + +func TestTableScopeCheckpointAfterCancelledSync(t *testing.T) { + for range 20 { + var scope checkpoint.Scope + ctx, cancel := context.WithCancel(context.Background()) + table := &schema.Table{ + Name: checkpointParentTable, + Columns: []schema.Column{idColumn()}, + Resolver: func(ctx context.Context, _ schema.ClientMeta, _ *schema.Resource, res chan<- any) error { + scope = checkpoint.FromContext(ctx) + res <- []int64{1, 2} + time.Sleep(20 * time.Millisecond) + cancel() + return nil + }, + } + s := NewScheduler(WithConcurrency(testConcurrency), WithStrategy(StrategyDFS), WithBatchOptions(WithBatchMaxRows(1000), WithBatchTimeout(time.Hour))) + res := make(chan message.SyncMessage) + wait := collectSyncMessages(res) + _ = s.Sync(ctx, &testExecutionClient{}, schema.Tables{table}, res) + // the resolver ended, so its scope sends through the sync, which was aborted + err := scope.Checkpoint(context.Background(), &testMarker{name: "after", parents: []int64{1, 2}}) + close(res) + require.ErrorIs(t, err, checkpoint.ErrSyncAborted) + require.Empty(t, requireMarkersAfterRows(t, wait())) + } +} + +func TestShuffleQueueMergesCheckpointsOfKeysAlsoSetOutsideResolvers(t *testing.T) { + var sent []*checkpoint.Sent + table := &schema.Table{ + Name: checkpointParentTable, + Columns: []schema.Column{idColumn()}, + Resolver: func(ctx context.Context, _ schema.ClientMeta, _ *schema.Resource, res chan<- any) error { + scope := checkpoint.FromContext(ctx) + var previous []*checkpoint.Sent + for page := range 6 { + res <- []int64{int64(page)} + // every other checkpoint is sent as if the key was set outside a resolver + var pageScope checkpoint.Scope + if page%2 == 0 { + pageScope = scope + } + s, err := checkpoint.Send(ctx, nil, pageScope, previous, cursorCheckpoint(strconv.Itoa(page), uint64(page+1))) + if err != nil { + return err + } + previous = []*checkpoint.Sent{s} + sent = append(sent, s) + } + return nil + }, + } + s := NewScheduler(WithConcurrency(testConcurrency), WithStrategy(StrategyShuffleQueue)) + res := make(chan message.SyncMessage) + wait := collectSyncMessages(res) + require.NoError(t, s.Sync(context.Background(), &testExecutionClient{}, schema.Tables{table}, res)) + close(res) + msgs := wait() + + var checkpoints int + for _, msg := range msgs { + if _, ok := msg.(*message.SyncCheckpoint); ok { + checkpoints++ + } + } + require.Equal(t, 1, checkpoints, "every checkpoint is merged") + require.Equal(t, map[string]string{"cursor": "5"}, stateValues(t, msgs)) + for _, s := range sent { + require.True(t, s.Done()) + } +} + +func TestAbortDropsCheckpointsOfStuckResolvers(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + var ( + release = make(chan struct{}) + queued = make(chan struct{}) + gate = make(chan struct{}) + resume = make(chan struct{}) + res = make(chan message.SyncMessage) + ) + defer close(release) + stateClient := newCheckpointStateClient(t, res) + stuck := &schema.Table{ + Name: checkpointParentTable, + Columns: []schema.Column{idColumn()}, + // the resolver goroutine never gets past this resource, so it never forwards the checkpoint after it + PreResourceResolver: func(context.Context, schema.ClientMeta, *schema.Resource) error { + <-release + return errors.New("released") + }, + Resolver: func(ctx context.Context, _ schema.ClientMeta, _ *schema.Resource, res chan<- any) error { + res <- int64(1) + defer close(queued) + if err := stateClient.SetKey(ctx, "cursor", "1"); err != nil { + return err + } + return stateClient.Flush(ctx) + }, + } + other := &schema.Table{ + Name: checkpointChildTable, + Columns: []schema.Column{idColumn()}, + Resolver: func(ctx context.Context, _ schema.ClientMeta, _ *schema.Resource, res chan<- any) error { + res <- int64(2) + return checkpoint.FromContext(ctx).Checkpoint(ctx, &testMarker{name: "other"}) + }, + } + + var cursorSent atomic.Bool + go func() { + for msg := range res { + if _, ok := msg.(*message.SyncCheckpoint); ok { + cursorSent.Store(true) + } + // stop reading once the other table's row arrives, so the sync is cancelled while sending its checkpoint + if insert, ok := msg.(*message.SyncInsert); ok && insert.GetTable().Name == checkpointChildTable { + close(gate) + <-resume + } + } + }() + defer close(res) + defer close(resume) + + s := NewScheduler(WithConcurrency(testConcurrency), WithStrategy(StrategyDFS), WithBatchOptions(WithBatchMaxRows(1000), WithBatchTimeout(time.Hour))) + syncErr := make(chan error, 1) + go func() { + syncErr <- s.Sync(ctx, &testExecutionClient{}, schema.Tables{stuck, other}, res) + }() + receiveWithin(t, queued) + receiveWithin(t, gate) + cancel() + require.ErrorIs(t, receiveWithin(t, syncErr), context.Canceled) + require.ErrorIs(t, stateClient.Flush(context.Background()), checkpoint.ErrSyncAborted, "the checkpoint behind the stuck resolver is aborted") + require.False(t, cursorSent.Load()) +} + +func TestShuffleQueueAbortDropsCheckpointsOfStuckResolvers(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + var ( + release = make(chan struct{}) + queued = make(chan struct{}) + cancelled = make(chan struct{}) + done = make(chan struct{}) + res = make(chan message.SyncMessage) + ) + defer close(release) + defer close(done) + stateClient := newCheckpointStateClient(t, res) + stuck := &schema.Table{ + Name: checkpointParentTable, + Columns: []schema.Column{idColumn()}, + PreResourceResolver: func(context.Context, schema.ClientMeta, *schema.Resource) error { + <-release + return errors.New("released") + }, + Resolver: func(ctx context.Context, _ schema.ClientMeta, _ *schema.Resource, res chan<- any) error { + res <- int64(1) + defer close(queued) + if err := stateClient.SetKey(ctx, "cursor", "1"); err != nil { + return err + } + return stateClient.Flush(ctx) + }, + } + // sends resources after the cancel, so that the sync notices it while the stuck resolver never returns + waker := &schema.Table{ + Name: checkpointChildTable, + Columns: []schema.Column{idColumn()}, + Resolver: func(_ context.Context, _ schema.ClientMeta, _ *schema.Resource, res chan<- any) error { + <-cancelled + ids := make([]int64, 64) + for i := range ids { + ids[i] = int64(i) + } + res <- ids + return nil + }, + } + + var cursorSent atomic.Bool + go func() { + for { + select { + case msg := <-res: + if _, ok := msg.(*message.SyncCheckpoint); ok { + cursorSent.Store(true) + } + case <-done: + return + } + } + }() + + s := NewScheduler(WithConcurrency(testConcurrency), WithStrategy(StrategyShuffleQueue), WithBatchOptions(WithBatchMaxRows(1000), WithBatchTimeout(time.Hour))) + syncErr := make(chan error, 1) + go func() { + syncErr <- s.Sync(ctx, &testExecutionClient{}, schema.Tables{stuck, waker}, res) + }() + receiveWithin(t, queued) + cancel() + close(cancelled) + require.ErrorIs(t, receiveWithin(t, syncErr), context.Canceled) + require.ErrorIs(t, stateClient.Flush(context.Background()), checkpoint.ErrSyncAborted, "the held checkpoint is aborted") + require.False(t, cursorSent.Load()) +} + +func TestTableScopeStartedAfterAbortDropsCheckpoints(t *testing.T) { + c := newSyncCheckpoints(make(chan message.SyncMessage)) + close(c.loopDone) + c.abort() + + sent, err := checkpoint.Send(context.Background(), nil, newTableScope(c), nil, &testMarker{name: "late"}) + require.ErrorIs(t, err, checkpoint.ErrSyncAborted) + require.True(t, sent.Aborted()) +} + +func TestEndedTableScopesCollapse(t *testing.T) { + scopes := make(chan checkpoint.Scope, 2) + collapsedWhileRunning := make(chan bool, 2) + table := func(name string) *schema.Table { + return &schema.Table{ + Name: name, + Columns: []schema.Column{idColumn()}, + Resolver: func(ctx context.Context, _ schema.ClientMeta, _ *schema.Resource, res chan<- any) error { + scope := checkpoint.FromContext(ctx) + collapsedWhileRunning <- checkpoint.Collapse(scope) != scope + scopes <- scope + res <- int64(1) + return nil + }, + } + } + s := NewScheduler(WithConcurrency(testConcurrency), WithStrategy(StrategyDFS)) + res := make(chan message.SyncMessage) + wait := collectSyncMessages(res) + require.NoError(t, s.Sync(context.Background(), &testExecutionClient{}, schema.Tables{table(checkpointParentTable), table(checkpointChildTable)}, res)) + close(res) + wait() + + require.False(t, <-collapsedWhileRunning) + require.False(t, <-collapsedWhileRunning) + a, b := <-scopes, <-scopes + require.NotSame(t, a, b) + require.Same(t, checkpoint.Collapse(a), checkpoint.Collapse(b)) + require.Same(t, checkpoint.Collapse(a), checkpoint.Join(a, b), "joining ended scopes waits on the sync once") +} + +func TestSyncAllDropsCheckpointsAfterItReturns(t *testing.T) { + var scope checkpoint.Scope + table := &schema.Table{ + Name: checkpointParentTable, + Columns: []schema.Column{idColumn()}, + Resolver: func(ctx context.Context, _ schema.ClientMeta, _ *schema.Resource, res chan<- any) error { + scope = checkpoint.FromContext(ctx) + res <- int64(1) + return nil + }, + } + for _, strategy := range []Strategy{StrategyDFS, StrategyShuffleQueue} { + t.Run(strategy.String(), func(t *testing.T) { + s := NewScheduler(WithConcurrency(testConcurrency), WithStrategy(strategy)) + _, err := s.SyncAll(context.Background(), &testExecutionClient{}, schema.Tables{table}) + require.NoError(t, err) + // SyncAll closed its channel, so sending would panic + sent, err := checkpoint.Send(context.Background(), nil, scope, nil, &testMarker{name: "late"}) + require.ErrorIs(t, err, errResClosed) + require.True(t, sent.Aborted()) + }) + } +} + +// receiveWithin fails the test instead of hanging when nothing arrives. +func receiveWithin[T any](t *testing.T, ch <-chan T) T { + t.Helper() + select { + case v := <-ch: + return v + case <-time.After(10 * time.Second): + t.Fatal("timed out waiting for a value") + var zero T + return zero + } +} + +func TestCheckpointOfEndedResolverFollowsItsRows(t *testing.T) { + res := make(chan message.SyncMessage) + stateClient := newCheckpointStateClient(t, res) + endedScope := make(chan checkpoint.Scope, 1) + ended := &schema.Table{ + Name: checkpointParentTable, + Columns: []schema.Column{idColumn()}, + // ends with its row still in the batcher, and the cursor not yet flushed + Resolver: func(ctx context.Context, _ schema.ClientMeta, _ *schema.Resource, res chan<- any) error { + endedScope <- checkpoint.FromContext(ctx) + res <- int64(1) + return stateClient.SetKey(ctx, "cursor", "1") + }, + } + flushing := &schema.Table{ + Name: checkpointChildTable, + Columns: []schema.Column{idColumn()}, + Resolver: func(ctx context.Context, _ schema.ClientMeta, _ *schema.Resource, _ chan<- any) error { + scope := <-endedScope + deadline := time.Now().Add(10 * time.Second) + for checkpoint.Collapse(scope) == scope { + if time.Now().After(deadline) { + return errors.New("the other resolver did not end") + } + time.Sleep(time.Millisecond) + } + return stateClient.Flush(ctx) + }, + } + + s := NewScheduler(WithConcurrency(testConcurrency), WithStrategy(StrategyDFS), WithBatchOptions(WithBatchMaxRows(1000), WithBatchTimeout(time.Hour))) + wait := collectSyncMessages(res) + require.NoError(t, s.Sync(context.Background(), &testExecutionClient{}, schema.Tables{ended, flushing}, res)) + close(res) + msgs := wait() + + rowAt, cursorAt := -1, -1 + for i, msg := range msgs { + switch m := msg.(type) { + case *message.SyncError: + t.Fatal(m.Error) + case *message.SyncInsert: + if slices.Contains(recordIDs(m.Record), 1) { + rowAt = i + } + case *message.SyncCheckpoint: + if stateValues(t, message.SyncMessages{m})["cursor"] == "1" { + cursorAt = i + } + } + } + require.NotEqual(t, -1, rowAt, "the row is sent") + require.NotEqual(t, -1, cursorAt, "the cursor is sent") + require.Less(t, rowAt, cursorAt, "the cursor follows the row of the resolver that set it") +} + +func TestSyncReportsCancellationWhileTheBatcherCloses(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + table := func(name string) *schema.Table { + return &schema.Table{ + Name: name, + Columns: []schema.Column{idColumn()}, + Resolver: func(_ context.Context, _ schema.ClientMeta, _ *schema.Resource, res chan<- any) error { + res <- []int64{1, 2} + return nil + }, + } + } + // the batches never fill, so their rows go out only as the batcher closes. The first one cancels the sync before + // the other can be sent. + s := NewScheduler(WithConcurrency(testConcurrency), WithBatchOptions(WithBatchMaxRows(1000), WithBatchTimeout(time.Hour))) + res := make(chan message.SyncMessage) + go func() { + for msg := range res { + if _, ok := msg.(*message.SyncInsert); ok { + cancel() + } + } + }() + err := s.Sync(ctx, &testExecutionClient{}, schema.Tables{table(checkpointParentTable), table(checkpointChildTable)}, res) + close(res) + require.ErrorIs(t, err, context.Canceled, "the sync may have dropped rows, so it must not report success") +} + +func TestReturnedWaitsForCheckpointsBeingSent(t *testing.T) { + res := make(chan message.SyncMessage) + c := newSyncCheckpoints(res) + c.closeRes = true + close(c.loopDone) + require.NoError(t, c.complete(context.Background())) + + sent := make(chan error, 1) + go func() { + sent <- c.send(context.Background(), []message.SyncMessage{&testMarker{name: "late"}}) + }() + // wait until the send holds the read lock, blocked on res + deadline := time.Now().Add(10 * time.Second) + for c.resMu.TryLock() { + c.resMu.Unlock() + require.True(t, time.Now().Before(deadline), "the send did not start") + time.Sleep(time.Millisecond) + } + returned := make(chan struct{}) + go func() { + c.returned() + close(returned) + }() + select { + case <-returned: + t.Fatal("returned did not wait for the checkpoint being sent") + case <-time.After(100 * time.Millisecond): + } + require.IsType(t, &testMarker{}, receiveWithin(t, res)) + require.NoError(t, receiveWithin(t, sent)) + receiveWithin(t, returned) + require.ErrorIs(t, c.send(context.Background(), []message.SyncMessage{&testMarker{name: "after"}}), errResClosed) +} + +// newCheckpointStateClient returns a state client that sends its checkpoints on res, backed by an in-memory plugin. +func newCheckpointStateClient(t *testing.T, res chan<- message.SyncMessage) *stateV3.Client { + t.Helper() + p := plugin.NewPlugin("memdb", "v1.0.0", memdb.NewMemDBClient) + require.NoError(t, p.Init(context.Background(), nil, plugin.NewClientOptions{})) + listener := bufconn.Listen(1024 * 1024) + server := grpc.NewServer() + pb.RegisterPluginServer(server, &serversv3.Server{Plugin: p, Logger: zerolog.Nop()}) + go func() { _ = server.Serve(listener) }() + t.Cleanup(server.Stop) + conn, err := grpc.NewClient("passthrough:///bufnet", + grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) { return listener.Dial() }), + grpc.WithTransportCredentials(insecure.NewCredentials()), + ) + require.NoError(t, err) + t.Cleanup(func() { _ = conn.Close() }) + client, err := stateV3.NewClient(context.Background(), conn, "cq_state", stateV3.WithCheckpoints(res)) + require.NoError(t, err) + return client +} + +func TestCheckpointsDroppedOnCancelAreAborted(t *testing.T) { + tracked := func() (*checkpoint.Sent, []message.SyncMessage) { + sent := &checkpoint.Sent{} + return sent, []message.SyncMessage{&checkpoint.Tracked{Message: &testMarker{name: "checkpoint"}, Sent: sent}} + } + cancelled, cancel := context.WithCancel(context.Background()) + cancel() + + t.Run("send while the sync runs", func(t *testing.T) { + c := newSyncCheckpoints(make(chan message.SyncMessage)) + sent, msgs := tracked() + require.ErrorIs(t, c.send(cancelled, msgs), context.Canceled) + require.True(t, sent.Aborted()) + }) + t.Run("send while the batcher closes", func(t *testing.T) { + c := newSyncCheckpoints(make(chan message.SyncMessage)) + close(c.loopDone) + sent, msgs := tracked() + require.ErrorIs(t, c.sendAfterSync(cancelled, msgs), context.Canceled) + require.True(t, sent.Aborted()) + }) + t.Run("held until an aborted sync ends", func(t *testing.T) { + c := newSyncCheckpoints(make(chan message.SyncMessage)) + sent, msgs := tracked() + require.NoError(t, c.Checkpoint(context.Background(), msgs[0])) + require.False(t, sent.Aborted()) + close(c.loopDone) + c.abort() + require.True(t, sent.Aborted()) + }) +} diff --git a/scheduler/scheduler.go b/scheduler/scheduler.go index b5e04aba75..d8cef522bc 100644 --- a/scheduler/scheduler.go +++ b/scheduler/scheduler.go @@ -112,6 +112,13 @@ func WithInvocationID(invocationID string) Option { } } +// withResClosedOnReturn tells Sync that the caller closes res once Sync returns, so later checkpoints must not be sent. +func withResClosedOnReturn() SyncOption { + return func(s *syncClient) { + s.checkpoints.closeRes = true + } +} + func WithShard(num int32, total int32) SyncOption { return func(s *syncClient) { s.shard = &shard{num: num, total: total} @@ -166,6 +173,8 @@ type syncClient struct { // message channel for sending SyncError messages msgChan chan<- message.SyncMessage + checkpoints *syncCheckpoints + shard *shard } @@ -213,7 +222,7 @@ func (s *Scheduler) SyncAll(ctx context.Context, client schema.ClientMeta, table var err error go func() { defer close(res) - err = s.Sync(ctx, client, tables, res) + err = s.Sync(ctx, client, tables, res, withResClosedOnReturn()) }() // nolint:prealloc var messages message.SyncMessages @@ -241,6 +250,7 @@ func (s *Scheduler) Sync(ctx context.Context, client schema.ClientMeta, tables s logger: s.logger, invocationID: s.invocationID, msgChan: res, + checkpoints: newSyncCheckpoints(res), } for _, opt := range opts { opt(syncClient) @@ -286,19 +296,55 @@ func (s *Scheduler) Sync(ctx context.Context, client schema.ClientMeta, tables s } }() + defer syncClient.checkpoints.returned() b := s.batchSettings.getBatcher(ctx, res, s.logger) - defer b.close() // wait for all resources to be processed + err := syncClient.processResources(ctx, b, resources) + close(syncClient.checkpoints.loopDone) + b.close() // wait for all resources to be processed + if err != nil || ctx.Err() != nil { + syncClient.checkpoints.abort() + if err == nil { + // cancelled while the batcher closed, which may have dropped resources + err = context.Cause(ctx) + } + return err + } + return syncClient.checkpoints.complete(ctx) +} + +// processResources hands resources to the batcher until resources is closed. Checkpoints go out after the resources +// received before them. +func (s *syncClient) processResources(ctx context.Context, b batcherInterface, resources <-chan *schema.Resource) error { done := ctx.Done() // no need to do the lookups in loop - for resource := range resources { + for { select { - case <-done: - s.logger.Debug().Msg("sync context cancelled") - return context.Cause(ctx) - default: - b.process(resource) + case resource, ok := <-resources: + if !ok { + return context.Cause(ctx) + } + select { + case <-done: + s.logger.Debug().Msg("sync context cancelled") + return context.Cause(ctx) + default: + b.process(resource) + } + case msgs := <-s.checkpoints.requests: + // Every table's batch goes out, not only the ones the checkpoint follows: those would be the tables of the + // scope's resources and their relations, which the batcher does not track. A state flush of the source sends a + // checkpoint per group of resolver scopes, so it can flush the batches a few times, which costs smaller batches + // only. + // A cancelled sync may drop resources, and their checkpoints must not go out. + if !b.flush() || ctx.Err() != nil { + s.logger.Debug().Msg("sync context cancelled") + dropMessages(msgs) + return context.Cause(ctx) + } + if err := sendMessages(ctx, s.msgChan, msgs); err != nil { + return err + } } } - return context.Cause(ctx) } func (s *syncClient) logTablesMetrics(tables schema.Tables, client Client) { diff --git a/scheduler/scheduler_dfs.go b/scheduler/scheduler_dfs.go index b810bedb1c..7aa04fe62d 100644 --- a/scheduler/scheduler_dfs.go +++ b/scheduler/scheduler_dfs.go @@ -8,6 +8,7 @@ import ( "time" "github.com/cloudquery/plugin-sdk/v4/helpers" + "github.com/cloudquery/plugin-sdk/v4/internal/checkpoint" "github.com/cloudquery/plugin-sdk/v4/message" "github.com/cloudquery/plugin-sdk/v4/scheduler/batchsender" "github.com/cloudquery/plugin-sdk/v4/scheduler/metrics" @@ -90,6 +91,8 @@ func (s *syncClient) resolveTableDfs(ctx context.Context, table *schema.Table, c defer span.End() logger := s.logger.With().Str("table", table.Name).Str("client", clientName).Logger() ctx = logger.WithContext(ctx) + scope := newTableScope(s.checkpoints) + ctx = checkpoint.WithScope(ctx, scope) startTime := time.Now() if parent == nil { // Log only for root tables, otherwise we spam too much. @@ -148,14 +151,32 @@ func (s *syncClient) resolveTableDfs(ctx context.Context, table *schema.Table, c batchSender := batchsender.NewBatchSender(func(item any) { s.resolveResourcesDfs(ctx, table, client, parent, item, resolvedResources, depth) }) - for r := range res { - batchSender.Send(r) + for done := false; !done; { + select { + case r, ok := <-res: + if !ok { + done = true + break + } + batchSender.Send(r) + case <-scope.notify: + // resolving is synchronous, so everything received so far is sent once Flush returns + batchSender.Flush() + s.sendCheckpoints(ctx, scope.take()) + } } batchSender.Close() + s.sendCheckpoints(ctx, scope.end()) // we don't need any waitgroups here because we are waiting for the channel to close } +func (s *syncClient) sendCheckpoints(ctx context.Context, msgs []message.SyncMessage) { + if err := s.checkpoints.send(ctx, msgs); err != nil { + s.logger.Warn().Err(err).Int("checkpoints", len(msgs)).Msg("dropped state checkpoints") + } +} + func (s *syncClient) resolveResourcesDfs(ctx context.Context, table *schema.Table, client schema.ClientMeta, parent *schema.Resource, resources any, resolvedResources chan<- *schema.Resource, depth int) { resourcesSlice := helpers.InterfaceSlice(resources) if len(resourcesSlice) == 0 { diff --git a/scheduler/scheduler_shuffle_queue.go b/scheduler/scheduler_shuffle_queue.go index 924b3e1eed..00290c022a 100644 --- a/scheduler/scheduler_shuffle_queue.go +++ b/scheduler/scheduler_shuffle_queue.go @@ -3,6 +3,8 @@ package scheduler import ( "context" + "github.com/cloudquery/plugin-sdk/v4/internal/checkpoint" + "github.com/cloudquery/plugin-sdk/v4/scheduler/queue" "github.com/cloudquery/plugin-sdk/v4/schema" ) @@ -46,5 +48,6 @@ func (s *syncClient) syncShuffleQueue(ctx context.Context, resolvedResources cha Client: tc.client, }) } - scheduler.Sync(ctx, queueClients, resolvedResources, s.msgChan) + // queue workers resolve relations on other goroutines, so checkpoints wait for the whole sync + scheduler.Sync(checkpoint.WithScope(ctx, s.checkpoints), queueClients, resolvedResources, s.msgChan) } diff --git a/serve/checkpoint_test.go b/serve/checkpoint_test.go new file mode 100644 index 0000000000..52938d30c4 --- /dev/null +++ b/serve/checkpoint_test.go @@ -0,0 +1,932 @@ +package serve + +import ( + "context" + "errors" + "fmt" + "io" + "net" + "slices" + "strconv" + "sync" + "testing" + "time" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/memory" + pbdiscoveryv1 "github.com/cloudquery/plugin-pb-go/pb/discovery/v1" + pb "github.com/cloudquery/plugin-pb-go/pb/plugin/v3" + "github.com/cloudquery/plugin-sdk/v4/internal/checkpoint" + stateV3 "github.com/cloudquery/plugin-sdk/v4/internal/clients/state/v3" + "github.com/cloudquery/plugin-sdk/v4/internal/memdb" + discoveryServerV1 "github.com/cloudquery/plugin-sdk/v4/internal/servers/discovery/v1" + serversv3 "github.com/cloudquery/plugin-sdk/v4/internal/servers/plugin/v3" + "github.com/cloudquery/plugin-sdk/v4/message" + "github.com/cloudquery/plugin-sdk/v4/plugin" + "github.com/cloudquery/plugin-sdk/v4/scheduler" + "github.com/cloudquery/plugin-sdk/v4/schema" + "github.com/cloudquery/plugin-sdk/v4/state" + "github.com/rs/zerolog" + "github.com/stretchr/testify/require" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/credentials/insecure" + "google.golang.org/grpc/status" +) + +// servePlugin serves p over an in-memory listener and returns a client connection to it. +func servePlugin(t *testing.T, p *plugin.Plugin) *grpc.ClientConn { + t.Helper() + srv := Plugin(p, WithArgs("serve"), WithTestListener()) + ctx, cancel := context.WithCancel(context.Background()) + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + if err := srv.Serve(ctx); err != nil { + t.Error(err) + } + }() + t.Cleanup(func() { + cancel() + wg.Wait() + }) + + // nolint:staticcheck + conn, err := grpc.DialContext(ctx, "bufnet", grpc.WithContextDialer(srv.bufPluginDialer), grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithBlock()) + require.NoError(t, err) + t.Cleanup(func() { _ = conn.Close() }) + return conn +} + +// serveBackend serves p on a TCP address, as state clients dial backends by address. +func serveBackend(t *testing.T, p *plugin.Plugin) string { + t.Helper() + listener, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + grpcServer := grpc.NewServer() + pb.RegisterPluginServer(grpcServer, &serversv3.Server{Plugin: p, Logger: zerolog.Nop()}) + pbdiscoveryv1.RegisterDiscoveryServer(grpcServer, &discoveryServerV1.Server{Versions: []int32{3}}) + go func() { _ = grpcServer.Serve(listener) }() + t.Cleanup(grpcServer.Stop) + return listener.Addr().String() +} + +func insertRecord(t *testing.T, table *schema.Table, ids ...int64) []byte { + t.Helper() + bldr := array.NewRecordBuilder(memory.DefaultAllocator, table.ToArrowSchema()) + defer bldr.Release() + bldr.Field(0).(*array.Int64Builder).AppendValues(ids, nil) + recordBytes, err := pb.RecordToBytes(bldr.NewRecordBatch()) + require.NoError(t, err) + return recordBytes +} + +func TestInitReportsFlushSupport(t *testing.T) { + for _, supported := range []bool{false, true} { + t.Run(fmt.Sprintf("supported=%t", supported), func(t *testing.T) { + var opts []plugin.Option + if supported { + opts = append(opts, plugin.WithFlushSupport()) + } + c := pb.NewPluginClient(servePlugin(t, plugin.NewPlugin("memdb", "v1.0.0", memdb.NewMemDBClient, opts...))) + res, err := c.Init(context.Background(), &pb.Init_Request{}) + require.NoError(t, err) + require.Equal(t, supported, res.SupportsFlush) + }) + } +} + +func TestWriteStreamFlush(t *testing.T) { + ctx := testContext(t) + c := pb.NewPluginClient(servePlugin(t, plugin.NewPlugin("memdb", "v1.0.0", memdb.NewMemDBClient, plugin.WithFlushSupport()))) + _, err := c.Init(ctx, &pb.Init_Request{}) + require.NoError(t, err) + + table := &schema.Table{Name: "test_flush", Columns: []schema.Column{{Name: "id", Type: arrow.PrimitiveTypes.Int64}}} + tableBytes, err := pb.SchemaToBytes(table.ToArrowSchema()) + require.NoError(t, err) + + stream, err := c.WriteStream(ctx) + require.NoError(t, err) + require.NoError(t, stream.Send(&pb.Write_Request{Message: &pb.Write_Request_MigrateTable{MigrateTable: &pb.Write_MessageMigrateTable{Table: tableBytes}}})) + require.NoError(t, stream.Send(&pb.Write_Request{Message: &pb.Write_Request_Insert{Insert: &pb.Write_MessageInsert{Record: insertRecord(t, table, 1, 2)}}})) + require.NoError(t, stream.Send(&pb.Write_Request{Message: &pb.Write_Request_Flush{Flush: &pb.Write_MessageFlush{Id: 7}}})) + res, err := stream.Recv() + require.NoError(t, err) + require.EqualValues(t, 7, res.FlushId) + + readClient, err := c.Read(ctx, &pb.Read_Request{Table: tableBytes}) + require.NoError(t, err) + var rows int64 + for { + res, err := readClient.Recv() + if errors.Is(err, io.EOF) { + break + } + require.NoError(t, err) + record, err := pb.NewRecordFromBytes(res.Record) + require.NoError(t, err) + rows += record.NumRows() + } + require.EqualValues(t, 2, rows) + + require.NoError(t, stream.CloseSend()) + _, err = stream.Recv() + require.ErrorIs(t, err, io.EOF) +} + +func TestWriteStreamWithoutFlushSupport(t *testing.T) { + ctx := testContext(t) + c := pb.NewPluginClient(servePlugin(t, plugin.NewPlugin("memdb", "v1.0.0", memdb.NewMemDBClient))) + _, err := c.Init(ctx, &pb.Init_Request{}) + require.NoError(t, err) + + stream, err := c.WriteStream(ctx) + require.NoError(t, err) + _, err = stream.Recv() + require.Equal(t, codes.Unimplemented, status.Code(err)) +} + +func TestWriteStreamFlushFailure(t *testing.T) { + ctx := testContext(t) + c := pb.NewPluginClient(servePlugin(t, plugin.NewPlugin("memdb", "v1.0.0", memdb.GetNewClient(memdb.WithErrOnWrite()), plugin.WithFlushSupport()))) + _, err := c.Init(ctx, &pb.Init_Request{}) + require.NoError(t, err) + + stream, err := c.WriteStream(ctx) + require.NoError(t, err) + // the plugin may fail before or after receiving the flush + _ = stream.Send(&pb.Write_Request{Message: &pb.Write_Request_Flush{Flush: &pb.Write_MessageFlush{Id: 1}}}) + _, err = stream.Recv() + require.Error(t, err) + require.NotErrorIs(t, err, io.EOF) +} + +func TestWriteRejectsFlush(t *testing.T) { + ctx := testContext(t) + c := pb.NewPluginClient(servePlugin(t, plugin.NewPlugin("memdb", "v1.0.0", memdb.NewMemDBClient, plugin.WithFlushSupport()))) + _, err := c.Init(ctx, &pb.Init_Request{}) + require.NoError(t, err) + + stream, err := c.Write(ctx) + require.NoError(t, err) + require.NoError(t, stream.Send(&pb.Write_Request{Message: &pb.Write_Request_Flush{Flush: &pb.Write_MessageFlush{Id: 1}}})) + _, err = stream.CloseAndRecv() + require.Equal(t, codes.InvalidArgument, status.Code(err)) +} + +const ( + checkpointPages = 3 + checkpointPageSize = 4 +) + +var checkpointTable = &schema.Table{ + Name: "test_checkpoint_items", + Columns: []schema.Column{{ + Name: "id", + Type: arrow.PrimitiveTypes.Int64, + Resolver: func(_ context.Context, _ schema.ClientMeta, resource *schema.Resource, c schema.Column) error { + return resource.Set(c.Name, resource.Item.(int64)) + }, + }}, +} + +type checkpointSourceClient struct { + plugin.UnimplementedDestination + strategy scheduler.Strategy + pages int +} + +func (*checkpointSourceClient) ID() string { return "checkpoint_source" } + +func (*checkpointSourceClient) Close(context.Context) error { return nil } + +func (*checkpointSourceClient) Tables(context.Context, plugin.TableOptions) (schema.Tables, error) { + return schema.Tables{checkpointTable}, nil +} + +// Sync syncs pages of items, checkpointing the page number as the cursor after each page. It sets one more key in the +// resolver and flushes it after the scheduler returns. +func (c *checkpointSourceClient) Sync(ctx context.Context, options plugin.SyncOptions, res chan<- message.SyncMessage) error { + stateClient, err := state.NewConnectedClientWithOptions(ctx, options.BackendOptions, state.ConnectionOptions{}, state.ClientOptions{Checkpoints: res}) + if err != nil { + return err + } + defer stateClient.Close() + + table := *checkpointTable + table.Resolver = func(ctx context.Context, _ schema.ClientMeta, _ *schema.Resource, res chan<- any) error { + for page := range c.pages { + ids := make([]int64, checkpointPageSize) + for i := range ids { + ids[i] = int64(page*checkpointPageSize + i) + } + res <- ids + if err := stateClient.SetKey(ctx, "cursor", strconv.Itoa(page)); err != nil { + return err + } + if err := stateClient.Flush(ctx); err != nil { + return err + } + } + return stateClient.SetKey(ctx, "finished", "true") + } + s := scheduler.NewScheduler(scheduler.WithStrategy(c.strategy), scheduler.WithBatchOptions(scheduler.WithBatchMaxRows(1000), scheduler.WithBatchTimeout(time.Hour))) + if err := s.Sync(ctx, c, schema.Tables{&table}, res); err != nil { + return err + } + return stateClient.Flush(ctx) +} + +func newCheckpointSource(ctx context.Context, logger zerolog.Logger, spec []byte, opts plugin.NewClientOptions) (plugin.Client, error) { + return newCheckpointSourceWith(scheduler.StrategyDFS, checkpointPages)(ctx, logger, spec, opts) +} + +func newCheckpointSourceWith(strategy scheduler.Strategy, pages int) plugin.NewClientFunc { + return func(context.Context, zerolog.Logger, []byte, plugin.NewClientOptions) (plugin.Client, error) { + return &checkpointSourceClient{strategy: strategy, pages: pages}, nil + } +} + +func readState(t *testing.T, conn *grpc.ClientConn, tableName string) map[string]string { + t.Helper() + tableBytes, err := pb.SchemaToBytes(stateV3.Table(tableName).ToArrowSchema()) + require.NoError(t, err) + readClient, err := pb.NewPluginClient(conn).Read(context.Background(), &pb.Read_Request{Table: tableBytes}) + require.NoError(t, err) + values := map[string]string{} + for { + res, err := readClient.Recv() + if errors.Is(err, io.EOF) { + return values + } + require.NoError(t, err) + record, err := pb.NewRecordFromBytes(res.Record) + require.NoError(t, err) + addStateRows(values, record) + } +} + +func addStateRows(values map[string]string, record arrow.RecordBatch) { + keys, vals := record.Column(0).(*array.String), record.Column(1).(*array.String) + for i := range keys.Len() { + values[keys.Value(i)] = vals.Value(i) + } +} + +func TestSyncCheckpoints(t *testing.T) { + for _, supportsCheckpoints := range []bool{true, false} { + t.Run(fmt.Sprintf("supports_checkpoints=%t", supportsCheckpoints), func(t *testing.T) { + ctx := testContext(t) + backendPlugin := plugin.NewPlugin("memdb", "v1.0.0", memdb.NewMemDBClient, plugin.WithFlushSupport()) + require.NoError(t, backendPlugin.Init(ctx, nil, plugin.NewClientOptions{})) + backendAddr := serveBackend(t, backendPlugin) + + c := pb.NewPluginClient(servePlugin(t, plugin.NewPlugin("checkpoint_source", "v1.0.0", newCheckpointSource))) + _, err := c.Init(ctx, &pb.Init_Request{}) + require.NoError(t, err) + + syncClient, err := c.Sync(ctx, &pb.Sync_Request{ + Backend: &pb.Sync_BackendOptions{TableName: "cq_state", Connection: backendAddr}, + SupportsCheckpoints: supportsCheckpoints, + }) + require.NoError(t, err) + + sent := map[int64]bool{} + checkpointed := map[string]string{} + for { + res, err := syncClient.Recv() + if errors.Is(err, io.EOF) { + break + } + require.NoError(t, err) + switch m := res.Message.(type) { + case *pb.Sync_Response_Insert: + record, err := pb.NewRecordFromBytes(m.Insert.Record) + require.NoError(t, err) + for _, id := range record.Column(record.Schema().FieldIndices("id")[0]).(*array.Int64).Int64Values() { + sent[id] = true + } + case *pb.Sync_Response_Checkpoint: + require.True(t, supportsCheckpoints) + record, err := pb.NewRecordFromBytes(m.Checkpoint.Record) + require.NoError(t, err) + values := map[string]string{} + addStateRows(values, record) + if cursor, ok := values["cursor"]; ok { + page, err := strconv.Atoi(cursor) + require.NoError(t, err) + for id := range int64((page + 1) * checkpointPageSize) { + require.Truef(t, sent[id], "checkpoint for page %d sent before row %d", page, id) + } + } + if _, ok := values["finished"]; ok { + require.Len(t, sent, checkpointPages*checkpointPageSize) + } + for k, v := range values { + checkpointed[k] = v + } + } + } + require.Len(t, sent, checkpointPages*checkpointPageSize) + + backendState := readState(t, servePluginConn(t, backendAddr), "cq_state") + want := map[string]string{"cursor": strconv.Itoa(checkpointPages - 1), "finished": "true"} + if supportsCheckpoints { + require.Equal(t, want, checkpointed) + require.Empty(t, backendState, "the client persists checkpoints") + } else { + require.Empty(t, checkpointed) + require.Equal(t, want, backendState) + } + }) + } +} + +func servePluginConn(t *testing.T, addr string) *grpc.ClientConn { + t.Helper() + conn, err := grpc.NewClient(addr, grpc.WithTransportCredentials(insecure.NewCredentials())) + require.NoError(t, err) + t.Cleanup(func() { _ = conn.Close() }) + return conn +} + +type recordingScope struct { + msgs []message.SyncMessage + err error +} + +func (s *recordingScope) Checkpoint(_ context.Context, msg message.SyncMessage) error { + if s.err != nil { + return s.err + } + s.msgs = append(s.msgs, msg) + return nil +} + +func checkpointValues(t *testing.T, msgs ...message.SyncMessage) map[string]string { + t.Helper() + values := map[string]string{} + for _, msg := range msgs { + if tracked, ok := msg.(*checkpoint.Tracked); ok { + msg = tracked.Message + } + require.IsType(t, &message.SyncCheckpoint{}, msg) + addStateRows(values, msg.(*message.SyncCheckpoint).Record) + } + return values +} + +// deliver sends what a scope received as the scheduler would, once the scope's resources are sent. +func deliver(t *testing.T, res chan<- message.SyncMessage, msgs ...message.SyncMessage) { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + for _, msg := range msgs { + require.NoError(t, checkpoint.Deliver(ctx, res, msg)) + } +} + +func newCheckpointStateClient(t *testing.T, tableName string) (*stateV3.Client, chan message.SyncMessage, *grpc.ClientConn) { + t.Helper() + ctx := testContext(t) + conn := servePlugin(t, plugin.NewPlugin("memdb", "v1.0.0", memdb.NewMemDBClient)) + _, err := pb.NewPluginClient(conn).Init(ctx, &pb.Init_Request{}) + require.NoError(t, err) + res := make(chan message.SyncMessage, 10) + stateClient, err := stateV3.NewClient(ctx, conn, tableName, stateV3.WithCheckpoints(res)) + require.NoError(t, err) + return stateClient, res, conn +} + +func TestStateClientCheckpointsPerScope(t *testing.T) { + ctx := testContext(t) + stateClient, res, conn := newCheckpointStateClient(t, "test_scopes") + + scopeA, scopeB := &recordingScope{}, &recordingScope{} + ctxA, ctxB := checkpoint.WithScope(ctx, scopeA), checkpoint.WithScope(ctx, scopeB) + require.NoError(t, stateClient.SetKey(ctxA, "a", "1")) + require.NoError(t, stateClient.SetKey(ctxB, "b", "1")) + require.NoError(t, stateClient.SetKey(ctx, "c", "1")) + require.NoError(t, stateClient.SetKey(ctxA, "shared", "1")) + require.NoError(t, stateClient.SetKey(ctxB, "shared", "2")) + + require.NoError(t, stateClient.Flush(ctxA)) + require.Len(t, res, 1) + require.Equal(t, map[string]string{"c": "1"}, checkpointValues(t, receive(t, res))) + + // the shared key waits for both scopes: each gets a barrier, and the last one to reach it sends the checkpoint + require.Len(t, scopeA.msgs, 2) + require.Len(t, scopeB.msgs, 2) + var barriers []message.SyncMessage + for _, scope := range []*recordingScope{scopeA, scopeB} { + var tracked []message.SyncMessage + for _, msg := range scope.msgs { + if _, ok := msg.(*checkpoint.Barrier); ok { + barriers = append(barriers, msg) + } else { + tracked = append(tracked, msg) + } + } + require.Len(t, tracked, 1) + } + require.Equal(t, map[string]string{"a": "1"}, checkpointValues(t, scopeA.msgs[slices.IndexFunc(scopeA.msgs, isCheckpoint)])) + require.Equal(t, map[string]string{"b": "1"}, checkpointValues(t, scopeB.msgs[slices.IndexFunc(scopeB.msgs, isCheckpoint)])) + require.Len(t, barriers, 2) + deliver(t, res, barriers[0]) + require.Empty(t, res) + deliver(t, res, barriers[1]) + require.Equal(t, map[string]string{"shared": "2"}, checkpointValues(t, receive(t, res))) + + require.NoError(t, stateClient.Flush(ctx)) + require.Len(t, scopeA.msgs, 2) + require.Len(t, scopeB.msgs, 2) + require.Empty(t, res) + + require.Empty(t, readState(t, conn, "test_scopes"), "checkpoints are not written to the backend") +} + +func isCheckpoint(msg message.SyncMessage) bool { + _, ok := msg.(*checkpoint.Tracked) + return ok +} + +func TestStateClientCheckpointWaitsForPreviousCheckpointOfKey(t *testing.T) { + ctx := testContext(t) + stateClient, res, _ := newCheckpointStateClient(t, "test_previous") + + scopeA, scopeB := &recordingScope{}, &recordingScope{} + require.NoError(t, stateClient.SetKey(checkpoint.WithScope(ctx, scopeA), "cursor", "1")) + require.NoError(t, stateClient.Flush(ctx)) + require.NoError(t, stateClient.SetKey(checkpoint.WithScope(ctx, scopeB), "cursor", "2")) + require.NoError(t, stateClient.Flush(ctx)) + require.Len(t, scopeA.msgs, 1) + require.Len(t, scopeB.msgs, 1) + + // scope B finishes first, but its checkpoint must not overtake scope A's resources + deliver(t, res, scopeB.msgs...) + require.Empty(t, res) + deliver(t, res, scopeA.msgs...) + require.Equal(t, map[string]string{"cursor": "1"}, checkpointValues(t, receive(t, res))) + require.Equal(t, map[string]string{"cursor": "2"}, checkpointValues(t, receive(t, res))) + + // once sent, a checkpoint holds back nothing + scopeC := &recordingScope{} + require.NoError(t, stateClient.SetKey(checkpoint.WithScope(ctx, scopeC), "cursor", "3")) + require.NoError(t, stateClient.Flush(ctx)) + require.Len(t, scopeC.msgs, 1) + require.IsType(t, &checkpoint.Tracked{}, scopeC.msgs[0]) +} + +func TestStateClientKeepsKeysOfFailedCheckpoints(t *testing.T) { + ctx := testContext(t) + stateClient, res, _ := newCheckpointStateClient(t, "test_failed_scopes") + + aborted := &recordingScope{err: checkpoint.ErrSyncAborted} + require.NoError(t, stateClient.SetKey(checkpoint.WithScope(ctx, aborted), "a", "1")) + require.NoError(t, stateClient.SetKey(ctx, "b", "1")) + + // the failed scope does not stop other keys from being sent + require.ErrorIs(t, stateClient.Flush(ctx), state.ErrSyncAborted) + require.Equal(t, map[string]string{"b": "1"}, checkpointValues(t, receive(t, res))) + + aborted.err = nil + require.NoError(t, stateClient.Flush(ctx)) + require.Equal(t, map[string]string{"a": "1"}, checkpointValues(t, aborted.msgs...)) + require.Empty(t, res) +} + +// heldFlushClient hands each flush to the test instead of completing it. +type heldFlushClient struct { + plugin.UnimplementedDestination + flushes chan *message.WriteFlush +} + +func (*heldFlushClient) Close(context.Context) error { return nil } +func (*heldFlushClient) Tables(context.Context, plugin.TableOptions) (schema.Tables, error) { + return nil, nil +} +func (*heldFlushClient) Sync(context.Context, plugin.SyncOptions, chan<- message.SyncMessage) error { + return plugin.ErrNotImplemented +} + +func (c *heldFlushClient) Write(_ context.Context, msgs <-chan message.WriteMessage) error { + for msg := range msgs { + if flush, ok := msg.(*message.WriteFlush); ok { + c.flushes <- flush + } + } + return nil +} + +func TestWriteStreamAcksFlushOnlyAfterCommit(t *testing.T) { + ctx := testContext(t) + client := &heldFlushClient{flushes: make(chan *message.WriteFlush)} + c := pb.NewPluginClient(servePlugin(t, plugin.NewPlugin("held", "v1.0.0", func(context.Context, zerolog.Logger, []byte, plugin.NewClientOptions) (plugin.Client, error) { + return client, nil + }, plugin.WithFlushSupport()))) + _, err := c.Init(ctx, &pb.Init_Request{}) + require.NoError(t, err) + + stream, err := c.WriteStream(ctx) + require.NoError(t, err) + responses := make(chan *pb.Write_Response) + errs := make(chan error, 1) + go func() { + for { + res, err := stream.Recv() + if err != nil { + errs <- err + return + } + responses <- res + } + }() + + require.NoError(t, stream.Send(&pb.Write_Request{Message: &pb.Write_Request_Flush{Flush: &pb.Write_MessageFlush{Id: 3}}})) + flush := receive(t, client.flushes) + select { + case res := <-responses: + t.Fatalf("flush %d acknowledged before it completed", res.FlushId) + case err := <-errs: + t.Fatal(err) + case <-time.After(200 * time.Millisecond): + } + flush.Done(nil) + require.EqualValues(t, 3, receive(t, responses).FlushId) + + require.NoError(t, stream.Send(&pb.Write_Request{Message: &pb.Write_Request_Flush{Flush: &pb.Write_MessageFlush{Id: 4}}})) + receive(t, client.flushes).Done(errors.New("disk full")) + res := receive(t, responses) + require.EqualValues(t, 4, res.FlushId) + require.Equal(t, "disk full", res.FlushError) + + // the stream stays open after a failed flush + require.NoError(t, stream.Send(&pb.Write_Request{Message: &pb.Write_Request_Flush{Flush: &pb.Write_MessageFlush{Id: 5}}})) + receive(t, client.flushes).Done(nil) + res = receive(t, responses) + require.EqualValues(t, 5, res.FlushId) + require.Empty(t, res.FlushError) + select { + case err := <-errs: + t.Fatalf("stream ended: %v", err) + default: + } +} + +func TestWriteStreamRejectsEmptyMessage(t *testing.T) { + ctx := testContext(t) + c := pb.NewPluginClient(servePlugin(t, plugin.NewPlugin("memdb", "v1.0.0", memdb.NewMemDBClient, plugin.WithFlushSupport()))) + _, err := c.Init(ctx, &pb.Init_Request{}) + require.NoError(t, err) + + stream, err := c.WriteStream(ctx) + require.NoError(t, err) + require.NoError(t, stream.Send(&pb.Write_Request{})) + _, err = stream.Recv() + require.Equal(t, codes.InvalidArgument, status.Code(err)) +} + +type noFlushClient struct { + plugin.Client +} + +func (noFlushClient) SupportsFlush() bool { return false } + +func TestClientCanTurnOffFlushSupport(t *testing.T) { + ctx := testContext(t) + c := pb.NewPluginClient(servePlugin(t, plugin.NewPlugin("memdb", "v1.0.0", func(ctx context.Context, logger zerolog.Logger, spec []byte, opts plugin.NewClientOptions) (plugin.Client, error) { + client, err := memdb.NewMemDBClient(ctx, logger, spec, opts) + return noFlushClient{Client: client}, err + }, plugin.WithFlushSupport()))) + res, err := c.Init(ctx, &pb.Init_Request{}) + require.NoError(t, err) + require.False(t, res.SupportsFlush) + + stream, err := c.WriteStream(ctx) + require.NoError(t, err) + _, err = stream.Recv() + require.Equal(t, codes.Unimplemented, status.Code(err)) +} + +func TestShuffleQueueCoalescesFlushedCheckpoints(t *testing.T) { + const pages = 100 + ctx := testContext(t) + backendPlugin := plugin.NewPlugin("memdb", "v1.0.0", memdb.NewMemDBClient, plugin.WithFlushSupport()) + require.NoError(t, backendPlugin.Init(ctx, nil, plugin.NewClientOptions{})) + backendAddr := serveBackend(t, backendPlugin) + + c := pb.NewPluginClient(servePlugin(t, plugin.NewPlugin("checkpoint_source", "v1.0.0", newCheckpointSourceWith(scheduler.StrategyShuffleQueue, pages)))) + _, err := c.Init(ctx, &pb.Init_Request{}) + require.NoError(t, err) + syncClient, err := c.Sync(ctx, &pb.Sync_Request{ + Backend: &pb.Sync_BackendOptions{TableName: "cq_state", Connection: backendAddr}, + SupportsCheckpoints: true, + }) + require.NoError(t, err) + + var rows int64 + var cursors []string + for { + res, err := syncClient.Recv() + if errors.Is(err, io.EOF) { + break + } + require.NoError(t, err) + switch m := res.Message.(type) { + case *pb.Sync_Response_Insert: + record, err := pb.NewRecordFromBytes(m.Insert.Record) + require.NoError(t, err) + rows += record.NumRows() + case *pb.Sync_Response_Checkpoint: + record, err := pb.NewRecordFromBytes(m.Checkpoint.Record) + require.NoError(t, err) + values := map[string]string{} + addStateRows(values, record) + if cursor, ok := values["cursor"]; ok { + require.EqualValues(t, pages*checkpointPageSize, rows, "the cursor follows every row") + cursors = append(cursors, cursor) + } + } + } + require.Equal(t, []string{strconv.Itoa(pages - 1)}, cursors, "flushes of the cursor are merged into one checkpoint") +} + +func TestStateClientJoinsScopesOfSameValueWrites(t *testing.T) { + ctx := testContext(t) + stateClient, res, _ := newCheckpointStateClient(t, "test_same_value") + + scopeA, scopeB := &recordingScope{}, &recordingScope{} + require.NoError(t, stateClient.SetKey(checkpoint.WithScope(ctx, scopeA), "cursor", "1")) + // the same value from another resolver summarizes its resources too + require.NoError(t, stateClient.SetKey(checkpoint.WithScope(ctx, scopeB), "cursor", "1")) + require.NoError(t, stateClient.Flush(ctx)) + + require.Len(t, scopeA.msgs, 1) + require.Len(t, scopeB.msgs, 1) + deliver(t, res, scopeA.msgs...) + require.Empty(t, res) + deliver(t, res, scopeB.msgs...) + require.Equal(t, map[string]string{"cursor": "1"}, checkpointValues(t, receive(t, res))) +} + +func TestStateClientKeepsKeysWhoseCheckpointWasDropped(t *testing.T) { + ctx := testContext(t) + stateClient, res, _ := newCheckpointStateClient(t, "test_dropped") + + scope := &recordingScope{} + require.NoError(t, stateClient.SetKey(checkpoint.WithScope(ctx, scope), "cursor", "1")) + require.NoError(t, stateClient.Flush(ctx)) + require.Len(t, scope.msgs, 1) + + // the sync is cancelled before the checkpoint is sent + checkpoint.Drop(scope.msgs[0]) + require.ErrorIs(t, stateClient.Flush(ctx), state.ErrSyncAborted) + require.NoError(t, stateClient.SetKey(ctx, "cursor", "2")) + require.ErrorIs(t, stateClient.Flush(ctx), state.ErrSyncAborted) + require.Empty(t, res) +} + +// collapsingScope stands for a table resolver's scope, which collapses into the sync's once the resolver ends. +type collapsingScope struct { + recordingScope + into checkpoint.Scope +} + +func (s *collapsingScope) Collapsed() checkpoint.Scope { return s.into } + +func TestStateClientCollapsesEndedScopes(t *testing.T) { + ctx := testContext(t) + stateClient, res, _ := newCheckpointStateClient(t, "test_collapsed") + + shared := &recordingScope{} + for i := range 1000 { + scope := &collapsingScope{into: shared} + require.NoError(t, stateClient.SetKey(checkpoint.WithScope(ctx, scope), "cursor", strconv.Itoa(i))) + } + require.NoError(t, stateClient.Flush(ctx)) + require.Len(t, shared.msgs, 1, "one checkpoint through the shared scope") + deliver(t, res, shared.msgs...) + require.Equal(t, map[string]string{"cursor": "999"}, checkpointValues(t, receive(t, res))) +} + +// receive fails the test instead of hanging when nothing arrives. +func receive[T any](t *testing.T, ch <-chan T) T { + t.Helper() + select { + case v := <-ch: + return v + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for a value") + var zero T + return zero + } +} + +// stoppingClient's Write returns without waiting for its messages to end, after taking takeFirst messages. +type stoppingClient struct { + heldFlushClient + takeFirst int +} + +func (c *stoppingClient) Write(_ context.Context, msgs <-chan message.WriteMessage) error { + for range c.takeFirst { + <-msgs + } + return nil +} + +func TestWriteStreamWhenThePluginStopsWriting(t *testing.T) { + for _, tc := range []struct { + name string + takeFirst int + err string + }{ + {name: "before taking the flush", takeFirst: 0, err: "plugin stopped writing"}, + {name: "before completing the flush", takeFirst: 1, err: "plugin stopped writing before the flush completed"}, + } { + t.Run(tc.name, func(t *testing.T) { + ctx := testContext(t) + client := &stoppingClient{takeFirst: tc.takeFirst} + c := pb.NewPluginClient(servePlugin(t, plugin.NewPlugin("stopping", "v1.0.0", func(context.Context, zerolog.Logger, []byte, plugin.NewClientOptions) (plugin.Client, error) { + return client, nil + }, plugin.WithFlushSupport()))) + _, err := c.Init(ctx, &pb.Init_Request{}) + require.NoError(t, err) + + stream, err := c.WriteStream(ctx) + require.NoError(t, err) + require.NoError(t, stream.Send(&pb.Write_Request{Message: &pb.Write_Request_Flush{Flush: &pb.Write_MessageFlush{Id: 1}}})) + errs := make(chan error, 1) + go func() { + _, err := stream.Recv() + errs <- err + }() + err = receive(t, errs) + require.Equal(t, codes.Internal, status.Code(err)) + require.ErrorContains(t, err, tc.err) + }) + } +} + +func TestInitDoesNotReportFlushSupportWithoutConnection(t *testing.T) { + c := pb.NewPluginClient(servePlugin(t, plugin.NewPlugin("memdb", "v1.0.0", memdb.NewMemDBClient, plugin.WithFlushSupport()))) + res, err := c.Init(context.Background(), &pb.Init_Request{NoConnection: true}) + require.NoError(t, err) + require.False(t, res.SupportsFlush) +} + +func TestWriteStreamFlushErrorIsValidUTF8(t *testing.T) { + ctx := testContext(t) + client := &heldFlushClient{flushes: make(chan *message.WriteFlush)} + c := pb.NewPluginClient(servePlugin(t, plugin.NewPlugin("held", "v1.0.0", func(context.Context, zerolog.Logger, []byte, plugin.NewClientOptions) (plugin.Client, error) { + return client, nil + }, plugin.WithFlushSupport()))) + _, err := c.Init(ctx, &pb.Init_Request{}) + require.NoError(t, err) + stream, err := c.WriteStream(ctx) + require.NoError(t, err) + + for id, tc := range []struct{ err, want string }{ + {err: "bad \xff byte", want: "bad � byte"}, + {err: "", want: "flush failed"}, + } { + require.NoError(t, stream.Send(&pb.Write_Request{Message: &pb.Write_Request_Flush{Flush: &pb.Write_MessageFlush{Id: uint64(id)}}})) + receive(t, client.flushes).Done(errors.New(tc.err)) + res, err := stream.Recv() + require.NoError(t, err) + require.EqualValues(t, id, res.FlushId) + require.Equal(t, tc.want, res.FlushError) + } +} + +// failOnceClient fails the first flush, then returns from Write without waiting for its messages to end. +type failOnceClient struct { + heldFlushClient +} + +func (*failOnceClient) Write(_ context.Context, msgs <-chan message.WriteMessage) error { + for msg := range msgs { + if flush, ok := msg.(*message.WriteFlush); ok { + flush.Done(errors.New("disk full")) + return nil + } + } + return nil +} + +func TestWriteStreamReportsFailedFlushWhenWriteReturns(t *testing.T) { + // the server sees the failure and the end of Write at once, and picks either first + for i := range 20 { + // stop at the first failure, so that a regression fails the test instead of running past the test timeout + if !t.Run(strconv.Itoa(i), func(t *testing.T) { + ctx := testContext(t) + c := pb.NewPluginClient(servePlugin(t, plugin.NewPlugin("fail_once", "v1.0.0", func(context.Context, zerolog.Logger, []byte, plugin.NewClientOptions) (plugin.Client, error) { + return &failOnceClient{}, nil + }, plugin.WithFlushSupport()))) + _, err := c.Init(ctx, &pb.Init_Request{}) + require.NoError(t, err) + stream, err := c.WriteStream(ctx) + require.NoError(t, err) + + require.NoError(t, stream.Send(&pb.Write_Request{Message: &pb.Write_Request_Flush{Flush: &pb.Write_MessageFlush{Id: 1}}})) + res, err := stream.Recv() + require.NoError(t, err) + require.EqualValues(t, 1, res.FlushId) + require.Equal(t, "disk full", res.FlushError) + require.NoError(t, stream.CloseSend()) + _, err = stream.Recv() + require.ErrorIs(t, err, io.EOF) + }) { + break + } + } +} + +func TestStateClientSendsKeysOfTheSameScopesTogether(t *testing.T) { + ctx := testContext(t) + stateClient, res, _ := newCheckpointStateClient(t, "test_grouped") + + scopeA, scopeB := &recordingScope{}, &recordingScope{} + ctxA, ctxB := checkpoint.WithScope(ctx, scopeA), checkpoint.WithScope(ctx, scopeB) + require.NoError(t, stateClient.SetKey(ctxA, "a1", "1")) + require.NoError(t, stateClient.SetKey(ctxA, "a2", "1")) + require.NoError(t, stateClient.SetKey(ctxB, "b", "1")) + require.NoError(t, stateClient.SetKey(ctx, "c1", "1")) + require.NoError(t, stateClient.SetKey(ctx, "c2", "1")) + require.NoError(t, stateClient.Flush(ctx)) + + require.Len(t, res, 1) + require.Equal(t, map[string]string{"c1": "1", "c2": "1"}, checkpointValues(t, receive(t, res))) + require.Len(t, scopeA.msgs, 1) + require.Equal(t, map[string]string{"a1": "1", "a2": "1"}, checkpointValues(t, scopeA.msgs...)) + require.Len(t, scopeB.msgs, 1) + require.Equal(t, map[string]string{"b": "1"}, checkpointValues(t, scopeB.msgs...)) +} + +// testContext bounds the streams of a test, so that it fails instead of hanging. +func testContext(t *testing.T) context.Context { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), time.Minute) + t.Cleanup(cancel) + return ctx +} + +func TestStateClientCompactionKeepsEveryScope(t *testing.T) { + ctx := testContext(t) + stateClient, res, _ := newCheckpointStateClient(t, "test_compaction") + + // ended resolvers first, so that compactions collapse them while live resolvers keep adding scopes + shared := &recordingScope{} + for i := range 100 { + require.NoError(t, stateClient.SetKey(checkpoint.WithScope(ctx, &collapsingScope{into: shared}), "cursor", strconv.Itoa(i))) + } + live := make([]*recordingScope, 100) + for i := range live { + live[i] = &recordingScope{} + require.NoError(t, stateClient.SetKey(checkpoint.WithScope(ctx, live[i]), "cursor", strconv.Itoa(100+i))) + } + require.NoError(t, stateClient.Flush(ctx)) + + require.Len(t, shared.msgs, 1, "the ended resolvers' resources are waited for") + for i, scope := range live { + require.Len(t, scope.msgs, 1, "live scope %d is waited for", i) + } + for _, scope := range live { + deliver(t, res, scope.msgs...) + } + require.Empty(t, res) + deliver(t, res, shared.msgs...) + require.Equal(t, map[string]string{"cursor": "199"}, checkpointValues(t, receive(t, res))) +} + +func TestStateClientCheckpointWaitsForThePreviousCheckpointOfEveryKey(t *testing.T) { + // each order catches a checkpoint that waits for only one of the keys' previous checkpoints + for _, firstKey := range []string{"a", "b"} { + t.Run("first "+firstKey, func(t *testing.T) { + ctx := testContext(t) + stateClient, res, _ := newCheckpointStateClient(t, "test_group_previous") + previous := map[string]*recordingScope{"a": {}, "b": {}} + for _, key := range []string{"a", "b"} { + require.NoError(t, stateClient.SetKey(checkpoint.WithScope(ctx, previous[key]), key, "1")) + require.NoError(t, stateClient.Flush(ctx)) + } + both := &recordingScope{} + require.NoError(t, stateClient.SetKey(checkpoint.WithScope(ctx, both), "a", "2")) + require.NoError(t, stateClient.SetKey(checkpoint.WithScope(ctx, both), "b", "2")) + require.NoError(t, stateClient.Flush(ctx)) + require.Len(t, both.msgs, 1, "the keys of the same scope go in one checkpoint") + + secondKey := map[string]string{"a": "b", "b": "a"}[firstKey] + deliver(t, res, both.msgs...) + deliver(t, res, previous[firstKey].msgs...) + require.Equal(t, map[string]string{firstKey: "1"}, checkpointValues(t, receive(t, res))) + require.Empty(t, res, "the checkpoint of both keys waits for the previous checkpoint of %s", secondKey) + deliver(t, res, previous[secondKey].msgs...) + require.Equal(t, map[string]string{secondKey: "1"}, checkpointValues(t, receive(t, res))) + require.Equal(t, map[string]string{"a": "2", "b": "2"}, checkpointValues(t, receive(t, res))) + }) + } +} diff --git a/state/client.go b/state/client.go index cca64bbb41..5b5808bc58 100644 --- a/state/client.go +++ b/state/client.go @@ -6,7 +6,9 @@ import ( "slices" pbDiscovery "github.com/cloudquery/plugin-pb-go/pb/discovery/v1" + "github.com/cloudquery/plugin-sdk/v4/internal/checkpoint" stateV3 "github.com/cloudquery/plugin-sdk/v4/internal/clients/state/v3" + "github.com/cloudquery/plugin-sdk/v4/message" "github.com/cloudquery/plugin-sdk/v4/plugin" "google.golang.org/grpc" "google.golang.org/grpc/credentials/insecure" @@ -14,6 +16,10 @@ import ( const defaultMaxMsgSizeInBytes = 100 * 1024 * 1024 // 100 MiB +// ErrSyncAborted is returned by Flush when the sync was cancelled before the state could be sent as a checkpoint. +// The state stays pending. +var ErrSyncAborted = checkpoint.ErrSyncAborted + type Client interface { SetKey(ctx context.Context, key string, value string) error GetKey(ctx context.Context, key string) (string, error) @@ -21,7 +27,15 @@ type Client interface { Close() error } -type ClientOptions struct{} +type ClientOptions struct { + // Checkpoints is the res channel of the plugin's Sync, the same one passed to the scheduler. + // When the CLI supports checkpoints, Flush sends state on it instead of writing to the backend, and the CLI persists + // the state once the destinations commit the resources synced before it. Changes made within a table resolver (using + // its ctx) wait only for the resources that resolver produced; other changes must be flushed after the resources + // they summarize are sent, e.g. once the scheduler's Sync returns. + // Only NewConnectedClientWithOptions uses it, as only the backend options say whether the CLI supports checkpoints. + Checkpoints chan<- message.SyncMessage +} type ConnectionOptions struct { MaxMsgSizeInBytes int @@ -32,6 +46,11 @@ func NewClient(ctx context.Context, conn *grpc.ClientConn, tableName string) (Cl } func NewClientWithOptions(ctx context.Context, conn *grpc.ClientConn, tableName string, _ ClientOptions) (Client, error) { + return newStateClient(ctx, conn, tableName, nil) +} + +// newStateClient makes Flush send checkpoints on checkpoints, if set, instead of writing to the backend. +func newStateClient(ctx context.Context, conn *grpc.ClientConn, tableName string, checkpoints chan<- message.SyncMessage) (Client, error) { discoveryClient := pbDiscovery.NewDiscoveryClient(conn) versions, err := discoveryClient.GetVersions(ctx, &pbDiscovery.GetVersions_Request{}) if err != nil { @@ -41,7 +60,11 @@ func NewClientWithOptions(ctx context.Context, conn *grpc.ClientConn, tableName return nil, fmt.Errorf("please upgrade your state backend plugin. state supporting version 3 plugin has %v", versions.Versions) } - return stateV3.NewClient(ctx, conn, tableName) + var opts []stateV3.Option + if checkpoints != nil { + opts = append(opts, stateV3.WithCheckpoints(checkpoints)) + } + return stateV3.NewClient(ctx, conn, tableName, opts...) } // NewConnectedClient returns a state client and initialises the gRPC connection to the state backend with a 100MiB max message size. @@ -65,6 +88,14 @@ func NewConnectedClientWithOptions(ctx context.Context, backendOpts *plugin.Back return nil, fmt.Errorf("backend_options must contain both connection and table_name: %v", backendOpts) } + var checkpoints chan<- message.SyncMessage + if backendOpts.SupportsCheckpoints { + checkpoints = clOpts.Checkpoints + } + if connOpts.MaxMsgSizeInBytes <= 0 { + connOpts.MaxMsgSizeInBytes = defaultMaxMsgSizeInBytes + } + backendConn, err := grpc.NewClient(backendOpts.Connection, grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithDefaultCallOptions( @@ -77,7 +108,7 @@ func NewConnectedClientWithOptions(ctx context.Context, backendOpts *plugin.Back return &NoOpClient{}, fmt.Errorf("failed to dial grpc source plugin at %s: %w", backendOpts.Connection, err) } - stateClient, err := NewClientWithOptions(ctx, backendConn, backendOpts.TableName, clOpts) + stateClient, err := newStateClient(ctx, backendConn, backendOpts.TableName, checkpoints) if err != nil { return &NoOpClient{}, fmt.Errorf("failed to create state client: %w", err) } diff --git a/writers/batchwriter/batchwriter.go b/writers/batchwriter/batchwriter.go index 99ab50d1c6..5a6c0946ea 100644 --- a/writers/batchwriter/batchwriter.go +++ b/writers/batchwriter/batchwriter.go @@ -3,6 +3,9 @@ package batchwriter import ( "context" "errors" + "fmt" + "maps" + "slices" "sync" "time" @@ -33,12 +36,21 @@ type BatchWriter struct { deleteRecordLock sync.Mutex deleteRecordMessages message.WriteDeleteRecords + // batches that failed to write, by table, reported on message.WriteFlush + failuresLock sync.Mutex + failures map[string]tableFailures + logger zerolog.Logger batchTimeout time.Duration batchSize int64 batchSizeBytes int64 } +type tableFailures struct { + count uint64 + last error +} + // Assert at compile-time that BatchWriter implements the Writer interface var _ writers.Writer = (*BatchWriter)(nil) @@ -83,6 +95,7 @@ func New(client Client, opts ...Option) (*BatchWriter, error) { c := &BatchWriter{ client: client, workers: make(map[string]*worker), + failures: make(map[string]tableFailures), logger: zerolog.Nop(), batchTimeout: defaultBatchTimeoutSeconds * time.Second, batchSize: defaultBatchSize, @@ -104,15 +117,38 @@ func (w *BatchWriter) Flush(ctx context.Context) error { <-done } w.workersLock.RUnlock() + // the deletes wait for every pending migration, so a failure leaves them unwritten too if err := w.flushMigrateTables(ctx); err != nil { - return err + return errors.Join(err, w.pendingDeleteStales(), w.pendingDeleteRecords()) } if err := w.flushDeleteStaleTables(ctx); err != nil { - return err + return errors.Join(err, w.pendingDeleteRecords()) } return w.flushDeleteRecordTables(ctx) } +var errNotWritten = errors.New("not written, as an earlier write failed") + +// pendingDeleteStales returns an error naming the tables of the delete-stale messages not yet written. +func (w *BatchWriter) pendingDeleteStales() error { + w.deleteStaleLock.Lock() + defer w.deleteStaleLock.Unlock() + if len(w.deleteStaleMessages) == 0 { + return nil + } + return writers.FailedTables(errNotWritten, w.deleteStaleMessages) +} + +// pendingDeleteRecords returns an error naming the tables of the delete-record messages not yet written. +func (w *BatchWriter) pendingDeleteRecords() error { + w.deleteRecordLock.Lock() + defer w.deleteRecordLock.Unlock() + if len(w.deleteRecordMessages) == 0 { + return nil + } + return writers.FailedTables(errNotWritten, w.deleteRecordMessages) +} + func (w *BatchWriter) Close(context.Context) error { w.workersLock.Lock() defer w.workersLock.Unlock() @@ -203,6 +239,9 @@ func (w *BatchWriter) flushTable(ctx context.Context, tableName string, resource duration := time.Since(start) if err != nil { w.logger.Err(err).Str("table", tableName).Int64("len", batchSize).Dur("duration", duration).Msg("failed to write batch") + w.failuresLock.Lock() + w.failures[tableName] = tableFailures{count: w.failures[tableName].count + 1, last: err} + w.failuresLock.Unlock() } else { w.logger.Debug().Str("table", tableName).Int64("len", batchSize).Dur("duration", duration).Msg("batch written successfully") } @@ -215,7 +254,7 @@ func (w *BatchWriter) flushMigrateTables(ctx context.Context) error { return nil } if err := w.client.MigrateTables(ctx, w.migrateTableMessages); err != nil { - return err + return writers.FailedTables(err, w.migrateTableMessages) } w.migrateTableMessages = w.migrateTableMessages[:0] return nil @@ -228,7 +267,7 @@ func (w *BatchWriter) flushDeleteStaleTables(ctx context.Context) error { return nil } if err := w.client.DeleteStale(ctx, w.deleteStaleMessages); err != nil { - return err + return writers.FailedTables(err, w.deleteStaleMessages) } w.deleteStaleMessages = w.deleteStaleMessages[:0] return nil @@ -241,7 +280,7 @@ func (w *BatchWriter) flushDeleteRecordTables(ctx context.Context) error { return nil } if err := w.client.DeleteRecord(ctx, w.deleteRecordMessages); err != nil { - return err + return writers.FailedTables(err, w.deleteRecordMessages) } w.deleteRecordMessages = w.deleteRecordMessages[:0] return nil @@ -261,6 +300,38 @@ func (w *BatchWriter) flushInsert(tableName string) { <-ch } +// trackFailures adds the table of msg to reported, so that the next flush reports the failures of that table from now +// on. +func (w *BatchWriter) trackFailures(reported map[string]uint64, msg *message.WriteInsert) { + tableName, _ := msg.Record.Schema().Metadata().GetValue(schema.MetadataTableName) + if _, ok := reported[tableName]; ok { + return + } + w.failuresLock.Lock() + defer w.failuresLock.Unlock() + reported[tableName] = w.failures[tableName].count +} + +// newFailures returns the batch writes of the tables in reported that failed since the previous flush, and marks them +// as reported. +func (w *BatchWriter) newFailures(reported map[string]uint64) error { + w.failuresLock.Lock() + defer w.failuresLock.Unlock() + var errs []error + for _, tableName := range slices.Sorted(maps.Keys(reported)) { + failures := w.failures[tableName] + if failures.count == reported[tableName] { + continue + } + errs = append(errs, &message.FailedTablesError{ + Tables: []string{tableName}, + Err: fmt.Errorf("table %s: %d batch writes failed since the previous flush, last error: %w", tableName, failures.count-reported[tableName], failures.last), + }) + reported[tableName] = failures.count + } + return errors.Join(errs...) +} + func (w *BatchWriter) writeAll(ctx context.Context, msgs []message.WriteMessage) error { ch := make(chan message.WriteMessage, len(msgs)) for _, msg := range msgs { @@ -271,8 +342,13 @@ func (w *BatchWriter) writeAll(ctx context.Context, msgs []message.WriteMessage) } func (w *BatchWriter) Write(ctx context.Context, msgs <-chan message.WriteMessage) error { + // failed batch writes reported so far, for each table this Write inserts to. Several Write calls can share the + // writer, as the state backend's does with the sync's, and each reports only the failures of its own tables. + reported := make(map[string]uint64) for msg := range msgs { switch m := msg.(type) { + case *message.WriteFlush: + m.Done(errors.Join(w.Flush(ctx), w.newFailures(reported))) case *message.WriteDeleteStale: if err := w.flushMigrateTables(ctx); err != nil { return err @@ -308,6 +384,8 @@ func (w *BatchWriter) Write(ctx context.Context, msgs <-chan message.WriteMessag } } case *message.WriteInsert: + // before the worker gets the insert, so that a failure of its batch counts against this Write + w.trackFailures(reported, m) if err := w.flushMigrateTables(ctx); err != nil { return err } diff --git a/writers/batchwriter/flush_test.go b/writers/batchwriter/flush_test.go new file mode 100644 index 0000000000..913fca2dd9 --- /dev/null +++ b/writers/batchwriter/flush_test.go @@ -0,0 +1,347 @@ +package batchwriter + +import ( + "context" + "errors" + "sync" + "testing" + "time" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/cloudquery/plugin-sdk/v4/message" + "github.com/cloudquery/plugin-sdk/v4/schema" + "github.com/stretchr/testify/require" +) + +type failingBatchClient struct { + testBatchClient +} + +func (*failingBatchClient) WriteTableBatch(context.Context, string, message.WriteInserts) error { + return errors.New("write failed") +} + +func startWrite(t *testing.T, wr *BatchWriter) (chan<- message.WriteMessage, func() error) { + t.Helper() + ch := make(chan message.WriteMessage) + errCh := make(chan error, 1) + go func() { + errCh <- wr.Write(context.Background(), ch) + }() + return ch, func() error { + close(ch) + return receive(t, errCh) + } +} + +func flush(t *testing.T, ch chan<- message.WriteMessage) error { + t.Helper() + msg := message.NewWriteFlush() + select { + case ch <- msg: + case <-time.After(10 * time.Second): + t.Fatal("Write did not take the flush") + } + return receive(t, msg.Result()) +} + +func TestFlushMessageCommitsBatches(t *testing.T) { + testClient := &testBatchClient{} + wr, err := New(testClient, WithBatchTimeout(time.Hour)) + require.NoError(t, err) + ch, stop := startWrite(t, wr) + + ch <- &message.WriteMigrateTable{Table: batchTestTables[0]} + ch <- &message.WriteInsert{Record: getRecord(batchTestTables[0].ToArrowSchema(), 3)} + require.NoError(t, flush(t, ch)) + require.Equal(t, 1, testClient.MigrateTablesLen()) + require.EqualValues(t, 3, testClient.InsertedRows()) + + ch <- &message.WriteInsert{Record: getRecord(batchTestTables[0].ToArrowSchema(), 2)} + require.NoError(t, flush(t, ch)) + require.EqualValues(t, 5, testClient.InsertedRows()) + + require.NoError(t, stop()) + require.NoError(t, wr.Close(context.Background())) +} + +func TestFlushMessageReportsFailedInserts(t *testing.T) { + wr, err := New(&failingBatchClient{}, WithBatchTimeout(time.Hour)) + require.NoError(t, err) + ch, stop := startWrite(t, wr) + + ch <- &message.WriteMigrateTable{Table: batchTestTables[0]} + ch <- &message.WriteInsert{Record: getRecord(batchTestTables[0].ToArrowSchema(), 3)} + require.ErrorContains(t, flush(t, ch), "write failed") + // the failure was reported already + require.NoError(t, flush(t, ch)) + + require.NoError(t, stop()) + require.NoError(t, wr.Close(context.Background())) +} + +func TestFlushMessageReportsFailuresFromTimeouts(t *testing.T) { + wr, err := New(&failingBatchClient{}, WithBatchTimeout(time.Millisecond)) + require.NoError(t, err) + ch, stop := startWrite(t, wr) + + ch <- &message.WriteMigrateTable{Table: batchTestTables[0]} + ch <- &message.WriteInsert{Record: getRecord(batchTestTables[0].ToArrowSchema(), 3)} + require.Eventually(t, func() bool { + wr.failuresLock.Lock() + defer wr.failuresLock.Unlock() + return wr.failures[batchTestTables[0].Name].count > 0 + }, 5*time.Second, time.Millisecond) + require.ErrorContains(t, flush(t, ch), "write failed") + + require.NoError(t, stop()) + require.NoError(t, wr.Close(context.Background())) +} + +func TestFlushMessageCommitsDeletes(t *testing.T) { + testClient := &testBatchClient{} + wr, err := New(testClient, WithBatchTimeout(time.Hour)) + require.NoError(t, err) + ch, stop := startWrite(t, wr) + + ch <- &message.WriteMigrateTable{Table: batchTestTables[0]} + ch <- &message.WriteDeleteStale{TableName: batchTestTables[0].Name, SourceName: "source", SyncTime: time.Now()} + require.NoError(t, flush(t, ch)) + require.Equal(t, 1, testClient.DeleteStalesLen()) + + ch <- &message.WriteDeleteRecord{DeleteRecord: message.DeleteRecord{TableName: batchTestTables[0].Name}} + require.NoError(t, flush(t, ch)) + testClient.mutex.Lock() + require.Len(t, testClient.deleteRecords, 1) + testClient.mutex.Unlock() + + require.NoError(t, stop()) + require.NoError(t, wr.Close(context.Background())) +} + +func TestFlushMessageErrorNamesTable(t *testing.T) { + wr, err := New(&failingBatchClient{}, WithBatchTimeout(time.Hour)) + require.NoError(t, err) + ch, stop := startWrite(t, wr) + + ch <- &message.WriteMigrateTable{Table: batchTestTables[0]} + ch <- &message.WriteInsert{Record: getRecord(batchTestTables[0].ToArrowSchema(), 1)} + err = flush(t, ch) + require.ErrorContains(t, err, "table "+batchTestTables[0].Name+": 1 batch writes failed since the previous flush, last error: write failed") + require.Equal(t, []string{batchTestTables[0].Name}, message.FailedTables(err)) + + require.NoError(t, stop()) + require.NoError(t, wr.Close(context.Background())) +} + +type failingOpsClient struct { + testBatchClient + migrateErr, deleteStaleErr, deleteRecordErr error +} + +func (c *failingOpsClient) MigrateTables(ctx context.Context, msgs message.WriteMigrateTables) error { + if c.migrateErr != nil { + return c.migrateErr + } + return c.testBatchClient.MigrateTables(ctx, msgs) +} + +func (c *failingOpsClient) DeleteStale(ctx context.Context, msgs message.WriteDeleteStales) error { + if c.deleteStaleErr != nil { + return c.deleteStaleErr + } + return c.testBatchClient.DeleteStale(ctx, msgs) +} + +func (c *failingOpsClient) DeleteRecord(ctx context.Context, msgs message.WriteDeleteRecords) error { + if c.deleteRecordErr != nil { + return c.deleteRecordErr + } + return c.testBatchClient.DeleteRecord(ctx, msgs) +} + +func TestFlushMessageReportsFailedMigrationsAndDeletes(t *testing.T) { + table := batchTestTables[0] + for _, tc := range []struct { + name string + client *failingOpsClient + msgs []message.WriteMessage + }{ + { + name: "migrate", + client: &failingOpsClient{migrateErr: errors.New("migrate failed")}, + msgs: []message.WriteMessage{&message.WriteMigrateTable{Table: table}}, + }, + { + name: "delete stale", + client: &failingOpsClient{deleteStaleErr: errors.New("delete stale failed")}, + msgs: []message.WriteMessage{ + &message.WriteMigrateTable{Table: table}, + &message.WriteDeleteStale{TableName: table.Name, SourceName: "source", SyncTime: time.Now()}, + }, + }, + { + name: "delete record", + client: &failingOpsClient{deleteRecordErr: errors.New("delete record failed")}, + msgs: []message.WriteMessage{ + &message.WriteMigrateTable{Table: table}, + &message.WriteDeleteRecord{DeleteRecord: message.DeleteRecord{TableName: table.Name}}, + }, + }, + } { + t.Run(tc.name, func(t *testing.T) { + wr, err := New(tc.client, WithBatchTimeout(time.Hour)) + require.NoError(t, err) + ch, stop := startWrite(t, wr) + for _, msg := range tc.msgs { + ch <- msg + } + err = flush(t, ch) + require.ErrorContains(t, err, tc.name+" failed") + require.Equal(t, []string{table.Name}, message.FailedTables(err)) + require.NoError(t, stop()) + }) + } +} + +// failingTableClient fails to write the batches of one table. +type failingTableClient struct { + testBatchClient + table string +} + +func (c *failingTableClient) WriteTableBatch(ctx context.Context, name string, msgs message.WriteInserts) error { + if name == c.table { + return errors.New("write failed") + } + return c.testBatchClient.WriteTableBatch(ctx, name, msgs) +} + +func TestFlushMessageReportsOnlyTheFailuresOfItsOwnTables(t *testing.T) { + table := func(name string) *schema.Table { + return &schema.Table{Name: name, Columns: schema.ColumnList{{Name: "id", Type: arrow.PrimitiveTypes.Int64}}} + } + failing, healthy := table("failing"), table("healthy") + wr, err := New(&failingTableClient{table: failing.Name}, WithBatchTimeout(time.Hour)) + require.NoError(t, err) + // two streams share the writer, as the state backend's does with the sync's when both use one destination + failingCh, stopFailing := startWrite(t, wr) + healthyCh, stopHealthy := startWrite(t, wr) + + failingCh <- &message.WriteMigrateTable{Table: failing} + healthyCh <- &message.WriteMigrateTable{Table: healthy} + failingCh <- &message.WriteInsert{Record: getRecord(failing.ToArrowSchema(), 1)} + healthyCh <- &message.WriteInsert{Record: getRecord(healthy.ToArrowSchema(), 1)} + + err = flush(t, failingCh) + require.ErrorContains(t, err, "write failed") + require.Equal(t, []string{failing.Name}, message.FailedTables(err)) + require.NoError(t, flush(t, healthyCh), "the failure of another stream's table is not reported") + + // a stream that starts writing a table later does not report its earlier failures + healthyCh <- &message.WriteInsert{Record: getRecord(failing.ToArrowSchema(), 1)} + err = flush(t, healthyCh) + require.Equal(t, []string{failing.Name}, message.FailedTables(err), "only the failure of its own insert") + require.ErrorContains(t, err, "1 batch writes failed") + + require.NoError(t, stopFailing()) + require.NoError(t, stopHealthy()) +} + +// receive fails the test instead of hanging when nothing arrives. +func receive[T any](t *testing.T, ch <-chan T) T { + t.Helper() + select { + case v := <-ch: + return v + case <-time.After(10 * time.Second): + t.Fatal("timed out waiting for a value") + var zero T + return zero + } +} + +// failOnceClient fails the first batch of each table. +type failOnceClient struct { + testBatchClient + mu sync.Mutex + failed map[string]bool +} + +func (c *failOnceClient) WriteTableBatch(ctx context.Context, name string, msgs message.WriteInserts) error { + c.mu.Lock() + first := !c.failed[name] + c.failed[name] = true + c.mu.Unlock() + if first { + return errors.New("write failed") + } + return c.testBatchClient.WriteTableBatch(ctx, name, msgs) +} + +func TestFlushMessageNamesEveryFailedTable(t *testing.T) { + table := func(name string) *schema.Table { + return &schema.Table{Name: name, Columns: schema.ColumnList{{Name: "id", Type: arrow.PrimitiveTypes.Int64}}} + } + first, second := table("first"), table("second") + wr, err := New(&failOnceClient{failed: map[string]bool{}}, WithBatchTimeout(time.Millisecond)) + require.NoError(t, err) + ch, stop := startWrite(t, wr) + + for _, table := range []*schema.Table{second, first} { + ch <- &message.WriteMigrateTable{Table: table} + ch <- &message.WriteInsert{Record: getRecord(table.ToArrowSchema(), 1)} + } + failedTables := func() int { + wr.failuresLock.Lock() + defer wr.failuresLock.Unlock() + return len(wr.failures) + } + require.Eventually(t, func() bool { return failedTables() == 2 }, 5*time.Second, time.Millisecond) + // later inserts of the tables succeed, and must not hide the failures before them + for _, table := range []*schema.Table{first, second} { + ch <- &message.WriteInsert{Record: getRecord(table.ToArrowSchema(), 1)} + } + err = flush(t, ch) + require.Equal(t, []string{"first", "second"}, message.FailedTables(err)) + require.NoError(t, flush(t, ch), "the failures were reported") + + require.NoError(t, stop()) +} + +func TestFlushMessageNamesDeletesLeftUnwrittenByAnEarlierFailure(t *testing.T) { + table := batchTestTables[0] + deleteRecord := &message.WriteDeleteRecord{DeleteRecord: message.DeleteRecord{TableName: "deleted"}} + for _, tc := range []struct { + name string + client *failingOpsClient + msgs []message.WriteMessage + want []string + }{ + { + name: "migrate", + client: &failingOpsClient{migrateErr: errors.New("migrate failed")}, + msgs: []message.WriteMessage{deleteRecord, &message.WriteMigrateTable{Table: table}}, + want: []string{"deleted", table.Name}, + }, + { + name: "delete stale", + client: &failingOpsClient{deleteStaleErr: errors.New("delete stale failed")}, + msgs: []message.WriteMessage{deleteRecord, &message.WriteDeleteStale{TableName: "stale", SourceName: "source", SyncTime: time.Now()}}, + want: []string{"deleted", "stale"}, + }, + } { + t.Run(tc.name, func(t *testing.T) { + wr, err := New(tc.client, WithBatchTimeout(time.Hour)) + require.NoError(t, err) + ch, stop := startWrite(t, wr) + for _, msg := range tc.msgs { + ch <- msg + } + err = flush(t, ch) + require.ErrorContains(t, err, tc.name+" failed") + require.Equal(t, tc.want, message.FailedTables(err)) + require.NoError(t, stop()) + }) + } +} diff --git a/writers/mixedbatchwriter/flush_test.go b/writers/mixedbatchwriter/flush_test.go new file mode 100644 index 0000000000..2b2bb0db49 --- /dev/null +++ b/writers/mixedbatchwriter/flush_test.go @@ -0,0 +1,158 @@ +package mixedbatchwriter + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/cloudquery/plugin-sdk/v4/message" + "github.com/cloudquery/plugin-sdk/v4/schema" + "github.com/stretchr/testify/require" +) + +type failingMixedBatchClient struct { + testMixedBatchClient +} + +func (*failingMixedBatchClient) InsertBatch(context.Context, message.WriteInserts) error { + return errors.New("insert failed") +} + +func TestFlushMessageCommitsBatches(t *testing.T) { + msgs := getTestMessages() + client := &testMixedBatchClient{} + wr, err := New(client, WithBatchTimeout(time.Hour)) + require.NoError(t, err) + + ch := make(chan message.WriteMessage) + errCh := make(chan error, 1) + go func() { + errCh <- wr.Write(context.Background(), ch) + }() + + ch <- msgs.migrateTable1 + ch <- msgs.insert1 + flush := message.NewWriteFlush() + ch <- flush + require.NoError(t, receive(t, flush.Result())) + require.Equal(t, [][]message.WriteMessage{{msgs.migrateTable1}, {msgs.insert1}}, client.receivedBatches) + + ch <- msgs.insert2 + flush = message.NewWriteFlush() + ch <- flush + require.NoError(t, receive(t, flush.Result())) + require.Equal(t, [][]message.WriteMessage{{msgs.migrateTable1}, {msgs.insert1}, {msgs.insert2}}, client.receivedBatches) + + close(ch) + require.NoError(t, receive(t, errCh)) +} + +func TestFlushMessageReportsFailedInserts(t *testing.T) { + msgs := getTestMessages() + wr, err := New(&failingMixedBatchClient{}, WithBatchTimeout(time.Hour)) + require.NoError(t, err) + + ch := make(chan message.WriteMessage, 3) + flush := message.NewWriteFlush() + ch <- msgs.migrateTable1 + ch <- msgs.insert1 + ch <- flush + close(ch) + + require.ErrorContains(t, wr.Write(context.Background(), ch), "insert failed") + err = receive(t, flush.Result()) + require.ErrorContains(t, err, "insert failed") + require.Equal(t, []string{msgs.migrateTable1.Table.Name}, message.FailedTables(err)) +} + +func TestFlushMessageCommitsDeletes(t *testing.T) { + msgs := getTestMessages() + client := &testMixedBatchClient{} + wr, err := New(client, WithBatchTimeout(time.Hour)) + require.NoError(t, err) + + ch := make(chan message.WriteMessage) + errCh := make(chan error, 1) + go func() { + errCh <- wr.Write(context.Background(), ch) + }() + + ch <- msgs.deleteStale1 + flush := message.NewWriteFlush() + ch <- flush + require.NoError(t, receive(t, flush.Result())) + require.Equal(t, [][]message.WriteMessage{{msgs.deleteStale1}}, client.receivedBatches) + + close(ch) + require.NoError(t, receive(t, errCh)) +} + +func TestFlushMessageCommitsMigrationsAndDeleteRecords(t *testing.T) { + msgs := getTestMessages() + deleteRecord := &message.WriteDeleteRecord{DeleteRecord: message.DeleteRecord{TableName: msgs.migrateTable1.Table.Name}} + client := &testMixedBatchClient{} + wr, err := New(client, WithBatchTimeout(time.Hour)) + require.NoError(t, err) + + ch := make(chan message.WriteMessage) + errCh := make(chan error, 1) + go func() { + errCh <- wr.Write(context.Background(), ch) + }() + + ch <- msgs.migrateTable1 + flush := message.NewWriteFlush() + ch <- flush + require.NoError(t, receive(t, flush.Result())) + require.Equal(t, [][]message.WriteMessage{{msgs.migrateTable1}}, client.receivedBatches) + + ch <- deleteRecord + flush = message.NewWriteFlush() + ch <- flush + require.NoError(t, receive(t, flush.Result())) + require.Equal(t, [][]message.WriteMessage{{msgs.migrateTable1}, {deleteRecord}}, client.receivedBatches) + + close(ch) + require.NoError(t, receive(t, errCh)) +} + +// receive fails the test instead of hanging when nothing arrives. +func receive[T any](t *testing.T, ch <-chan T) T { + t.Helper() + select { + case v := <-ch: + return v + case <-time.After(10 * time.Second): + t.Fatal("timed out waiting for a value") + var zero T + return zero + } +} + +func TestFlushMessageNamesEveryTableOfTheFailedBatch(t *testing.T) { + table := func(name string) *schema.Table { + return &schema.Table{Name: name, Columns: schema.ColumnList{{Name: "id", Type: arrow.PrimitiveTypes.Int64}}} + } + insert := func(table *schema.Table) *message.WriteInsert { + bldr := array.NewRecordBuilder(memory.DefaultAllocator, table.ToArrowSchema()) + defer bldr.Release() + bldr.Field(0).(*array.Int64Builder).Append(1) + return &message.WriteInsert{Record: bldr.NewRecordBatch()} + } + wr, err := New(&failingMixedBatchClient{}, WithBatchTimeout(time.Hour)) + require.NoError(t, err) + + ch := make(chan message.WriteMessage, 3) + flush := message.NewWriteFlush() + ch <- insert(table("table2")) + ch <- insert(table("table1")) + ch <- flush + close(ch) + + require.ErrorContains(t, wr.Write(context.Background(), ch), "insert failed") + require.Equal(t, []string{"table1", "table2"}, message.FailedTables(receive(t, flush.Result()))) +} diff --git a/writers/mixedbatchwriter/mixedbatchwriter.go b/writers/mixedbatchwriter/mixedbatchwriter.go index 38033c05b4..f4b5ea48cf 100644 --- a/writers/mixedbatchwriter/mixedbatchwriter.go +++ b/writers/mixedbatchwriter/mixedbatchwriter.go @@ -133,6 +133,16 @@ loop: if !ok { break loop } + if flushMsg, ok := msg.(*message.WriteFlush); ok { + // batches hold a single message type, so only the previous type can have pending messages + err := flush(prevMsgType) + flushMsg.Done(err) + if err != nil { + return err + } + prevMsgType = writers.MsgTypeUnset + continue + } msgType := writers.MsgID(msg) if prevMsgType != msgType { if err := flush(prevMsgType); err != nil { @@ -189,7 +199,7 @@ func (m *batchManager[A, T]) flush(ctx context.Context) error { err := m.writeFunc(ctx, m.batch) if err != nil { - return err + return writers.FailedTables(err, m.batch) } clear(m.batch) // GC can work m.batch = m.batch[:0] @@ -244,7 +254,7 @@ func (m *insertBatchManager) flush(ctx context.Context) error { duration := time.Since(start) if err != nil { m.logger.Err(err).Int64("len", rows).Dur("duration", duration).Msg("failed to write batch") - return err + return writers.FailedTables(err, m.batch) } m.logger.Debug().Int64("len", rows).Dur("duration", duration).Msg("batch written successfully") diff --git a/writers/streamingbatchwriter/flush_test.go b/writers/streamingbatchwriter/flush_test.go new file mode 100644 index 0000000000..cdca3430cf --- /dev/null +++ b/writers/streamingbatchwriter/flush_test.go @@ -0,0 +1,203 @@ +package streamingbatchwriter + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/cloudquery/plugin-sdk/v4/message" + "github.com/cloudquery/plugin-sdk/v4/schema" + "github.com/stretchr/testify/require" +) + +func TestFlushMessageCommitsBatches(t *testing.T) { + testClient := newClient() + wr, err := New(testClient, WithBatchTimeout(time.Hour)) + require.NoError(t, err) + + ch := make(chan message.WriteMessage) + errCh := make(chan error, 1) + go func() { + errCh <- wr.Write(context.Background(), ch) + }() + + ch <- &message.WriteMigrateTable{Table: streamingBatchTestTable} + ch <- &message.WriteInsert{Record: getRecord(streamingBatchTestTable.ToArrowSchema(), 1)} + ch <- &message.WriteInsert{Record: getRecord(streamingBatchTestTable.ToArrowSchema(), 1)} + flush := message.NewWriteFlush() + ch <- flush + require.NoError(t, receive(t, flush.Result())) + require.Equal(t, 1, testClient.MessageLen(messageTypeMigrateTable)) + require.Equal(t, 2, testClient.MessageLen(messageTypeInsert)) + require.Equal(t, 0, testClient.OpenLen(messageTypeInsert)) + + ch <- &message.WriteInsert{Record: getRecord(streamingBatchTestTable.ToArrowSchema(), 1)} + flush = message.NewWriteFlush() + ch <- flush + require.NoError(t, receive(t, flush.Result())) + require.Equal(t, 3, testClient.MessageLen(messageTypeInsert)) + + close(ch) + require.NoError(t, receive(t, errCh)) +} + +func TestFlushMessageReportsFailedWrites(t *testing.T) { + testClient := newClient() + testClient.writeCommitErr = errors.New("commit failed") + wr, err := New(testClient, WithBatchTimeout(time.Hour)) + require.NoError(t, err) + + ch := make(chan message.WriteMessage) + errCh := make(chan error, 1) + go func() { + errCh <- wr.Write(context.Background(), ch) + }() + + ch <- &message.WriteInsert{Record: getRecord(streamingBatchTestTable.ToArrowSchema(), 1)} + flush := message.NewWriteFlush() + ch <- flush + err = receive(t, flush.Result()) + require.ErrorContains(t, err, "failed to write messages for table1") + require.Equal(t, []string{"table1"}, message.FailedTables(err)) + + close(ch) + require.ErrorContains(t, receive(t, errCh), "commit failed") +} + +func TestFlushMessageReportsRejectedMessages(t *testing.T) { + testClient := newClient() + wr, err := New(testClient, WithBatchTimeout(time.Hour)) + require.NoError(t, err) + + ch := make(chan message.WriteMessage) + errCh := make(chan error, 1) + go func() { + errCh <- wr.Write(context.Background(), ch) + }() + + noTableName := arrow.NewSchema([]arrow.Field{{Name: "id", Type: arrow.PrimitiveTypes.Int64}}, nil) + ch <- &message.WriteInsert{Record: getRecord(noTableName, 1)} + flush := message.NewWriteFlush() + ch <- flush + err = receive(t, flush.Result()) + require.ErrorContains(t, err, "table name not found in metadata") + require.Nil(t, message.FailedTables(err), "a message without a table may belong to any table") + + close(ch) + require.ErrorContains(t, receive(t, errCh), "table name not found in metadata") +} + +// blockingCommitClient holds each WriteTable call open after its messages, until release is closed. +type blockingCommitClient struct { + *testStreamingBatchClient + committing chan struct{} + release chan struct{} +} + +func (c *blockingCommitClient) WriteTable(_ context.Context, msgs <-chan *message.WriteInsert) error { + // nolint:revive + for range msgs { + } + c.committing <- struct{}{} + <-c.release + return nil +} + +func TestFlushMessageCancelledDuringCommit(t *testing.T) { + testClient := &blockingCommitClient{testStreamingBatchClient: newClient(), committing: make(chan struct{}, 1), release: make(chan struct{})} + wr, err := New(testClient, WithBatchTimeout(time.Hour)) + require.NoError(t, err) + + ctx, cancel := context.WithCancel(context.Background()) + ch := make(chan message.WriteMessage) + errCh := make(chan error, 1) + go func() { + errCh <- wr.Write(ctx, ch) + }() + + ch <- &message.WriteInsert{Record: getRecord(streamingBatchTestTable.ToArrowSchema(), 1)} + flush := message.NewWriteFlush() + ch <- flush + receive(t, testClient.committing) + cancel() + require.ErrorIs(t, receive(t, flush.Result()), context.Canceled) + close(testClient.release) + close(ch) + + select { + case <-errCh: + case <-time.After(5 * time.Second): + t.Fatal("Write did not return after the flush was cancelled") + } +} + +func TestFlushMessageCommitsDeletes(t *testing.T) { + testClient := newClient() + wr, err := New(testClient, WithBatchTimeout(time.Hour)) + require.NoError(t, err) + + ch := make(chan message.WriteMessage) + errCh := make(chan error, 1) + go func() { + errCh <- wr.Write(context.Background(), ch) + }() + + ch <- &message.WriteDeleteStale{TableName: streamingBatchTestTable.Name, SourceName: "source", SyncTime: time.Now()} + flush := message.NewWriteFlush() + ch <- flush + require.NoError(t, receive(t, flush.Result())) + require.Equal(t, 1, testClient.MessageLen(messageTypeDeleteStale)) + + ch <- &message.WriteDeleteRecord{DeleteRecord: message.DeleteRecord{TableName: streamingBatchTestTable.Name}} + flush = message.NewWriteFlush() + ch <- flush + require.NoError(t, receive(t, flush.Result())) + require.Equal(t, 1, testClient.MessageLen(messageTypeDeleteRecord)) + + close(ch) + require.NoError(t, receive(t, errCh)) +} + +// receive fails the test instead of hanging when nothing arrives. +func receive[T any](t *testing.T, ch <-chan T) T { + t.Helper() + select { + case v := <-ch: + return v + case <-time.After(10 * time.Second): + t.Fatal("timed out waiting for a value") + var zero T + return zero + } +} + +func TestFlushMessageNamesEveryFailedTable(t *testing.T) { + testClient := newClient() + testClient.writeCommitErr = errors.New("commit failed") + wr, err := New(testClient, WithBatchTimeout(time.Hour)) + require.NoError(t, err) + + ch := make(chan message.WriteMessage) + errCh := make(chan error, 1) + go func() { + errCh <- wr.Write(context.Background(), ch) + }() + + other := &schema.Table{Name: "table2", Columns: streamingBatchTestTable.Columns} + ch <- &message.WriteInsert{Record: getRecord(streamingBatchTestTable.ToArrowSchema(), 1)} + ch <- &message.WriteInsert{Record: getRecord(other.ToArrowSchema(), 1)} + flush := message.NewWriteFlush() + ch <- flush + err = receive(t, flush.Result()) + require.Equal(t, []string{"table1", "table2"}, message.FailedTables(err)) + + // a failed handler keeps failing its flushes, as its earlier messages are lost + flush = message.NewWriteFlush() + ch <- flush + require.Equal(t, []string{"table1", "table2"}, message.FailedTables(receive(t, flush.Result()))) + + close(ch) + require.ErrorContains(t, receive(t, errCh), "commit failed") +} diff --git a/writers/streamingbatchwriter/streamingbatchwriter.go b/writers/streamingbatchwriter/streamingbatchwriter.go index 3951a01a9c..93bb0dfb99 100644 --- a/writers/streamingbatchwriter/streamingbatchwriter.go +++ b/writers/streamingbatchwriter/streamingbatchwriter.go @@ -14,6 +14,7 @@ // 2. The batch size is reached // 3. The batch size in bytes is reached // 4. A different message type is received +// 5. A message.WriteFlush is received: every open batch completes, and the flush is done once their handlers return // // Each handler can get invoked multiple times as new batches are processed. // Handlers get invoked only if there's a message of that type at hand: First message of the batch is immediately available in the channel. @@ -154,6 +155,43 @@ func (w *StreamingBatchWriter) Flush(context.Context) error { return nil // not checked below } +// commit ends the open batches and waits for their handlers to return. +func (w *StreamingBatchWriter) commit(ctx context.Context) error { + w.workersLock.RLock() + defer w.workersLock.RUnlock() + errs := make([]error, 0, 3+len(w.insertWorkers)) + errs = append(errs, + commitWorker(ctx, w.migrateWorker), + commitWorker(ctx, w.deleteStaleWorker), + commitWorker(ctx, w.deleteRecordWorker), + ) + for tableName, worker := range w.insertWorkers { + if err := commitWorker(ctx, worker); err != nil { + errs = append(errs, &message.FailedTablesError{Tables: []string{tableName}, Err: err}) + } + } + return errors.Join(errs...) +} + +func commitWorker[T message.WriteMessage](ctx context.Context, s *streamingWorkerManager[T]) error { + if s == nil { + return nil + } + // buffered, so the worker does not block if ctx is done before it replies + done := make(chan error, 1) + select { + case s.commit <- done: + case <-ctx.Done(): + return ctx.Err() + } + select { + case err := <-done: + return err + case <-ctx.Done(): + return ctx.Err() + } +} + func (w *StreamingBatchWriter) Close(context.Context) error { w.workersLock.Lock() defer w.workersLock.Unlock() @@ -187,13 +225,22 @@ func (w *StreamingBatchWriter) Write(ctx context.Context, msgs <-chan message.Wr go func() { defer close(errCh) defer w.Close(ctx) + // a rejected message is never written, so every later flush must fail + var rejectErr error for msg := range msgs { + if flush, ok := msg.(*message.WriteFlush); ok { + flush.Done(errors.Join(rejectErr, w.commit(ctx))) + continue + } msgType := writers.MsgID(msg) if w.lastMsgType != writers.MsgTypeUnset && w.lastMsgType != msgType { _ = w.Flush(ctx) } w.lastMsgType = msgType if err := w.startWorker(ctx, errCh, msg); err != nil { + if rejectErr == nil { + rejectErr = err + } errCh <- err } } @@ -239,8 +286,9 @@ func (w *StreamingBatchWriter) startWorker(ctx context.Context, errCh chan<- err ch: make(chan *message.WriteMigrateTable), writeFunc: w.client.MigrateTable, - flush: make(chan chan bool), - errCh: errCh, + flush: make(chan chan bool), + commit: make(chan chan error), + errCh: errCh, tableName: tableName, limit: batch.CappedAt(0, w.batchSizeRows), @@ -268,8 +316,9 @@ func (w *StreamingBatchWriter) startWorker(ctx context.Context, errCh chan<- err writeFunc: w.client.DeleteStale, tableName: tableName, - flush: make(chan chan bool), - errCh: errCh, + flush: make(chan chan bool), + commit: make(chan chan error), + errCh: errCh, limit: batch.CappedAt(0, w.batchSizeRows), batchTimeout: w.batchTimeout, @@ -305,8 +354,9 @@ func (w *StreamingBatchWriter) startWorker(ctx context.Context, errCh chan<- err writeFunc: w.client.WriteTable, tableName: tableName, - flush: make(chan chan bool), - errCh: errCh, + flush: make(chan chan bool), + commit: make(chan chan error), + errCh: errCh, limit: batch.CappedAt(w.batchSizeBytes, w.batchSizeRows), batchTimeout: w.batchTimeout, @@ -337,8 +387,9 @@ func (w *StreamingBatchWriter) startWorker(ctx context.Context, errCh chan<- err writeFunc: w.client.DeleteRecords, tableName: tableName, - flush: make(chan chan bool), - errCh: errCh, + flush: make(chan chan bool), + commit: make(chan chan error), + errCh: errCh, limit: batch.CappedAt(w.batchSizeBytes, w.batchSizeRows), batchTimeout: w.batchTimeout, @@ -362,8 +413,11 @@ type streamingWorkerManager[T message.WriteMessage] struct { writeFunc func(context.Context, <-chan T) error tableName string - flush chan chan bool - errCh chan<- error + flush chan chan bool + commit chan chan error + errCh chan<- error + // a handler failed or panicked, so some messages may not be written + errored atomic.Bool limit *batch.Cap batchTimeout time.Duration @@ -411,6 +465,7 @@ func (s *streamingWorkerManager[T]) send(ctx context.Context, data T) { defer s.workerWg.Done() defer func() { if msg := recover(); msg != nil { + s.errored.Store(true) switch v := msg.(type) { case error: s.errCh <- fmt.Errorf("panic processing %s: %w [recovered]", s.tableName, v) @@ -430,6 +485,7 @@ func (s *streamingWorkerManager[T]) send(ctx context.Context, data T) { err := s.writeFunc(ctx, ch) if err != nil { + s.errored.Store(true) s.failed.Store(true) go func() { // nolint:revive @@ -510,6 +566,15 @@ func (s *streamingWorkerManager[T]) run(ctx context.Context, wg *sync.WaitGroup) ticker.Reset(s.batchTimeout) } done <- true + case done := <-s.commit: + s.closeFlush() + s.workerWg.Wait() + ticker.Reset(s.batchTimeout) + if s.errored.Load() { + done <- fmt.Errorf("failed to write messages for %s", s.tableName) + } else { + done <- nil + } case <-ctxDone: // this means the request was cancelled return // after this NO other call will succeed diff --git a/writers/tables.go b/writers/tables.go new file mode 100644 index 0000000000..45d71adb83 --- /dev/null +++ b/writers/tables.go @@ -0,0 +1,39 @@ +package writers + +import ( + "github.com/cloudquery/plugin-sdk/v4/message" + "github.com/cloudquery/plugin-sdk/v4/schema" +) + +// FailedTables names the tables of msgs in err, so that a flush that fails with it reports them. It returns err as is +// if it is nil, or if a message has no table name. +func FailedTables[T message.WriteMessage](err error, msgs []T) error { + if err == nil { + return nil + } + tables := make([]string, 0, len(msgs)) + for _, msg := range msgs { + name := tableName(msg) + if name == "" { + return err + } + tables = append(tables, name) + } + return &message.FailedTablesError{Tables: tables, Err: err} +} + +func tableName(msg message.WriteMessage) string { + switch m := msg.(type) { + case *message.WriteInsert: + name, _ := m.Record.Schema().Metadata().GetValue(schema.MetadataTableName) + return name + case *message.WriteMigrateTable: + return m.Table.Name + case *message.WriteDeleteStale: + return m.TableName + case *message.WriteDeleteRecord: + return m.TableName + default: + return "" + } +} diff --git a/writers/tables_test.go b/writers/tables_test.go new file mode 100644 index 0000000000..eb07ab28c3 --- /dev/null +++ b/writers/tables_test.go @@ -0,0 +1,37 @@ +package writers + +import ( + "errors" + "testing" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/cloudquery/plugin-sdk/v4/message" + "github.com/cloudquery/plugin-sdk/v4/schema" + "github.com/stretchr/testify/require" +) + +func TestFailedTables(t *testing.T) { + table := &schema.Table{Name: "inserted", Columns: schema.ColumnList{{Name: "id", Type: arrow.PrimitiveTypes.Int64}}} + bldr := array.NewRecordBuilder(memory.DefaultAllocator, table.ToArrowSchema()) + defer bldr.Release() + insert := &message.WriteInsert{Record: bldr.NewRecordBatch()} + failed := errors.New("write failed") + + err := FailedTables(failed, []message.WriteMessage{ + insert, + &message.WriteMigrateTable{Table: &schema.Table{Name: "migrated"}}, + &message.WriteDeleteStale{TableName: "stale"}, + &message.WriteDeleteRecord{DeleteRecord: message.DeleteRecord{TableName: "deleted"}}, + }) + require.ErrorIs(t, err, failed) + require.Equal(t, []string{"deleted", "inserted", "migrated", "stale"}, message.FailedTables(err)) + + require.NoError(t, FailedTables(nil, []message.WriteMessage{insert})) + + noName := array.NewRecordBuilder(memory.DefaultAllocator, arrow.NewSchema([]arrow.Field{{Name: "id", Type: arrow.PrimitiveTypes.Int64}}, nil)) + defer noName.Release() + err = FailedTables(failed, []message.WriteMessage{insert, &message.WriteInsert{Record: noName.NewRecordBatch()}}) + require.Equal(t, failed, err, "a message without a table may belong to any table") +}