feat: Save unflushed state on Close unless the sync was cancelled

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 {

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 := &quotaMonitor{}
+	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
+}

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
 }