[management] reduce calls to update last usage for users and pats (#8290)

management/server/store/sql_store_personal_access_token.go

@@ -12,7 +12,6 @@ import (
 	"gorm.io/gorm/clause"
 
 	"github.com/netbirdio/netbird/management/server/types"
-	"github.com/netbirdio/netbird/management/server/util"
 	"github.com/netbirdio/netbird/shared/management/status"
 )
 
@@ -130,22 +129,29 @@ func (s *SqlStore) GetUserPATs(ctx context.Context, lockStrength LockingStrength
 	return pats, nil
 }
 
-// MarkPATUsed marks a personal access token as used.
+// MarkPATUsed marks a personal access token as used. The write is skipped when LastUsed is newer than
+// types.PATLastUsedUpdateInterval, which also turns concurrent calls for the same token into no-ops. It returns a
+// NotFound error when the token doesn't exist.
 func (s *SqlStore) MarkPATUsed(ctx context.Context, patID string) error {
-	patCopy := types.PersonalAccessToken{
-		LastUsed: util.ToPtr(time.Now().UTC()),
-	}
-
-	fieldsToUpdate := []string{"last_used"}
-	result := s.db.Select(fieldsToUpdate).
-		Where(idQueryCondition, patID).Updates(&patCopy)
+	now := time.Now().UTC()
+	result := s.db.Model(&types.PersonalAccessToken{}).
+		Where(idQueryCondition, patID).
+		Where("last_used IS NULL OR last_used < ?", now.Add(-types.PATLastUsedUpdateInterval)).
+		Update("last_used", now)
 	if result.Error != nil {
 		log.WithContext(ctx).Errorf("failed to mark pat as used: %s", result.Error)
 		return status.Errorf(status.Internal, "failed to mark pat as used")
 	}
 
 	if result.RowsAffected == 0 {
-		return status.NewPATNotFoundError(patID)
+		var count int64
+		if err := s.db.Model(&types.PersonalAccessToken{}).Where(idQueryCondition, patID).Count(&count).Error; err != nil {
+			log.WithContext(ctx).Errorf("failed to check pat existence: %s", err)
+			return status.Errorf(status.Internal, "failed to mark pat as used")
+		}
+		if count == 0 {
+			return status.NewPATNotFoundError(patID)
+		}
 	}
 
 	return nil

management/server/store/sql_store_personal_access_token_test.go

@@ -136,6 +136,25 @@ func TestSqlStore_MarkPATUsed(t *testing.T) {
 	require.NoError(t, err)
 	now := time.Now().UTC()
 	require.WithinRange(t, pat.LastUsed.UTC(), now.Add(-15*time.Second), now, "LastUsed should be within 1 second of now")
+
+	recent := now.Add(-types.PATLastUsedUpdateInterval / 2).Truncate(time.Second)
+	require.NoError(t, store.(*SqlStore).db.Model(&types.PersonalAccessToken{}).Where("id = ?", patID).Update("last_used", recent).Error)
+	require.NoError(t, store.MarkPATUsed(context.Background(), patID))
+	pat, err = store.GetPATByID(context.Background(), LockingStrengthNone, userID, patID)
+	require.NoError(t, err)
+	assert.True(t, recent.Equal(pat.LastUsed.UTC()), "LastUsed within the update interval must not be rewritten")
+
+	stale := now.Add(-2 * types.PATLastUsedUpdateInterval)
+	require.NoError(t, store.(*SqlStore).db.Model(&types.PersonalAccessToken{}).Where("id = ?", patID).Update("last_used", stale).Error)
+	require.NoError(t, store.MarkPATUsed(context.Background(), patID))
+	pat, err = store.GetPATByID(context.Background(), LockingStrengthNone, userID, patID)
+	require.NoError(t, err)
+	assert.True(t, pat.LastUsed.UTC().After(stale), "LastUsed older than the update interval must be rewritten")
+
+	err = store.MarkPATUsed(context.Background(), "unknown-pat")
+	sErr, ok := status.FromError(err)
+	require.True(t, ok, "unknown pat must return a status error")
+	assert.Equal(t, status.NotFound, sErr.Type())
 }
 
 func TestSqlStore_SavePAT(t *testing.T) {

management/server/store/sql_store_user.go

@@ -252,20 +252,32 @@ func (s *SqlStore) GetAccountIDByUserID(ctx context.Context, lockStrength Lockin
 	return accountID, nil
 }
 
-// SaveUserLastLogin stores the last login time for a user in DB.
+// SaveUserLastLogin stores the last login time for a user in DB. It only writes when the stored value is unset or
+// older than lastLogin, so repeating the same login doesn't rewrite the row and an older login never overwrites a newer one.
+// It returns a NotFound error when the user doesn't exist.
 func (s *SqlStore) SaveUserLastLogin(ctx context.Context, accountID, userID string, lastLogin time.Time) error {
-	var user types.User
-	result := s.db.Take(&user, accountAndIDQueryCondition, accountID, userID)
-	if result.Error != nil {
-		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
-			return status.NewUserNotFoundError(userID)
-		}
-		return status.NewGetUserFromStoreError()
+	if lastLogin.IsZero() {
+		return nil
 	}
 
-	if !lastLogin.IsZero() {
-		user.LastLogin = &lastLogin
-		return s.db.Save(&user).Error
+	result := s.db.Model(&types.User{}).
+		Where(accountAndIDQueryCondition, accountID, userID).
+		Where("last_login IS NULL OR last_login < ?", lastLogin).
+		Update("last_login", lastLogin)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to save user last login: %s", result.Error)
+		return status.Errorf(status.Internal, "failed to save user last login")
+	}
+
+	if result.RowsAffected == 0 {
+		var count int64
+		if err := s.db.Model(&types.User{}).Where(accountAndIDQueryCondition, accountID, userID).Count(&count).Error; err != nil {
+			log.WithContext(ctx).Errorf("failed to check user existence: %s", err)
+			return status.NewGetUserFromStoreError()
+		}
+		if count == 0 {
+			return status.NewUserNotFoundError(userID)
+		}
 	}
 
 	return nil

management/server/store/sql_store_user_test.go

@@ -341,3 +341,50 @@ func TestSqlStore_SaveUsers_LargeBatch(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, 8002, len(accountUsers))
 }
+
+func TestSqlStore_SaveUserLastLogin(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	ctx := context.Background()
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	userID := "f4f6d672-63fb-11ec-90d6-0242ac120003"
+
+	before, err := store.GetUserByUserID(ctx, LockingStrengthNone, userID)
+	require.NoError(t, err)
+
+	lastLoginOf := func() time.Time {
+		user, err := store.GetUserByUserID(ctx, LockingStrengthNone, userID)
+		require.NoError(t, err)
+		return user.GetLastLogin().UTC()
+	}
+
+	first := time.Now().UTC().Truncate(time.Second)
+	require.NoError(t, store.SaveUserLastLogin(ctx, accountID, userID, first))
+	assert.True(t, first.Equal(lastLoginOf()), "first login must be stored")
+
+	require.NoError(t, store.SaveUserLastLogin(ctx, accountID, userID, first))
+	assert.True(t, first.Equal(lastLoginOf()), "repeating the same login must keep the stored value")
+
+	later := first.Add(30 * time.Second)
+	require.NoError(t, store.SaveUserLastLogin(ctx, accountID, userID, later))
+	assert.True(t, later.Equal(lastLoginOf()), "a newer login must be written")
+
+	require.NoError(t, store.SaveUserLastLogin(ctx, accountID, userID, first))
+	assert.True(t, later.Equal(lastLoginOf()), "an older login must not overwrite a newer one")
+
+	require.NoError(t, store.SaveUserLastLogin(ctx, accountID, userID, time.Time{}))
+	assert.True(t, later.Equal(lastLoginOf()), "zero login must be ignored")
+
+	err = store.SaveUserLastLogin(ctx, accountID, "unknown-user", later.Add(time.Hour))
+	sErr, ok := status.FromError(err)
+	require.True(t, ok, "unknown user must return a status error")
+	assert.Equal(t, status.NotFound, sErr.Type())
+
+	after, err := store.GetUserByUserID(ctx, LockingStrengthNone, userID)
+	require.NoError(t, err)
+	assert.Equal(t, before.Role, after.Role)
+	assert.Equal(t, before.AutoGroups, after.AutoGroups)
+	assert.Equal(t, before.Blocked, after.Blocked)
+}