sql: ensure that descriptor lease drains are traced

Release note: None

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

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)
 	}
 }
 

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.

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")