diff --git a/internal/clients/state/v3/state.go b/internal/clients/state/v3/state.go index e61441ec66..24860f2e10 100644 --- a/internal/clients/state/v3/state.go +++ b/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 diff --git a/scheduler/checkpoint_test.go b/scheduler/checkpoint_test.go index 06084d9b04..41a3b57684 100644 --- a/scheduler/checkpoint_test.go +++ b/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()) +}