fix: Keep checkpoint order for unchanged state values

internal/clients/state/v3/state.go

@@ -166,22 +166,22 @@ 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))
+	changed := c.mem[key].value != value
+	if changed {
+		c.mem[key] = versionedValue{
+			value:   value,
+			version: c.mem[key].version + 1,
 		}
-		return nil // don't update if the value is the same
-	}
-	c.mem[key] = versionedValue{
-		value:   value,
-		version: c.mem[key].version + 1,
 	}
 	if c.checkpoints == nil {
-		c.changes[key] = nil
+		if changed {
+			c.changes[key] = nil
+		}
 		return nil
 	}
+	// with checkpoints, even an unchanged value covers the caller's resources, so later checkpoints of the key must
+	// follow them
+	scopes, pending := c.changes[key]
 	if !pending {
 		scopes = newScopeSet()
 		c.changes[key] = scopes

scheduler/checkpoint_test.go

@@ -936,3 +936,59 @@ func TestCheckpointsDroppedOnCancelAreAborted(t *testing.T) {
 		require.True(t, sent.Aborted())
 	})
 }
+
+// gatedScope holds the checkpoints passed to it until released, as a table resolver's scope holds them until the rows
+// it produced are sent.
+type gatedScope struct {
+	mu   sync.Mutex
+	held []message.SyncMessage
+}
+
+func (g *gatedScope) Checkpoint(_ context.Context, msg message.SyncMessage) error {
+	g.mu.Lock()
+	defer g.mu.Unlock()
+	g.held = append(g.held, msg)
+	return nil
+}
+
+func (g *gatedScope) release(ctx context.Context, res chan<- message.SyncMessage) error {
+	g.mu.Lock()
+	held := g.held
+	g.held = nil
+	g.mu.Unlock()
+	for _, msg := range held {
+		if err := checkpoint.Deliver(ctx, res, msg); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func TestUnchangedStateValueOrdersLaterCheckpointsAfterTheCallersResources(t *testing.T) {
+	ctx := context.Background()
+	res := make(chan message.SyncMessage, 10)
+	stateClient := newCheckpointStateClient(t, res)
+	sent := func() []string {
+		var values []string
+		for len(res) > 0 {
+			values = append(values, stateValues(t, message.SyncMessages{<-res})["cursor"])
+		}
+		return values
+	}
+
+	require.NoError(t, stateClient.SetKey(ctx, "cursor", "1"))
+	require.NoError(t, stateClient.Flush(ctx))
+	require.Equal(t, []string{"1"}, sent())
+
+	// the same value, set by a resolver whose rows are not sent yet: it covers them too
+	unsent := &gatedScope{}
+	resolverCtx := checkpoint.WithScope(ctx, unsent)
+	require.NoError(t, stateClient.SetKey(resolverCtx, "cursor", "1"))
+	require.NoError(t, stateClient.Flush(resolverCtx))
+	require.NoError(t, stateClient.SetKey(ctx, "cursor", "2"))
+	require.NoError(t, stateClient.Flush(ctx))
+	require.Empty(t, sent(), "a later value must not overtake the rows the unchanged value covers")
+
+	require.NoError(t, unsent.release(ctx, res))
+	require.Equal(t, []string{"1", "2"}, sent())
+}