From 568cb582153f75fc557c8f5c87b7b350ae41c428 Mon Sep 17 00:00:00 2001 From: Raphael 'kena' Poss Date: Thu, 11 Nov 2021 13:46:12 +0100 Subject: [PATCH] sql: ensure that descriptor lease drains are traced Release note: None --- pkg/server/drain.go | 2 +- pkg/sql/catalog/lease/descriptor_state.go | 2 +- pkg/sql/catalog/lease/lease.go | 13 +++++++------ pkg/sql/catalog/lease/lease_test.go | 7 ++++--- 4 files changed, 13 insertions(+), 11 deletions(-) diff --git a/pkg/server/drain.go b/pkg/server/drain.go index dd733acc1ffc..80755261e607 100644 --- a/pkg/server/drain.go +++ b/pkg/server/drain.go @@ -207,7 +207,7 @@ func (s *Server) drainClients(ctx context.Context, reporter func(int, redact.Saf // Drain the SQL leases. This must be done after the pgServer has // given sessions a chance to finish ongoing work. - s.sqlServer.leaseMgr.SetDraining(true /* drain */, reporter) + s.sqlServer.leaseMgr.SetDraining(ctx, true /* drain */, reporter) // Done. This executes the defers set above to drain SQL leases. return nil diff --git a/pkg/sql/catalog/lease/descriptor_state.go b/pkg/sql/catalog/lease/descriptor_state.go index 30d08df0d5b8..a6c75aa7b7c3 100644 --- a/pkg/sql/catalog/lease/descriptor_state.go +++ b/pkg/sql/catalog/lease/descriptor_state.go @@ -260,7 +260,7 @@ func (t *descriptorState) release(ctx context.Context, s *descriptorVersionState return nil } if l := maybeRemoveLease(); l != nil { - releaseLease(l, t.m) + releaseLease(ctx, l, t.m) } } diff --git a/pkg/sql/catalog/lease/lease.go b/pkg/sql/catalog/lease/lease.go index f707688cc8ed..3d2522be273d 100644 --- a/pkg/sql/catalog/lease/lease.go +++ b/pkg/sql/catalog/lease/lease.go @@ -459,7 +459,7 @@ func acquireNodeLease(ctx context.Context, m *Manager, id descpb.ID) (bool, erro m.names.insert(newDescVersionState) } if toRelease != nil { - releaseLease(toRelease, m) + releaseLease(ctx, toRelease, m) } return true, nil }) @@ -475,8 +475,7 @@ func acquireNodeLease(ctx context.Context, m *Manager, id descpb.ID) (bool, erro } // releaseLease from store. -func releaseLease(lease *storedLease, m *Manager) { - ctx := context.TODO() +func releaseLease(ctx context.Context, lease *storedLease, m *Manager) { if m.isDraining() { // Release synchronously to guarantee release before exiting. m.storage.release(ctx, m.stopper, lease) @@ -530,7 +529,7 @@ func purgeOldVersions( leases := t.removeInactiveVersions() t.mu.Unlock() for _, l := range leases { - releaseLease(l, m) + releaseLease(ctx, l, m) } } @@ -952,7 +951,9 @@ func (m *Manager) isDraining() bool { // to report work that needed to be done and which may or may not have // been done by the time this call returns. See the explanation in // pkg/server/drain.go for details. -func (m *Manager) SetDraining(drain bool, reporter func(int, redact.SafeString)) { +func (m *Manager) SetDraining( + ctx context.Context, drain bool, reporter func(int, redact.SafeString), +) { m.draining.Store(drain) if !drain { return @@ -965,7 +966,7 @@ func (m *Manager) SetDraining(drain bool, reporter func(int, redact.SafeString)) leases := t.removeInactiveVersions() t.mu.Unlock() for _, l := range leases { - releaseLease(l, m) + releaseLease(ctx, l, m) } if reporter != nil { // Report progress through the Drain RPC. diff --git a/pkg/sql/catalog/lease/lease_test.go b/pkg/sql/catalog/lease/lease_test.go index db60bd19260f..128145126a64 100644 --- a/pkg/sql/catalog/lease/lease_test.go +++ b/pkg/sql/catalog/lease/lease_test.go @@ -487,6 +487,7 @@ func TestLeaseManagerDrain(testingT *testing.T) { t := newLeaseTest(testingT, params) defer t.cleanup() + ctx := context.Background() const descID = keys.LeaseTableID { @@ -499,8 +500,8 @@ func TestLeaseManagerDrain(testingT *testing.T) { // starts draining. l1RemovalTracker := leaseRemovalTracker.TrackRemoval(l1.Underlying()) - t.nodes[1].SetDraining(true, nil /* reporter */) - t.nodes[2].SetDraining(true, nil /* reporter */) + t.nodes[1].SetDraining(ctx, true, nil /* reporter */) + t.nodes[2].SetDraining(ctx, true, nil /* reporter */) // Leases cannot be acquired when in draining mode. if _, err := t.acquire(1, descID); !testutils.IsError(err, "cannot acquire lease when draining") { @@ -523,7 +524,7 @@ func TestLeaseManagerDrain(testingT *testing.T) { { // Check that leases with a refcount of 0 are correctly kept in the // store once the drain mode has been exited. - t.nodes[1].SetDraining(false, nil /* reporter */) + t.nodes[1].SetDraining(ctx, false, nil /* reporter */) l1 := t.mustAcquire(1, descID) t.mustRelease(1, l1, nil) t.expectLeases(descID, "/1/1")