MM 65084 server-side (#33861) (#34006) (#34044)

* MM 65084 server-side (#33861) (#34006)

Automatic Merge

* Add ConsumeOnce method to store layers

---------

Co-authored-by: Mattermost Build <build@mattermost.com>

server/channels/store/retrylayer/retrylayer.go

@@ -14138,6 +14138,27 @@ func (s *RetryLayerTokenStore) Cleanup(expiryTime int64) {
 
 }
 
+func (s *RetryLayerTokenStore) ConsumeOnce(tokenStr string) (*model.Token, error) {
+
+	tries := 0
+	for {
+		result, err := s.TokenStore.ConsumeOnce(tokenStr)
+		if err == nil {
+			return result, nil
+		}
+		if !isRepeatableError(err) {
+			return result, err
+		}
+		tries++
+		if tries >= 3 {
+			err = errors.Wrap(err, "giving up after 3 consecutive repeatable transaction failures")
+			return result, err
+		}
+		timepkg.Sleep(100 * timepkg.Millisecond)
+	}
+
+}
+
 func (s *RetryLayerTokenStore) Delete(token string) error {
 
 	tries := 0

server/channels/store/sqlstore/tokens_store.go

@@ -78,6 +78,21 @@ func (s SqlTokenStore) GetByToken(tokenString string) (*model.Token, error) {
 	return &token, nil
 }
 
+func (s SqlTokenStore) ConsumeOnce(tokenStr string) (*model.Token, error) {
+	var token model.Token
+
+	query := `DELETE FROM Tokens WHERE Token = ? RETURNING *`
+
+	if err := s.GetMaster().Get(&token, query, tokenStr); err != nil {
+		if err == sql.ErrNoRows {
+			return nil, store.NewErrNotFound("Token", tokenStr)
+		}
+		return nil, errors.Wrapf(err, "failed to consume token")
+	}
+
+	return &token, nil
+}
+
 func (s SqlTokenStore) Cleanup(expiryTime int64) {
 	if _, err := s.GetMaster().Exec("DELETE FROM Tokens WHERE CreateAt < ?", expiryTime); err != nil {
 		mlog.Error("Unable to cleanup token store.")

server/channels/store/store.go

@@ -693,6 +693,7 @@ type TokenStore interface {
 	Save(recovery *model.Token) error
 	Delete(token string) error
 	GetByToken(token string) (*model.Token, error)
+	ConsumeOnce(tokenStr string) (*model.Token, error)
 	Cleanup(expiryTime int64)
 	GetAllTokensByType(tokenType string) ([]*model.Token, error)
 	RemoveAllTokensByType(tokenType string) error

server/channels/store/storetest/mocks/TokenStore.go

@@ -19,6 +19,36 @@ func (_m *TokenStore) Cleanup(expiryTime int64) {
 	_m.Called(expiryTime)
 }
 
+// ConsumeOnce provides a mock function with given fields: tokenStr
+func (_m *TokenStore) ConsumeOnce(tokenStr string) (*model.Token, error) {
+	ret := _m.Called(tokenStr)
+
+	if len(ret) == 0 {
+		panic("no return value specified for ConsumeOnce")
+	}
+
+	var r0 *model.Token
+	var r1 error
+	if rf, ok := ret.Get(0).(func(string) (*model.Token, error)); ok {
+		return rf(tokenStr)
+	}
+	if rf, ok := ret.Get(0).(func(string) *model.Token); ok {
+		r0 = rf(tokenStr)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*model.Token)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(string) error); ok {
+		r1 = rf(tokenStr)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
 // Delete provides a mock function with given fields: token
 func (_m *TokenStore) Delete(token string) error {
 	ret := _m.Called(token)

server/channels/store/timerlayer/timerlayer.go

@@ -11106,6 +11106,22 @@ func (s *TimerLayerTokenStore) Cleanup(expiryTime int64) {
 	}
 }
 
+func (s *TimerLayerTokenStore) ConsumeOnce(tokenStr string) (*model.Token, error) {
+	start := time.Now()
+
+	result, err := s.TokenStore.ConsumeOnce(tokenStr)
+
+	elapsed := float64(time.Since(start)) / float64(time.Second)
+	if s.Root.Metrics != nil {
+		success := "false"
+		if err == nil {
+			success = "true"
+		}
+		s.Root.Metrics.ObserveStoreMethodDuration("TokenStore.ConsumeOnce", success, elapsed)
+	}
+	return result, err
+}
+
 func (s *TimerLayerTokenStore) Delete(token string) error {
 	start := time.Now()