diff --git a/client/internal/engine.go b/client/internal/engine.go index 5f5b57eaf..2bc6e723f 100644 --- a/client/internal/engine.go +++ b/client/internal/engine.go @@ -1586,13 +1586,18 @@ func (e *Engine) receiveManagementEvents() { log.Infof("sync finished in %s", d) e.clientMetrics.RecordSyncDuration(e.ctx, d) }) - e.shutdownWg.Add(1) + // The manager only has work while the stream can feed it, so it lives as + // long as Sync and not as long as the engine. + managerCtx, cancelManager := context.WithCancel(e.ctx) + managerDone := make(chan struct{}) go func() { - defer e.shutdownWg.Done() - manager.run(e.ctx) + defer close(managerDone) + manager.run(managerCtx) }() err := e.mgmClient.Sync(e.ctx, e.syncInfoFunc(info), manager.SetTarget) + cancelManager() + <-managerDone if err != nil { // happens if management is unavailable for a long time. // We want to cancel the operation of the whole client diff --git a/client/internal/mapsync.go b/client/internal/mapsync.go index a60f55f7a..f91814078 100644 --- a/client/internal/mapsync.go +++ b/client/internal/mapsync.go @@ -143,6 +143,10 @@ func (m *mapStateManager) run(ctx context.Context) { // true on a fresh target and on a coalesced switch to a newer target mid-flight. var passGen uint64 for { + if ctx.Err() != nil { + return + } + m.mu.Lock() target, tg, ag := m.target, m.targetGen, m.appliedGen m.mu.Unlock() diff --git a/client/internal/mapsync_test.go b/client/internal/mapsync_test.go index d776fef5a..30b838b2b 100644 --- a/client/internal/mapsync_test.go +++ b/client/internal/mapsync_test.go @@ -279,3 +279,33 @@ func TestMapStateManager_ReappliesOnNewTarget(t *testing.T) { t.Fatal("new target not applied") } } + +// cancelling the context while a target still has passes left stops the loop +// between passes instead of draining the target to convergence. +func TestMapStateManager_StopsBetweenPassesOnCancel(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + var passes atomic.Int32 + apply := func(_ *mgmProto.SyncResponse, _ bool) (bool, error) { + if passes.Add(1) == 1 { + cancel() + } + return true, nil // never converges on its own + } + m := newMapStateManager(apply, nil, nil) + require.NoError(t, m.SetTarget(&mgmProto.SyncResponse{})) + + done := make(chan struct{}) + go func() { + defer close(done) + m.run(ctx) + }() + + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("run kept applying passes after the context was cancelled") + } + require.EqualValues(t, 1, passes.Load(), "no pass should run after the cancelling one") +}