Merge #150991 #151041

150991: pgwire: limit concurrent updates to last_login_time r=rafiss a=rafiss

###  pgwire: limit concurrent updates to last_login_time

This patch uses singleflight to ensure that each node only has one
request in flight for updating the last_login_time. If there already is
a request in flight, then that user is added to a set of pending users
to update, and their time gets populated in the next singleflight.

This will reduce write traffic to the system.users table.

This includes a minor refactor to separate the login_time updating logic
into a separate component.

As part of this, we also now skip updating the last_login_time if running in
a read-only tenant.

###  sql/pgwire: use a single query to update last_login_time

Rather than using a separate query for each user, now we will run one
query to update multiple users at once.

###  sql/pgwire: avoid frequent updates to last_login_time

This patch adds a small LRU cache to avoid updating the
estimated_last_login_time overly frequently for a single user. The cache
will cause us to skip the update if the user's login time was last
updated within the past 10 minutes.

This requires us to update the tests to use a different user, since the
existing test drops `testuser` and recreates, which means that the
last_login_time will not be updated for that user as it's still in the
LRU cache.

----

fixes #151051
fixes #150560

Release note: None

151041: security: fix memory accounting issue in client certificate cache r=angles-n-daemons a=angles-n-daemons

The client certificate cache is a simple utility cache which keeps track of the nearest certificate expiration times by user. It does this by reading in certificates as they are used, and then updating the gauge associated with the user so that our operators can alert on expiring certificates.

It has been identified that there is a memory accounting issue with this utility class, in that if we see the same certificate multiple times, we report that our cache is growing for each one. In reality however, though we may change the reference in our cache, the old value will be cleaned up, and memory should remain relatively low.

This PR both fixes the accounting bug, and adds a limit so that the client certificate cache can't negatively affect SQL operation.

Epic: none
Fixes: #151040
Release note: Fixes a memory accounting issue with the client certificate cache, where we were accidentally reporting multiple allocations for the same memory.

Co-authored-by: Rafi Shamim <[email protected]>
Co-authored-by: Brian Dillmann <[email protected]>

Author: 3 people

pkg/security/certificate_manager.go

@@ -13,6 +13,8 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/security/certnames"
 	"github.com/cockroachdb/cockroach/pkg/security/clientcert"
 	"github.com/cockroachdb/cockroach/pkg/security/username"
+	"github.com/cockroachdb/cockroach/pkg/settings"
+	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/log/eventpb"
 	"github.com/cockroachdb/cockroach/pkg/util/log/severity"
@@ -196,15 +198,24 @@ func (cm *CertificateManager) RegisterSignalHandler(
 	})
 }
 
+var CertCacheMemLimit = settings.RegisterIntSetting(
+	settings.SystemVisible,
+	"security.clienc_cert.cache_memory_limit",
+	"memory limit for the client certificate expiration cache",
+	1<<29, // 512MiB
+)
+
 // RegisterExpirationCache registers a cache for client certificate expiration.
 // It is called during server startup.
 func (cm *CertificateManager) RegisterExpirationCache(
 	ctx context.Context,
 	stopper *stop.Stopper,
 	timeSrc timeutil.TimeSource,
 	parentMon *mon.BytesMonitor,
+	st *cluster.Settings,
 ) error {
-	m := mon.NewMonitorInheritWithLimit(mon.MakeName("client-expiration-caches"), 0 /* limit */, parentMon, true /* longLiving */)
+	limit := CertCacheMemLimit.Get(&st.SV)
+	m := mon.NewMonitorInheritWithLimit(mon.MakeName("client-expiration-caches"), limit, parentMon, true /* longLiving */)
 	acc := m.MakeConcurrentBoundAccount()
 	m.StartNoReserved(ctx, parentMon)
 

pkg/security/clientcert/cert_expiry_cache.go

@@ -105,10 +105,13 @@ func (c *Cache) Upsert(ctx context.Context, user string, serial string, newExpir
 		c.cache[user] = map[string]certInfo{}
 	}
 
-	err := c.account.Grow(ctx, CertInfoSize)
-	if err != nil {
-		log.Warningf(ctx, "no memory available to cache cert expiry: %v", err)
-		c.evictLocked(ctx, user, serial)
+	// if the serial hasn't been seen, report it in the memory accounting.
+	if _, ok := c.cache[user][serial]; !ok {
+		err := c.account.Grow(ctx, CertInfoSize)
+		if err != nil {
+			log.Warningf(ctx, "no memory available to cache cert expiry: %v", err)
+			return
+		}
 	}
 
 	// insert / update the certificate expiration time.

pkg/security/clientcert/cert_expiry_cache_test.go

@@ -313,6 +313,7 @@ func TestAllocationTracking(t *testing.T) {
 		cache.Upsert(ctx, "user1", "serial2", 120)
 		cache.Upsert(ctx, "user2", "serial3", 65)
 		clock.Advance(time.Minute)
+		require.Equal(t, (3*certInfoSize)+(4*gaugeSize), account.Used())
 		cache.Purge(ctx)
 		cache.Upsert(ctx, "user2", "serial4", 75)
 		require.Equal(t, certInfoSize+(2*gaugeSize), account.Used())
@@ -324,9 +325,21 @@ func TestAllocationTracking(t *testing.T) {
 		cache.Upsert(ctx, "user1", "serial2", 120)
 		cache.Upsert(ctx, "user2", "serial3", 65)
 		cache.Upsert(ctx, "user2", "serial4", 75)
+		require.Equal(t, (4*certInfoSize)+(4*gaugeSize), account.Used())
 		cache.Clear()
 		require.Equal(t, int64(0), account.Used())
 	})
+
+	t.Run("overwriting an existing certificate does not change the reported memory allocation", func(t *testing.T) {
+		cache, account := newCacheAndAccount(ctx, clock, stopper)
+		require.Equal(t, int64(0), account.Used())
+		cache.Upsert(ctx, "user1", "serial1", 100)
+		require.Equal(t, certInfoSize+(2*gaugeSize), account.Used())
+		cache.Upsert(ctx, "user1", "serial1", 100)
+		require.Equal(t, certInfoSize+(2*gaugeSize), account.Used())
+		cache.Upsert(ctx, "user1", "serial1", 100)
+		require.Equal(t, certInfoSize+(2*gaugeSize), account.Used())
+	})
 }
 
 func TestExpiration(t *testing.T) {

pkg/server/server_sql.go

@@ -947,7 +947,7 @@ func newSQLServer(ctx context.Context, cfg sqlServerArgs) (*SQLServer, error) {
 			return nil, errors.Wrap(err, "initializing certificate manager")
 		}
 		err = certMgr.RegisterExpirationCache(
-			ctx, cfg.stopper, &timeutil.DefaultTimeSource{}, rootSQLMemoryMonitor,
+			ctx, cfg.stopper, &timeutil.DefaultTimeSource{}, rootSQLMemoryMonitor, cfg.Settings,
 		)
 		if err != nil {
 			return nil, errors.Wrap(err, "adding clientcert cache")

pkg/sql/logictest/testdata/logic_test/user

@@ -240,9 +240,9 @@ onlyif config local-mixed-25.2
 statement ok
 DROP user testuser
 
-onlyif config local-mixed-25.2
 user testuser nodeidx=0 newsession
 
+onlyif config local-mixed-25.2
 statement error pq: password authentication failed for user testuser
 SHOW session_user
 
@@ -259,7 +259,6 @@ skipif config local-mixed-25.2
 statement ok
 DROP user testuser
 
-skipif config local-mixed-25.2
 user testuser nodeidx=0 newsession
 
 # The logictest suite currently doesn't work for external authentication methods
@@ -271,38 +270,38 @@ skipif config local-mixed-25.2
 statement error pq: user identity unknown for identity provider
 SHOW session_user
 
-subtest end
-
-subtest validate_estimated_last_login_time
-
 user root
 
 statement ok
 CREATE user IF NOT EXISTS testuser
 
-user testuser
+subtest end
+
+subtest validate_estimated_last_login_time
+
+user testuser2
 
 # Since the logictest framework does not connect with the user until a command
 # is executed, we need to run a command before checking estimated_last_login_time.
 query T
 SHOW session_user
 ----
-testuser
+testuser2
 
 user root
 
 # The estimated_last_login_time is not guaranteed to be populated synchronously,
 # so we poll until testuser's entry was updated with the last login time.
 skipif config local-mixed-25.2
 query I retry
-SELECT count(*) FROM system.users WHERE estimated_last_login_time IS NOT NULL AND username = 'testuser'
+SELECT count(*) FROM system.users WHERE estimated_last_login_time IS NOT NULL AND username = 'testuser2'
 ----
 1
 
 user root
 
 onlyif config local-mixed-25.2
 statement error pq: column "estimated_last_login_time" does not exist
-select estimated_last_login_time from [SHOW USERS] where username = 'testuser';
+select estimated_last_login_time from [SHOW USERS] where username = 'testuser2';
 
 subtest end

pkg/sql/pgwire/BUILD.bazel

@@ -12,6 +12,7 @@ go_library(
         "conn.go",
         "hba_conf.go",
         "ident_map_conf.go",
+        "last_login_updater.go",
         "pre_serve.go",
         "pre_serve_options.go",
         "role_mapper.go",
@@ -63,6 +64,7 @@ go_library(
         "//pkg/sql/sqltelemetry",
         "//pkg/sql/types",
         "//pkg/util",
+        "//pkg/util/cache",
         "//pkg/util/ctxlog",
         "//pkg/util/duration",
         "//pkg/util/envutil",
@@ -81,6 +83,7 @@ go_library(
         "//pkg/util/ring",
         "//pkg/util/stop",
         "//pkg/util/syncutil",
+        "//pkg/util/syncutil/singleflight",
         "//pkg/util/system",
         "//pkg/util/timeofday",
         "//pkg/util/timeutil",

pkg/sql/pgwire/auth.go

@@ -90,14 +90,17 @@ type authOptions struct {
 // authentication and update c.sessionArgs with the authenticated user's name,
 // if different from the one given initially.
 func (c *conn) handleAuthentication(
-	ctx context.Context, ac AuthConn, authOpt authOptions, execCfg *sql.ExecutorConfig,
+	ctx context.Context, ac AuthConn, authOpt authOptions, server *Server,
 ) (connClose func(), _ error) {
 	if authOpt.testingSkipAuth {
 		return nil, nil
 	}
 	if authOpt.testingAuthHook != nil {
 		return nil, authOpt.testingAuthHook(ctx)
 	}
+	// Get execCfg from the server.
+	execCfg := server.execCfg
+
 	// To book-keep the authentication start time.
 	authStartTime := timeutil.Now()
 
@@ -278,28 +281,11 @@ func (c *conn) handleAuthentication(
 	duration := timeutil.Since(authStartTime).Nanoseconds()
 	c.publishConnLatencyMetric(duration, hbaEntry.Method.String())
 
-	c.populateLastLoginTime(ctx, execCfg, dbUser)
+	server.lastLoginUpdater.updateLastLoginTime(ctx, dbUser)
 
 	return connClose, nil
 }
 
-// populateLastLoginTime updates the last login time for the sql user
-// asynchronously. This not guaranteed to succeed and we log any errors obtained
-// from the update transaction to the DEV channel.
-func (c *conn) populateLastLoginTime(
-	ctx context.Context, execCfg *sql.ExecutorConfig, dbUser username.SQLUsername,
-) {
-	// Update last login time in async. This is done asynchronously to avoid
-	// blocking the connection.
-	if err := execCfg.Stopper.RunAsyncTask(ctx, "write_last_login_time", func(ctx context.Context) {
-		if err := sql.UpdateLastLoginTime(ctx, execCfg, dbUser.SQLIdentifier()); err != nil {
-			log.Warningf(ctx, "failed to update last login time for user %s: %v", dbUser, err)
-		}
-	}); err != nil {
-		log.Warningf(ctx, "failed to create async task to update last login time for user %s: %v", dbUser, err)
-	}
-}
-
 // publishConnLatencyMetric publishes the latency  of the connection
 // based on the authentication method.
 func (c *conn) publishConnLatencyMetric(duration int64, authMethod string) {

pkg/sql/pgwire/conn.go

@@ -151,7 +151,7 @@ func (c *conn) processCommands(
 	ctx context.Context,
 	authOpt authOptions,
 	ac AuthConn,
-	sqlServer *sql.Server,
+	server *Server,
 	reserved *mon.BoundAccount,
 	onDefaultIntSizeChange func(newSize int32),
 	sessionID clusterunique.ID,
@@ -175,7 +175,7 @@ func (c *conn) processCommands(
 		// network read on the connection's goroutine.
 		c.cancelConn()
 
-		pgwireKnobs := sqlServer.GetExecutorConfig().PGWireTestingKnobs
+		pgwireKnobs := server.SQLServer.GetExecutorConfig().PGWireTestingKnobs
 		if pgwireKnobs != nil && pgwireKnobs.CatchPanics {
 			if r := recover(); r != nil {
 				// Catch the panic and return it to the client as an error.
@@ -203,14 +203,14 @@ func (c *conn) processCommands(
 
 	// Authenticate the connection.
 	if connCloseAuthHandler, retErr = c.handleAuthentication(
-		ctx, ac, authOpt, sqlServer.GetExecutorConfig(),
+		ctx, ac, authOpt, server,
 	); retErr != nil {
 		// Auth failed or some other error.
 		return
 	}
 
 	var decrementConnectionCount func()
-	if decrementConnectionCount, retErr = sqlServer.IncrementConnectionCount(c.sessionArgs); retErr != nil {
+	if decrementConnectionCount, retErr = server.SQLServer.IncrementConnectionCount(c.sessionArgs); retErr != nil {
 		// This will return pgcode.TooManyConnections which is used by the sql proxy
 		// to skip failed auth throttle (as in this case the auth was fine but the
 		// error occurred before sending back auth ok msg)
@@ -224,7 +224,7 @@ func (c *conn) processCommands(
 	}
 
 	// Inform the client of the default session settings.
-	connHandler, retErr = c.sendInitialConnData(ctx, sqlServer, onDefaultIntSizeChange, sessionID)
+	connHandler, retErr = c.sendInitialConnData(ctx, server.SQLServer, onDefaultIntSizeChange, sessionID)
 	if retErr != nil {
 		return
 	}
@@ -250,7 +250,7 @@ func (c *conn) processCommands(
 
 	// Now actually process commands.
 	reservedOwned = false // We're about to pass ownership away.
-	retErr = sqlServer.ServeConn(
+	retErr = server.SQLServer.ServeConn(
 		ctx,
 		connHandler,
 		reserved,