From b6b561a8f14c9a7a4f7f8d25df6356b6fe7492fc Mon Sep 17 00:00:00 2001
From: Ben Schumacher <ben.schumacher@mattermost.com>
Date: Mon, 5 Jun 2023 10:32:23 +0200
Subject: [PATCH] [MM-52973] Avoid thundering herd problem in
 IsFirstUserAccount (#23549)

---
 server/channels/app/platform/config.go      | 36 ++++++++-----
 server/channels/app/platform/config_test.go | 59 +++++++++++++++++++--
 server/channels/app/platform/service.go     | 14 ++---
 3 files changed, 86 insertions(+), 23 deletions(-)

diff --git a/server/channels/app/platform/config.go b/server/channels/app/platform/config.go
index c6aeb6afdf..eb656285eb 100644
--- a/server/channels/app/platform/config.go
+++ b/server/channels/app/platform/config.go
@@ -323,21 +323,31 @@ func (ps *PlatformService) LimitedClientConfig() map[string]string {
 }
 
 func (ps *PlatformService) IsFirstUserAccount() bool {
-	if ps.fetchUserCountForFirstUserAccountCheck.Load() {
-		ps.logger.Debug("Fetching user count for first user account check")
-		count, err := ps.Store.User().Count(model.UserCountOptions{IncludeDeleted: true})
-		if err != nil {
-			return false
-		}
-		// Avoid calling the user count query in future if we get a count > 0
-		if count > 0 {
-			ps.fetchUserCountForFirstUserAccountCheck.Store(false)
-			return false
-		}
-		return true
+	if !ps.isFirstUserAccount.Load() {
+		return false
 	}
 
-	return false
+	ps.isFirstUserAccountLock.Lock()
+	defer ps.isFirstUserAccountLock.Unlock()
+	// Retry under lock as another call might have already succeeded.
+	if !ps.isFirstUserAccount.Load() {
+		return false
+	}
+
+	ps.logger.Debug("Fetching user count for first user account check")
+	count, err := ps.Store.User().Count(model.UserCountOptions{IncludeDeleted: true})
+	if err != nil {
+		return false
+	}
+
+	// Avoid calling the user count query in future if we get a count > 0
+	if count > 0 {
+		ps.isFirstUserAccount.Store(false)
+		return false
+	}
+
+	return true
+
 }
 
 func (ps *PlatformService) MaxPostSize() int {
diff --git a/server/channels/app/platform/config_test.go b/server/channels/app/platform/config_test.go
index f982ef794d..76c9164ba4 100644
--- a/server/channels/app/platform/config_test.go
+++ b/server/channels/app/platform/config_test.go
@@ -5,6 +5,7 @@ package platform
 
 import (
 	"errors"
+	"sync"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -131,12 +132,62 @@ func TestIsFirstUserAccount(t *testing.T) {
 		t.Run(te.name, func(t *testing.T) {
 			*userStoreMock = smocks.UserStore{}
 
-			userStoreMock.On("Count", model.UserCountOptions{IncludeDeleted: true}).Return(te.count, te.err).RunFn = func(args mock.Arguments) {
-				if !te.shouldCallStore {
-					assert.Fail(t, "should not have called the store")
-				}
+			if te.shouldCallStore {
+				userStoreMock.On("Count", model.UserCountOptions{IncludeDeleted: true}).Return(te.count, te.err).Once()
+			} else {
+				userStoreMock.On("Count", model.UserCountOptions{IncludeDeleted: true}).Unset()
 			}
+
 			require.Equal(t, te.result, th.Service.IsFirstUserAccount())
 		})
 	}
 }
+
+func TestIsFirstUserAccountThunderingHerd(t *testing.T) {
+	th := SetupWithStoreMock(t)
+	defer th.TearDown()
+	storeMock := th.Service.Store.(*smocks.Store)
+	userStoreMock := &smocks.UserStore{}
+	storeMock.On("User").Return(userStoreMock)
+
+	tests := []struct {
+		name               string
+		count              int64
+		err                error
+		concurrentRequest  int
+		result             bool
+		numberOfStoreCalls int
+	}{
+		{"failed request", 0, errors.New("error"), 10, false, 10},
+		{"success negative users", -100, nil, 10, true, 10},
+		{"success no users", 0, nil, 10, true, 10},
+		{"success one user - lot of requests", 1, nil, 1000, false, 1},
+		{"success multiple users - no store call", 42, nil, 10, false, 0},
+	}
+
+	for _, te := range tests {
+		t.Run(te.name, func(t *testing.T) {
+			*userStoreMock = smocks.UserStore{}
+
+			if te.numberOfStoreCalls != 0 {
+				userStoreMock.On("Count", model.UserCountOptions{IncludeDeleted: true}).Return(te.count, te.err).Times(te.numberOfStoreCalls)
+			} else {
+				userStoreMock.On("Count", model.UserCountOptions{IncludeDeleted: true}).Unset()
+			}
+			defer userStoreMock.AssertExpectations(t)
+
+			var wg sync.WaitGroup
+			for i := 0; i < te.concurrentRequest; i++ {
+				wg.Add(1)
+
+				go func() {
+					defer wg.Done()
+					require.Equal(t, te.result, th.Service.IsFirstUserAccount())
+				}()
+			}
+
+			wg.Wait()
+		})
+	}
+
+}
diff --git a/server/channels/app/platform/service.go b/server/channels/app/platform/service.go
index 7315833e0f..5e0a99de52 100644
--- a/server/channels/app/platform/service.go
+++ b/server/channels/app/platform/service.go
@@ -50,11 +50,13 @@ type PlatformService struct {
 	sessionCache  cache.Cache
 	sessionPool   sync.Pool
 
-	asymmetricSigningKey                   atomic.Pointer[ecdsa.PrivateKey]
-	clientConfig                           atomic.Value
-	clientConfigHash                       atomic.Value
-	limitedClientConfig                    atomic.Value
-	fetchUserCountForFirstUserAccountCheck atomic.Bool
+	asymmetricSigningKey atomic.Pointer[ecdsa.PrivateKey]
+	clientConfig         atomic.Value
+	clientConfigHash     atomic.Value
+	limitedClientConfig  atomic.Value
+
+	isFirstUserAccountLock sync.Mutex
+	isFirstUserAccount     atomic.Bool
 
 	logger              *mlog.Logger
 	notificationsLogger *mlog.Logger
@@ -129,7 +131,7 @@ func New(sc ServiceConfig, options ...Option) (*PlatformService, error) {
 	}
 
 	// Assume the first user account has not been created yet. A call to the DB will later check if this is really the case.
-	ps.fetchUserCountForFirstUserAccountCheck.Store(true)
+	ps.isFirstUserAccount.Store(true)
 
 	// Step 1: Cache provider.
 	// At the moment we only have this implementation