From 07cb0217db75b19d2db3fc3d1f80fb5a2c055b29 Mon Sep 17 00:00:00 2001 From: Kemal Hadimli Date: Tue, 13 Oct 2026 12:27:28 +0100 Subject: [PATCH] feat: Save unflushed state on Close unless the sync was cancelled --- internal/clients/state/v3/state.go | 32 ++++++- serve/state_test.go | 134 +++++++++++++++++++++++++++++ state/client.go | 2 + 3 files changed, 166 insertions(+), 2 deletions(-) diff --git a/internal/clients/state/v3/state.go b/internal/clients/state/v3/state.go index cdd9aae524..46748a5418 100644 --- a/internal/clients/state/v3/state.go +++ b/internal/clients/state/v3/state.go @@ -3,6 +3,7 @@ package state import ( "bytes" "context" + "errors" "io" "sync" @@ -11,6 +12,7 @@ import ( "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/premium" "github.com/cloudquery/plugin-sdk/v4/schema" "google.golang.org/grpc" ) @@ -22,6 +24,7 @@ const ( ) type Client struct { + ctx context.Context client pb.PluginClient mem map[string]versionedValue changes map[string]struct{} // changed keys @@ -39,6 +42,7 @@ type versionedValue struct { func NewClient(ctx context.Context, conn *grpc.ClientConn, tableName string) (*Client, error) { table := Table(tableName) c := &Client{ + ctx: ctx, conn: conn, client: pb.NewPluginClient(conn), mem: make(map[string]versionedValue), @@ -149,6 +153,13 @@ func (c *Client) SetKey(_ context.Context, key string, value string) error { func (c *Client) Flush(ctx context.Context) error { c.mutex.Lock() defer c.mutex.Unlock() + return c.flushLocked(ctx) +} + +func (c *Client) flushLocked(ctx context.Context) error { + if len(c.changes) == 0 { + return nil + } bldr := array.NewRecordBuilder(memory.DefaultAllocator, c.schema) keys, values := bldr.Field(0).(*array.StringBuilder), bldr.Field(1).(*array.StringBuilder) var version *array.Uint64Builder @@ -195,11 +206,28 @@ func (c *Client) GetKey(_ context.Context, key string) (string, error) { return c.mem[key].value, nil } +// Close saves changes not yet flushed, then closes the connection. +// It drops them instead when the context passed to NewClient was cancelled for a reason other than quota: +// the sync stopped early, and moving the cursors forward would skip records that were never written. func (c *Client) Close() error { + c.mutex.Lock() + var err error + if !stoppedEarly(c.ctx) { + err = c.flushLocked(context.WithoutCancel(c.ctx)) + } + c.mutex.Unlock() if c.conn != nil { - return c.conn.Close() + return errors.Join(err, c.conn.Close()) } - return nil + return err +} + +func stoppedEarly(ctx context.Context) bool { + if ctx.Err() == nil { + return false + } + var noQuota premium.ErrNoQuota + return !errors.As(context.Cause(ctx), &noQuota) } func Table(name string) *schema.Table { diff --git a/serve/state_test.go b/serve/state_test.go index 02c2d21a13..37bb168deb 100644 --- a/serve/state_test.go +++ b/serve/state_test.go @@ -3,12 +3,15 @@ package serve import ( "context" "sync" + "sync/atomic" "testing" + "time" pb "github.com/cloudquery/plugin-pb-go/pb/plugin/v3" "github.com/cloudquery/plugin-sdk/v4/internal/clients/state/v3" "github.com/cloudquery/plugin-sdk/v4/internal/memdb" "github.com/cloudquery/plugin-sdk/v4/plugin" + "github.com/cloudquery/plugin-sdk/v4/premium" "google.golang.org/grpc" "google.golang.org/grpc/credentials/insecure" ) @@ -156,3 +159,134 @@ func TestStateOverwriteGetLatest(t *testing.T) { t.Fatal(serverErr) } } + +func TestStateCloseSavesUnflushedChanges(t *testing.T) { + ctx := context.Background() + dial := serveStateBackend(t) + + stateClient, err := state.NewClient(ctx, dial(), "test") + if err != nil { + t.Fatal(err) + } + if err := stateClient.SetKey(ctx, "key", "value"); err != nil { + t.Fatal(err) + } + if err := stateClient.Close(); err != nil { + t.Fatal(err) + } + + assertStateValue(ctx, t, dial(), "value") +} + +func TestStateCloseDropsUnflushedChangesWhenSyncIsCancelled(t *testing.T) { + ctx := context.Background() + dial := serveStateBackend(t) + + syncCtx, cancel := context.WithCancel(ctx) + stateClient, err := state.NewClient(syncCtx, dial(), "test") + if err != nil { + t.Fatal(err) + } + if err := stateClient.SetKey(ctx, "key", "saved"); err != nil { + t.Fatal(err) + } + if err := stateClient.Flush(ctx); err != nil { + t.Fatal(err) + } + if err := stateClient.SetKey(ctx, "key", "dropped"); err != nil { + t.Fatal(err) + } + cancel() + if err := stateClient.Close(); err != nil { + t.Fatal(err) + } + + assertStateValue(ctx, t, dial(), "saved") +} + +func TestStateCloseSavesUnflushedChangesWhenQuotaRunsOut(t *testing.T) { + ctx := context.Background() + dial := serveStateBackend(t) + + quota := "aMonitor{} + syncCtx, err := premium.WithCancelOnQuotaExceeded(ctx, quota, premium.WithQuotaCheckPeriod(time.Millisecond)) + if err != nil { + t.Fatal(err) + } + stateClient, err := state.NewClient(syncCtx, dial(), "test") + if err != nil { + t.Fatal(err) + } + if err := stateClient.SetKey(ctx, "key", "value"); err != nil { + t.Fatal(err) + } + quota.exceeded.Store(true) + <-syncCtx.Done() + if err := stateClient.Close(); err != nil { + t.Fatal(err) + } + + assertStateValue(ctx, t, dial(), "value") +} + +type quotaMonitor struct { + exceeded atomic.Bool +} + +func (*quotaMonitor) TeamName() string { + return "test-team" +} + +func (q *quotaMonitor) CheckQuota(context.Context) (premium.CheckQuotaResult, error) { + return premium.CheckQuotaResult{HasQuota: !q.exceeded.Load()}, nil +} + +func assertStateValue(ctx context.Context, t *testing.T, conn *grpc.ClientConn, want string) { + t.Helper() + stateClient, err := state.NewClient(ctx, conn, "test") + if err != nil { + t.Fatal(err) + } + val, err := stateClient.GetKey(ctx, "key") + if err != nil { + t.Fatal(err) + } + if val != want { + t.Fatalf("expected value to be %q but got %q", want, val) + } +} + +// serveStateBackend serves a memdb destination and returns a func that opens a new connection to it. +// Each state client needs its own connection, because Close closes it. +func serveStateBackend(t *testing.T) func() *grpc.ClientConn { + t.Helper() + p := plugin.NewPlugin("testPluginV3", "v1.0.0", memdb.NewMemDBClient) + 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() + }) + + dial := func() *grpc.ClientConn { + // nolint:staticcheck + conn, err := grpc.DialContext(ctx, "bufnet", grpc.WithContextDialer(srv.bufPluginDialer), grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithBlock()) + if err != nil { + t.Fatalf("Failed to dial bufnet: %v", err) + } + t.Cleanup(func() { _ = conn.Close() }) + return conn + } + if _, err := pb.NewPluginClient(dial()).Init(ctx, &pb.Init_Request{}); err != nil { + t.Fatal(err) + } + return dial +} diff --git a/state/client.go b/state/client.go index cca64bbb41..0824be5d08 100644 --- a/state/client.go +++ b/state/client.go @@ -18,6 +18,8 @@ type Client interface { SetKey(ctx context.Context, key string, value string) error GetKey(ctx context.Context, key string) (string, error) Flush(ctx context.Context) error + // Close saves changes not yet flushed, unless the context the client was created with + // was cancelled for a reason other than quota, then closes the connection. Close() error }