Fix race condition in telemetry (#18734)

* Fix race condition in telemetry

* Fix tests

api4/cors_test.go

@@ -127,11 +127,8 @@ func TestCORSRequestHandling(t *testing.T) {
 				*cfg.ServiceSettings.CorsAllowCredentials = testcase.CorsAllowCredentials
 			})
 			defer th.TearDown()
-			systemStore := mocks.SystemStore{}
-			systemStore.On("Get").Return(make(model.StringMap), nil)
 			licenseStore := mocks.LicenseStore{}
 			licenseStore.On("Get", "").Return(&model.LicenseRecord{}, nil)
-			th.App.Srv().Store.(*mocks.Store).On("System").Return(&systemStore)
 			th.App.Srv().Store.(*mocks.Store).On("License").Return(&licenseStore)
 
 			port := th.App.Srv().ListenAddr.Port

app/app_test.go

@@ -50,7 +50,6 @@ func TestUnitUpdateConfig(t *testing.T) {
 	mockSystemStore.On("GetByName", "UpgradedFromTE").Return(&model.System{Name: "UpgradedFromTE", Value: "false"}, nil)
 	mockSystemStore.On("GetByName", "InstallationDate").Return(&model.System{Name: "InstallationDate", Value: "10"}, nil)
 	mockSystemStore.On("GetByName", "FirstServerRunTimestamp").Return(&model.System{Name: "FirstServerRunTimestamp", Value: "10"}, nil)
-	mockSystemStore.On("Get").Return(make(model.StringMap), nil)
 	mockLicenseStore := mocks.LicenseStore{}
 	mockLicenseStore.On("Get", "").Return(&model.LicenseRecord{}, nil)
 	mockStore.On("User").Return(&mockUserStore)

services/telemetry/telemetry.go

@@ -122,20 +122,16 @@ func (ts *TelemetryService) ensureTelemetryID() {
 	if ts.TelemetryID != "" {
 		return
 	}
-	props, err := ts.dbStore.System().Get()
+
+	id := model.NewId()
+	systemID := &model.System{Name: model.SystemTelemetryId, Value: id}
+	systemID, err := ts.dbStore.System().InsertIfExists(systemID)
 	if err != nil {
 		mlog.Error("unable to get the telemetry ID", mlog.Err(err))
 		return
 	}
 
-	id := props[model.SystemTelemetryId]
+	ts.TelemetryID = systemID.Value
-	if id == "" {
-		id = model.NewId()
-		systemID := &model.System{Name: model.SystemTelemetryId, Value: id}
-		ts.dbStore.System().Save(systemID)
-	}
-
-	ts.TelemetryID = id
 }
 
 func (ts *TelemetryService) getRudderConfig() RudderConfig {

services/telemetry/telemetry_test.go

@@ -7,6 +7,7 @@ import (
 	"context"
 	"crypto/ecdsa"
 	"encoding/json"
+	"errors"
 	"io/ioutil"
 	"net/http"
 	"net/http/httptest"
@@ -81,9 +82,9 @@ func initializeMocks(cfg *model.Config) (*mocks.ServerIface, *storeMocks.Store,
 	storeMock.On("GetDbVersion", false).Return("5.24.0", nil)
 
 	systemStore := storeMocks.SystemStore{}
-	props := model.StringMap{}
+	systemStore.On("Get").Return(make(model.StringMap), nil)
-	props[model.SystemTelemetryId] = "test"
+	systemID := &model.System{Name: model.SystemTelemetryId, Value: "test"}
-	systemStore.On("Get").Return(props, nil)
+	systemStore.On("InsertIfExists", mock.Anything).Return(systemID, nil)
 	systemStore.On("GetByName", model.AdvancedPermissionsMigrationKey).Return(nil, nil)
 	systemStore.On("GetByName", model.MigrationKeyAdvancedPermissionsPhase2).Return(nil, nil)
 
@@ -159,6 +160,83 @@ func initializeMocks(cfg *model.Config) (*mocks.ServerIface, *storeMocks.Store,
 	}, cleanUp
 }
 
+func TestEnsureTelemetryID(t *testing.T) {
+	t.Run("test ID in database and does not run twice", func(t *testing.T) {
+		storeMock := &storeMocks.Store{}
+
+		systemStore := storeMocks.SystemStore{}
+		returnValue := &model.System{
+			Name:  model.SystemTelemetryId,
+			Value: "test",
+		}
+		systemStore.On("InsertIfExists", mock.AnythingOfType("*model.System")).Return(returnValue, nil).Once()
+
+		storeMock.On("System").Return(&systemStore)
+
+		serverIfaceMock := &mocks.ServerIface{}
+		cfg := &model.Config{}
+		cfg.SetDefaults()
+
+		testLogger, _ := mlog.NewLogger()
+
+		telemetryService := New(serverIfaceMock, storeMock, searchengine.NewBroker(cfg, nil), testLogger)
+		assert.Equal(t, "test", telemetryService.TelemetryID)
+
+		telemetryService.ensureTelemetryID()
+		assert.Equal(t, "test", telemetryService.TelemetryID)
+
+		// No more calls to the store if we try to ensure it again
+		telemetryService.ensureTelemetryID()
+		assert.Equal(t, "test", telemetryService.TelemetryID)
+	})
+
+	t.Run("new test ID created", func(t *testing.T) {
+		storeMock := &storeMocks.Store{}
+
+		systemStore := storeMocks.SystemStore{}
+		returnValue := &model.System{
+			Name: model.SystemTelemetryId,
+		}
+
+		var generatedID string
+		systemStore.On("InsertIfExists", mock.AnythingOfType("*model.System")).Return(returnValue, nil).Once().Run(func(args mock.Arguments) {
+			s := args.Get(0).(*model.System)
+			returnValue.Value = s.Value
+			generatedID = s.Value
+		})
+		storeMock.On("System").Return(&systemStore)
+
+		serverIfaceMock := &mocks.ServerIface{}
+		cfg := &model.Config{}
+		cfg.SetDefaults()
+
+		testLogger, _ := mlog.NewLogger()
+
+		telemetryService := New(serverIfaceMock, storeMock, searchengine.NewBroker(cfg, nil), testLogger)
+		assert.Equal(t, generatedID, telemetryService.TelemetryID)
+	})
+
+	t.Run("fail to save test ID", func(t *testing.T) {
+		storeMock := &storeMocks.Store{}
+
+		systemStore := storeMocks.SystemStore{}
+
+		insertError := errors.New("insert error")
+		systemStore.On("InsertIfExists", mock.AnythingOfType("*model.System")).Return(nil, insertError).Once()
+
+		storeMock.On("System").Return(&systemStore)
+
+		serverIfaceMock := &mocks.ServerIface{}
+		cfg := &model.Config{}
+		cfg.SetDefaults()
+
+		testLogger, _ := mlog.NewLogger()
+
+		telemetryService := New(serverIfaceMock, storeMock, searchengine.NewBroker(cfg, nil), testLogger)
+		assert.Equal(t, "", telemetryService.TelemetryID)
+	})
+}
+
 func TestPluginSetting(t *testing.T) {
 	settings := &model.PluginSettings{
 		Plugins: map[string]map[string]interface{}{

testlib/store.go

@@ -62,7 +62,7 @@ func GetMockStoreForSetupFunctions() *mocks.Store {
 	systemStore.On("GetByName", model.MigrationKeyAddIntegrationsSubsectionPermissions).Return(&model.System{Name: model.MigrationKeyAddIntegrationsSubsectionPermissions, Value: "true"}, nil)
 	systemStore.On("GetByName", model.MigrationKeyAddManageSharedChannelPermissions).Return(&model.System{Name: model.MigrationKeyAddManageSharedChannelPermissions, Value: "true"}, nil)
 	systemStore.On("GetByName", model.MigrationKeyAddManageSecureConnectionsPermissions).Return(&model.System{Name: model.MigrationKeyAddManageSecureConnectionsPermissions, Value: "true"}, nil)
-	systemStore.On("Get").Return(make(model.StringMap), nil)
+	systemStore.On("InsertIfExists", mock.AnythingOfType("*model.System")).Return(&model.System{}, nil).Once()
 	systemStore.On("Save", mock.AnythingOfType("*model.System")).Return(nil)
 
 	userStore := mocks.UserStore{}