Merge #149264

149264: catalog/lease: add memory monitoring for lease manager r=fqazi a=fqazi

Previously, there was no memory monitoring to track how much memory was being used by the lease manager. This made it hard to understand how memory was being utilized with a large number of schema objects. This patch adds support for memory monitoring inside the lease manager.

Fixes: #71282

Release note: None

Co-authored-by: Faizan Qazi <[email protected]>

Author: craig[bot]

pkg/server/drain.go

@@ -486,7 +486,7 @@ func (s *drainServer) drainClients(
 	// Drain all SQL table leases. This must be done after the pgServer has
 	// given sessions a chance to finish ongoing work and after the background
 	// tasks that may issue SQL statements have shut down.
-	s.sqlServer.leaseMgr.SetDraining(ctx, true /* drain */, reporter)
+	s.sqlServer.leaseMgr.SetDraining(ctx, true /* drain */, reporter, true /*assertOnLeakedDescriptor*/)
 
 	session, err := s.sqlServer.sqlLivenessProvider.Release(ctx)
 	if err != nil {

pkg/server/server_sql.go

@@ -667,6 +667,7 @@ func newSQLServer(ctx context.Context, cfg sqlServerArgs) (*SQLServer, error) {
 	}
 
 	leaseMgr := lease.NewLeaseManager(
+		ctx,
 		cfg.AmbientCtx,
 		cfg.nodeIDContainer,
 		cfg.internalDB,
@@ -678,6 +679,7 @@ func newSQLServer(ctx context.Context, cfg sqlServerArgs) (*SQLServer, error) {
 		lmKnobs,
 		cfg.stopper,
 		cfg.rangeFeedFactory,
+		cfg.monitorAndMetrics.rootSQLMemoryMonitor,
 	)
 	cfg.registry.AddMetricStruct(leaseMgr.MetricsStruct())
 

pkg/sql/catalog/lease/BUILD.bazel

@@ -54,6 +54,7 @@ go_library(
         "//pkg/util/log",
         "//pkg/util/log/logcrash",
         "//pkg/util/metric",
+        "//pkg/util/mon",
         "//pkg/util/quotapool",
         "//pkg/util/retry",
         "//pkg/util/startup",

pkg/sql/catalog/lease/descriptor_state.go

@@ -135,6 +135,12 @@ func (t *descriptorState) upsertLeaseLocked(
 			log.Infof(ctx, "new lease: %s", desc)
 		}
 		descState := newDescriptorVersionState(t, desc, hlc.Timestamp{}, session, regionEnumPrefix, true /* isLease */)
+		if err := t.m.boundAccount.Grow(ctx, descState.getByteSize()); err != nil {
+			// If we don't have sufficient memory, then release the lease so
+			// that the system.lease table doesn't have a reference.
+			t.m.storage.release(ctx, t.m.stopper, descState.mu.lease)
+			return wrapMemoryError(err)
+		}
 		t.mu.active.insert(descState)
 		t.m.names.insert(ctx, descState)
 		return nil
@@ -196,7 +202,7 @@ func newDescriptorVersionState(
 
 // removeInactiveVersions removes inactive versions in t.mu.active.data with
 // refcount 0. t.mu must be locked. It returns leases that need to be released.
-func (t *descriptorState) removeInactiveVersions() []*storedLease {
+func (t *descriptorState) removeInactiveVersions(ctx context.Context) []*storedLease {
 	var leases []*storedLease
 	// A copy of t.mu.active.data must be made since t.mu.active.data will be changed
 	// within the loop.
@@ -206,10 +212,12 @@ func (t *descriptorState) removeInactiveVersions() []*storedLease {
 			func() {
 				desc.mu.Lock()
 				defer desc.mu.Unlock()
+				// Ensure we have a lock to allow us to clean up the usage
+				// by the stored lease.
+				t.m.boundAccount.Shrink(ctx, desc.getByteSize())
 				if l := desc.mu.lease; l != nil {
 					desc.mu.lease = nil
 					leases = append(leases, l)
-
 				}
 			}()
 		}
@@ -272,14 +280,15 @@ func (t *descriptorState) release(ctx context.Context, s *descriptorVersionState
 			leaseReleased := true
 			// For testing, we will synchronously release leases, but that
 			// exposes us to the danger of the context getting cancelled. To
-			// eliminate this risk, we are going first remove the lease from
-			// storage and then delete if from mqemory.
+			// eliminate this risk, we are going to first remove the lease from
+			// storage and then delete it from memory.
 			if t.m.storage.testingKnobs.RemoveOnceDereferenced {
 				leaseReleased = releaseLease(ctx, l, t.m)
 				l = nil
 			}
 			if leaseReleased {
 				t.mu.active.remove(s)
+				s.t.m.boundAccount.Shrink(ctx, s.getByteSize())
 			}
 			return l
 		}

pkg/sql/catalog/lease/descriptor_version_state.go

@@ -8,6 +8,7 @@ package lease
 import (
 	"context"
 	"sync/atomic"
+	"unsafe"
 
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
@@ -39,6 +40,12 @@ func (s *storedLease) SafeFormat(w redact.SafePrinter, _ rune) {
 	w.Printf("ID=%d ver=%d expiration=%s", s.id, s.version, s.expiration)
 }
 
+// getByteSize returns the full size of stored lease.
+func (s *storedLease) getByteSize() int64 {
+	return int64(len(s.prefix)+len(s.sessionID)+int(s.expiration.Size())) +
+		int64(unsafe.Sizeof(*s))
+}
+
 // descriptorVersionState holds the state for a descriptor version. This
 // includes the lease information for a descriptor version.
 // TODO(vivek): A node only needs to manage lease information on what it
@@ -94,6 +101,15 @@ func (s *descriptorVersionState) SafeFormat(w redact.SafePrinter, _ rune) {
 	w.Print(s.redactedString())
 }
 
+// getByteSize returns the full size of a descriptor version state structure.
+func (s *descriptorVersionState) getByteSize() int64 {
+	size := s.ByteSize() + int64(unsafe.Sizeof(*s))
+	if s.mu.lease != nil {
+		size += s.mu.lease.getByteSize()
+	}
+	return size
+}
+
 func (s *descriptorVersionState) String() string {
 	return redact.StringWithoutMarkers(s)
 }

pkg/sql/catalog/lease/lease.go

@@ -41,10 +41,12 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlliveness"
 	kvstorage "github.com/cockroachdb/cockroach/pkg/storage"
+	"github.com/cockroachdb/cockroach/pkg/util/buildutil"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/log/logcrash"
 	"github.com/cockroachdb/cockroach/pkg/util/metric"
+	"github.com/cockroachdb/cockroach/pkg/util/mon"
 	"github.com/cockroachdb/cockroach/pkg/util/quotapool"
 	"github.com/cockroachdb/cockroach/pkg/util/retry"
 	"github.com/cockroachdb/cockroach/pkg/util/startup"
@@ -777,6 +779,12 @@ func (m *Manager) readOlderVersionForTimestamp(
 	return descs, nil
 }
 
+// wrapMemoryError adds a hint on memory errors to indicate
+// which setting should be bumped.
+func wrapMemoryError(err error) error {
+	return errors.WithHint(err, "Consider increasing --max-sql-memory startup parameter.")
+}
+
 // Insert descriptor versions. The versions provided are not in
 // any particular order.
 func (m *Manager) insertDescriptorVersions(
@@ -789,16 +797,25 @@ func (m *Manager) insertDescriptorVersions(
 	}
 	t.mu.Lock()
 	defer t.mu.Unlock()
+	newVersionsToInsert := make([]*descriptorVersionState, 0, len(versions))
 	for i := range versions {
 		// Since we gave up the lock while reading the versions from
 		// the store we have to ensure that no one else inserted the
 		// same version.
 		existingVersion := t.mu.active.findVersion(versions[i].desc.GetVersion())
 		if existingVersion == nil {
-			t.mu.active.insert(
-				newDescriptorVersionState(t, versions[i].desc, versions[i].expiration, session, nil, false))
+			descState := newDescriptorVersionState(t, versions[i].desc, versions[i].expiration, session, nil, false)
+			if err := t.m.boundAccount.Grow(ctx, descState.getByteSize()); err != nil {
+				return wrapMemoryError(err)
+			}
+			newVersionsToInsert = append(newVersionsToInsert, descState)
 		}
 	}
+	// Only insert if all versions were allocated.
+	for _, descState := range newVersionsToInsert {
+		t.mu.active.insert(descState)
+	}
+
 	return nil
 }
 
@@ -978,7 +995,7 @@ func purgeOldVersions(
 			t.mu.Lock()
 			defer t.mu.Unlock()
 			t.mu.takenOffline = dropped
-			return t.removeInactiveVersions(), t.mu.active.findPreviousToExpire(dropped)
+			return t.removeInactiveVersions(ctx), t.mu.active.findPreviousToExpire(dropped)
 		}()
 		for _, l := range leases {
 			releaseLease(ctx, l, m)
@@ -1122,6 +1139,11 @@ type Manager struct {
 	// initComplete is a fast check to confirm that initialization is complete, since
 	// performance testing showed select on the waitForInit channel can be expensive.
 	initComplete atomic.Bool
+
+	// bytesMonitor tracks the memory usage from leased descriptors.
+	bytesMonitor *mon.BytesMonitor
+	// boundAccount tracks the memory usage from leased descriptors.
+	boundAccount *mon.ConcurrentBoundAccount
 }
 
 const leaseConcurrencyLimit = 5
@@ -1133,6 +1155,7 @@ const leaseConcurrencyLimit = 5
 //
 // stopper is used to run async tasks. Can be nil in tests.
 func NewLeaseManager(
+	ctx context.Context,
 	ambientCtx log.AmbientContext,
 	nodeIDContainer *base.SQLIDContainer,
 	db isql.DB,
@@ -1144,7 +1167,11 @@ func NewLeaseManager(
 	testingKnobs ManagerTestingKnobs,
 	stopper *stop.Stopper,
 	rangeFeedFactory *rangefeed.Factory,
+	rootBytesMonitor *mon.BytesMonitor,
 ) *Manager {
+	// See pkg/sql/mem_metrics.go
+	// log10int64times1000 = log10(math.MaxInt64) * 1000, rounded up somewhat
+	const log10int64times1000 = 19 * 1000
 	lm := &Manager{
 		storage: storage{
 			nodeIDContainer:  nodeIDContainer,
@@ -1199,6 +1226,24 @@ func NewLeaseManager(
 					Measurement: "Number of wait for initial version routines executing",
 					Unit:        metric.Unit_COUNT,
 				}),
+				leaseCurBytesCount: metric.NewGauge(metric.Metadata{
+					Name:        "sql.leases.lease_cur_bytes_count",
+					Help:        "The current number of bytes used by the lease manager.",
+					Measurement: "Number of bytes used by the lease manager.",
+					Unit:        metric.Unit_BYTES,
+				}),
+				leaseMaxBytesHist: metric.NewHistogram(metric.HistogramOptions{
+					Metadata: metric.Metadata{
+						Name:        "sql.leases.lease_max_bytes_hist",
+						Help:        "Memory used by the lease manager.",
+						Measurement: "Number of bytes used by the lease manager.",
+						Unit:        metric.Unit_BYTES,
+					},
+					Duration:     base.DefaultHistogramWindowInterval(),
+					MaxVal:       log10int64times1000,
+					SigFigs:      3,
+					BucketConfig: metric.MemoryUsage64MBBuckets,
+				}),
 			},
 		},
 		settings:         settings,
@@ -1226,6 +1271,22 @@ func NewLeaseManager(
 	lm.descUpdateCh = make(chan catalog.Descriptor)
 	lm.descDelCh = make(chan descpb.ID)
 	lm.rangefeedErrCh = make(chan error)
+	lm.bytesMonitor = mon.NewMonitor(mon.Options{
+		Name:       mon.MakeName("leased-descriptors"),
+		CurCount:   lm.storage.leasingMetrics.leaseCurBytesCount,
+		MaxHist:    lm.storage.leasingMetrics.leaseMaxBytesHist,
+		Res:        mon.MemoryResource,
+		Settings:   settings,
+		LongLiving: true,
+	})
+	lm.bytesMonitor.StartNoReserved(context.Background(), rootBytesMonitor)
+	lm.boundAccount = lm.bytesMonitor.MakeConcurrentBoundAccount()
+	// Add a stopper for the bound account that we are using to
+	// track memory usage.
+	lm.stopper.AddCloser(stop.CloserFn(func() {
+		lm.boundAccount.Close(ctx)
+		lm.bytesMonitor.Stop(ctx)
+	}))
 	return lm
 }
 
@@ -1517,7 +1578,10 @@ func (m *Manager) IsDraining() bool {
 // been done by the time this call returns. See the explanation in
 // pkg/server/drain.go for details.
 func (m *Manager) SetDraining(
-	ctx context.Context, drain bool, reporter func(int, redact.SafeString),
+	ctx context.Context,
+	drain bool,
+	reporter func(int, redact.SafeString),
+	assertOnLeakedDescriptor bool,
 ) {
 	m.draining.Store(drain)
 	if !drain {
@@ -1530,7 +1594,13 @@ func (m *Manager) SetDraining(
 		leases := func() []*storedLease {
 			t.mu.Lock()
 			defer t.mu.Unlock()
-			return t.removeInactiveVersions()
+			leasesToRelease := t.removeInactiveVersions(ctx)
+			// Ensure that all leases are released at this time.
+			if buildutil.CrdbTestBuild && assertOnLeakedDescriptor && len(t.mu.active.data) > 0 {
+				// Panic that a descriptor may have leaked.
+				panic(errors.AssertionFailedf("descriptor leak was detected for: %d (%s)", t.id, t.mu.active))
+			}
+			return leasesToRelease
 		}()
 		for _, l := range leases {
 			releaseLease(ctx, l, m)
@@ -2478,3 +2548,8 @@ func (m *Manager) deleteOrphanedLeasesWithSameInstanceID(
 		}
 	}
 }
+
+// TestingGetBoundAccount returns the bound account used by the lease manager.
+func (m *Manager) TestingGetBoundAccount() *mon.ConcurrentBoundAccount {
+	return m.boundAccount
+}

pkg/sql/catalog/lease/lease_test.go

@@ -271,6 +271,7 @@ func (t *leaseTest) node(nodeID uint32) *lease.Manager {
 		)
 		cfgCpy.SQLLiveness.Start(context.Background(), nil)
 		mgr = lease.NewLeaseManager(
+			context.Background(),
 			ambientCtx,
 			nc,
 			cfgCpy.InternalDB,
@@ -282,6 +283,7 @@ func (t *leaseTest) node(nodeID uint32) *lease.Manager {
 			t.leaseManagerTestingKnobs,
 			t.server.AppStopper(),
 			cfgCpy.RangeFeedFactory,
+			cfgCpy.RootMemoryMonitor,
 		)
 		ctx := logtags.AddTag(context.Background(), "leasemgr", nodeID)
 		mgr.RunBackgroundLeasingTask(ctx)
@@ -520,8 +522,8 @@ func TestLeaseManagerDrain(testingT *testing.T) {
 		// starts draining.
 		l1RemovalTracker := leaseRemovalTracker.TrackRemoval(l1.Underlying())
 
-		t.node(1).SetDraining(ctx, true, nil /* reporter */)
-		t.node(2).SetDraining(ctx, true, nil /* reporter */)
+		t.node(1).SetDraining(ctx, true, nil /* reporter */, false)
+		t.node(2).SetDraining(ctx, true, nil /* reporter */, false)
 
 		// Leases cannot be acquired when in draining mode.
 		if _, err := t.acquire(1, descID); !testutils.IsError(err, "cannot acquire lease when draining") {
@@ -544,7 +546,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.node(1).SetDraining(ctx, false, nil /* reporter */)
+		t.node(1).SetDraining(ctx, false, nil /* reporter */, false /* assertOnLeakedDescriptor */)
 		l1 := t.mustAcquire(1, descID)
 		t.mustRelease(1, l1, nil)
 		t.expectLeases(descID, "/1/1")
@@ -3635,3 +3637,36 @@ func BenchmarkAcquireLeaseConcurrent(b *testing.B) {
 	}
 
 }
+
+// TestLeaseManagerIsMemoryMonitored basic sanity test to confirm memory monitoring
+// is working.
+func TestLeaseManagerIsMemoryMonitored(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	// This test creates a large number of objects, which
+	// can timeout under stress / race.
+	skip.UnderDuress(t)
+
+	ctx := context.Background()
+	s, sqlDB, _ := serverutils.StartServer(t, base.TestServerArgs{})
+	defer s.Stopper().Stop(ctx)
+	lm := s.LeaseManager().(*lease.Manager)
+	startBytes := lm.TestingGetBoundAccount().Used()
+	lastBytes := startBytes
+	runner := sqlutils.MakeSQLRunner(sqlDB)
+	// First, acquire leases on all the tables
+	for i := 0; i < 100; i++ {
+		runner.Exec(t, fmt.Sprintf("CREATE TABLE t%d(n int)", i))
+		runner.Exec(t, fmt.Sprintf("INSERT INTO t%d VALUES (1)", i))
+		currentBytes := lm.TestingGetBoundAccount().Used()
+		require.Greaterf(t, currentBytes, lastBytes, "memory usage should increase after using a table")
+		lastBytes = currentBytes
+	}
+	// Next we will release the leases on all of them.
+	lastBytes = lm.TestingGetBoundAccount().Used()
+	for i := 0; i < 100; i++ {
+		runner.Exec(t, fmt.Sprintf("DROP TABLE t%d", i))
+	}
+	currentBytes := lm.TestingGetBoundAccount().Used()
+	require.Lessf(t, currentBytes, lastBytes, "memory usage should be decreasing after dropping a table")
+}

pkg/sql/catalog/lease/storage.go

@@ -72,6 +72,8 @@ type leasingMetrics struct {
 	longWaitForNoVersionsActive                *metric.Gauge
 	longTwoVersionInvariantViolationWaitActive *metric.Gauge
 	longWaitForInitialVersionActive            *metric.Gauge
+	leaseMaxBytesHist                          metric.IHistogram
+	leaseCurBytesCount                         *metric.Gauge
 }
 
 type leaseFields struct {