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

management/server/http/middleware/auth_middleware.go

@@ -185,9 +185,10 @@ func (m *AuthMiddleware) checkPATFromRequest(r *http.Request, authHeaderParts []
 		return fmt.Errorf("token expired")
 	}
 
-	err = m.authManager.MarkPATUsed(ctx, pat.ID)
-	if err != nil {
-		return err
+	if pat.LastUsed == nil || time.Since(*pat.LastUsed) >= types.PATLastUsedUpdateInterval {
+		if err := m.authManager.MarkPATUsed(ctx, pat.ID); err != nil {
+			return err
+		}
 	}
 
 	userAuth := auth.UserAuth{

management/server/http/middleware/auth_middleware_test.go

@@ -779,3 +779,66 @@ func TestAuthMiddleware_Handler_Child(t *testing.T) {
 		})
 	}
 }
+
+func TestAuthMiddleware_MarkPATUsedThrottled(t *testing.T) {
+	tt := []struct {
+		name       string
+		lastUsed   *time.Time
+		expectMark bool
+	}{
+		{name: "never used", lastUsed: nil, expectMark: true},
+		{name: "used just now", lastUsed: util.ToPtr(time.Now().UTC()), expectMark: false},
+		{name: "used before the update interval", lastUsed: util.ToPtr(time.Now().UTC().Add(-2 * types.PATLastUsedUpdateInterval)), expectMark: true},
+	}
+
+	for _, tc := range tt {
+		t.Run(tc.name, func(t *testing.T) {
+			pat := &types.PersonalAccessToken{
+				ID:             tokenID,
+				ExpirationDate: util.ToPtr(time.Now().UTC().AddDate(0, 0, 7)),
+				LastUsed:       tc.lastUsed,
+			}
+
+			marked := 0
+			mockAuth := &auth.MockManager{
+				GetPATInfoFunc: func(_ context.Context, _ string) (*types.User, *types.PersonalAccessToken, string, string, error) {
+					return testAccount.Users[userID], pat, testAccount.Domain, testAccount.DomainCategory, nil
+				},
+				MarkPATUsedFunc: func(_ context.Context, _ string) error {
+					marked++
+					return nil
+				},
+			}
+
+			disabledLimiter := ratelimit.NewAPIRateLimiter(nil)
+			disabledLimiter.SetEnabled(false)
+			authMiddleware := NewAuthMiddleware(
+				mockAuth,
+				func(ctx context.Context, userAuth nbauth.UserAuth) (string, string, error) {
+					return userAuth.AccountId, userAuth.UserId, nil
+				},
+				func(ctx context.Context, userAuth nbauth.UserAuth) error {
+					return nil
+				},
+				func(ctx context.Context, userAuth nbauth.UserAuth) (*types.User, error) {
+					return &types.User{}, nil
+				},
+				disabledLimiter,
+				nil,
+			)
+
+			handler := authMiddleware.Handler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
+			req := httptest.NewRequest("GET", "http://testing/test", nil)
+			req.Header.Set("Authorization", "Token "+PAT)
+			rec := httptest.NewRecorder()
+			handler.ServeHTTP(rec, req)
+
+			assert.Equal(t, http.StatusOK, rec.Code)
+			if tc.expectMark {
+				assert.Equal(t, 1, marked, "MarkPATUsed must be called")
+			} else {
+				assert.Equal(t, 0, marked, "MarkPATUsed must be skipped")
+			}
+		})
+	}
+}

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

management/server/types/personal_access_token.go

@@ -23,6 +23,8 @@ const (
 	PATChecksumLength = 6
 	// PATLength total number of characters used for the token
 	PATLength = 40
+	// PATLastUsedUpdateInterval is the minimum age of a token's LastUsed before it is written again
+	PATLastUsedUpdateInterval = time.Minute
 )
 
 // PersonalAccessToken holds all information about a PAT including a hashed version of it for verification

management/server/user.go

@@ -238,9 +238,10 @@ func (am *DefaultAccountManager) GetUserFromUserAuth(ctx context.Context, userAu
 	// server when user authenticates a device. And we need to separate the Dashboard login event from the Device login event.
 	newLogin := user.LastDashboardLoginChanged(userAuth.LastLogin)
 
-	err = am.Store.SaveUserLastLogin(ctx, userAuth.AccountId, userAuth.UserId, userAuth.LastLogin)
-	if err != nil {
-		log.WithContext(ctx).Debugf("failed to update user last login: %v", err)
+	if userAuth.LastLogin.After(user.GetLastLogin()) {
+		if err := am.Store.SaveUserLastLogin(ctx, userAuth.AccountId, userAuth.UserId, userAuth.LastLogin); err != nil {
+			log.WithContext(ctx).Debugf("failed to update user last login: %v", err)
+		}
 	}
 
 	if newLogin {