|
| 1 | +// Copyright 2025 The Cockroach Authors. |
| 2 | +// |
| 3 | +// Use of this software is governed by the CockroachDB Software License |
| 4 | +// included in the /LICENSE file. |
| 5 | + |
| 6 | +package pgwire |
| 7 | + |
| 8 | +import ( |
| 9 | + "context" |
| 10 | + |
| 11 | + "github.com/cockroachdb/cockroach/pkg/security/username" |
| 12 | + "github.com/cockroachdb/cockroach/pkg/sql" |
| 13 | + "github.com/cockroachdb/cockroach/pkg/util/log" |
| 14 | + "github.com/cockroachdb/cockroach/pkg/util/syncutil" |
| 15 | + "github.com/cockroachdb/cockroach/pkg/util/syncutil/singleflight" |
| 16 | +) |
| 17 | + |
| 18 | +// lastLoginUpdater handles updating the last login time for SQL users with |
| 19 | +// deduplication via singleflight to reduce concurrent updates. |
| 20 | +type lastLoginUpdater struct { |
| 21 | + // group ensures that there is at most one last login time update |
| 22 | + // in-flight at any given time. |
| 23 | + group *singleflight.Group |
| 24 | + // execCfg is the executor configuration used for running update operations. |
| 25 | + execCfg *sql.ExecutorConfig |
| 26 | + |
| 27 | + mu struct { |
| 28 | + syncutil.Mutex |
| 29 | + // pendingUsers tracks users that need their last login time updated |
| 30 | + // in the next singleflight operation. |
| 31 | + pendingUsers map[username.SQLUsername]struct{} |
| 32 | + } |
| 33 | +} |
| 34 | + |
| 35 | +// newLastLoginUpdater creates a new lastLoginUpdater instance. |
| 36 | +func newLastLoginUpdater(execCfg *sql.ExecutorConfig) *lastLoginUpdater { |
| 37 | + u := &lastLoginUpdater{ |
| 38 | + group: singleflight.NewGroup("update last login time", ""), |
| 39 | + execCfg: execCfg, |
| 40 | + } |
| 41 | + u.mu.pendingUsers = make(map[username.SQLUsername]struct{}) |
| 42 | + return u |
| 43 | +} |
| 44 | + |
| 45 | +// updateLastLoginTime updates the last login time for the SQL user |
| 46 | +// asynchronously. This is not guaranteed to succeed and we log any errors |
| 47 | +// obtained from the update transaction to the DEV channel. |
| 48 | +func (u *lastLoginUpdater) updateLastLoginTime(ctx context.Context, dbUser username.SQLUsername) { |
| 49 | + // Avoid updating if we're in a read-only tenant. |
| 50 | + if u.execCfg.TenantReadOnly { |
| 51 | + return |
| 52 | + } |
| 53 | + |
| 54 | + // Use singleflight to ensure at most one last login time update batch |
| 55 | + // is in-flight at any given time. |
| 56 | + future, leader := u.group.DoChan(ctx, "UpdateLastLoginTime", |
| 57 | + singleflight.DoOpts{ |
| 58 | + Stop: u.execCfg.Stopper, |
| 59 | + InheritCancelation: false, |
| 60 | + }, |
| 61 | + func(ctx context.Context) (interface{}, error) { |
| 62 | + // The leader adds itself to pending users, and processes all |
| 63 | + // pending updates. |
| 64 | + u.mu.Lock() |
| 65 | + u.mu.pendingUsers[dbUser] = struct{}{} |
| 66 | + u.mu.Unlock() |
| 67 | + return nil, u.processPendingUpdates(ctx) |
| 68 | + }) |
| 69 | + |
| 70 | + // Add this user to the pending set if not the leader, |
| 71 | + if !leader { |
| 72 | + u.mu.Lock() |
| 73 | + u.mu.pendingUsers[dbUser] = struct{}{} |
| 74 | + u.mu.Unlock() |
| 75 | + } else { |
| 76 | + // Leader waits for the result in an async task to avoid blocking authentication |
| 77 | + if err := u.execCfg.Stopper.RunAsyncTask(ctx, "wait_last_login_update", func(ctx context.Context) { |
| 78 | + result := future.WaitForResult(ctx) |
| 79 | + if result.Err != nil { |
| 80 | + log.Warningf(ctx, "could not update last login times: %v", result.Err) |
| 81 | + } |
| 82 | + }); err != nil { |
| 83 | + log.Warningf(ctx, "could not create async task to wait for last login update: %v", err) |
| 84 | + } |
| 85 | + } |
| 86 | +} |
| 87 | + |
| 88 | +// processPendingUpdates processes all users in the pending set and clears it. |
| 89 | +func (u *lastLoginUpdater) processPendingUpdates(ctx context.Context) error { |
| 90 | + u.mu.Lock() |
| 91 | + users := make([]username.SQLUsername, 0, len(u.mu.pendingUsers)) |
| 92 | + for user := range u.mu.pendingUsers { |
| 93 | + users = append(users, user) |
| 94 | + } |
| 95 | + // Clear the pending users set. |
| 96 | + u.mu.pendingUsers = make(map[username.SQLUsername]struct{}) |
| 97 | + u.mu.Unlock() |
| 98 | + |
| 99 | + // Update last login time for all pending users. |
| 100 | + for _, user := range users { |
| 101 | + if err := sql.UpdateLastLoginTime(ctx, u.execCfg, user.SQLIdentifier()); err != nil { |
| 102 | + return err |
| 103 | + } |
| 104 | + } |
| 105 | + return nil |
| 106 | +} |
0 commit comments