[MM-63314] Fix ClaimJob in HA environments (#30383)

* ClaimJob now returns newly claimed job * internal code affected by change * test changes required * two branches: for mysql, use transaction; for postgres, use returning * two branches: for mysql, use transaction; for postgres, use returning * use same millis value for LastActivityAt and StartAt * blank commit --------- Co-authored-by: Mattermost Build <build@mattermost.com>
2025-03-14 10:24:26 -04:00
--- a/server/channels/store/retrylayer/retrylayer.go
+++ b/server/channels/store/retrylayer/retrylayer.go
@@ -6384,7 +6384,7 @@ func (s *RetryLayerJobStore) UpdateStatus(id string, status string) (*model.Job,

 }

-func (s *RetryLayerJobStore) UpdateStatusOptimistically(id string, currentStatus string, newStatus string) (bool, error) {
+func (s *RetryLayerJobStore) UpdateStatusOptimistically(id string, currentStatus string, newStatus string) (*model.Job, error) {

 	tries := 0
 	for {
--- a/server/channels/store/sqlstore/job_store.go
+++ b/server/channels/store/sqlstore/job_store.go
@@ -172,34 +172,86 @@ func (jss SqlJobStore) UpdateStatus(id string, status string) (*model.Job, error
 	return job, nil
 }

-func (jss SqlJobStore) UpdateStatusOptimistically(id string, currentStatus string, newStatus string) (bool, error) {
+func (jss SqlJobStore) UpdateStatusOptimistically(id string, currentStatus string, newStatus string) (*model.Job, error) {
+	lastActivityAndStartTime := model.GetMillis()
+
+	if jss.DriverName() == model.DatabaseDriverMysql {
+		tx, err := jss.GetMaster().Beginx()
+		if err != nil {
+			return nil, errors.Wrap(err, "begin_transaction")
+		}
+		defer finalizeTransactionX(tx, &err)
+
+		builder := jss.getQueryBuilder().
+			Update("Jobs").
+			Set("LastActivityAt", lastActivityAndStartTime).
+			Set("Status", newStatus).
+			Where(sq.Eq{"Id": id, "Status": currentStatus})
+
+		if newStatus == model.JobStatusInProgress {
+			builder = builder.Set("StartAt", lastActivityAndStartTime)
+		}
+
+		sqlResult, err := tx.ExecBuilder(builder)
+		if err != nil {
+			return nil, errors.Wrapf(err, "failed to update Job with id=%s", id)
+		}
+		rows, err := sqlResult.RowsAffected()
+		if err != nil {
+			return nil, errors.Wrap(err, "unable to get rows affected")
+		}
+		if rows != 1 {
+			return nil, nil
+		}
+
+		getBuilder := jss.getQueryBuilder().
+			Select("*").
+			From("Jobs").
+			Where(sq.Eq{"Id": id, "Status": newStatus})
+
+		var job model.Job
+		if err = tx.GetBuilder(&job, getBuilder); err != nil {
+			if errors.Is(err, sql.ErrNoRows) {
+				return nil, store.NewErrNotFound("Job", id)
+			}
+			return nil, errors.Wrapf(err, "failed to get Job with id=%s", id)
+		}
+
+		err = tx.Commit()
+		if err != nil {
+			return nil, errors.Wrap(err, "commit_transaction")
+		}
+
+		return &job, nil
+	}
+
+	// For PostgreSQL, use RETURNING to get the updated job in a single query
 	builder := jss.getQueryBuilder().
 		Update("Jobs").
-		Set("LastActivityAt", model.GetMillis()).
+		Set("LastActivityAt", lastActivityAndStartTime).
 		Set("Status", newStatus).
-		Where(sq.Eq{"Id": id, "Status": currentStatus})
+		Where(sq.Eq{"Id": id, "Status": currentStatus}).
+		Suffix("RETURNING *")

 	if newStatus == model.JobStatusInProgress {
-		builder = builder.Set("StartAt", model.GetMillis())
-	}
-	query, args, err := builder.ToSql()
-	if err != nil {
-		return false, errors.Wrap(err, "job_tosql")
+		builder = builder.Set("StartAt", lastActivityAndStartTime)
 	}

-	sqlResult, err := jss.GetMaster().Exec(query, args...)
-	if err != nil {
-		return false, errors.Wrapf(err, "failed to update Job with id=%s", id)
-	}
-	rows, err := sqlResult.RowsAffected()
-	if err != nil {
-		return false, errors.Wrap(err, "unable to get rows affected")
-	}
-	if rows != 1 {
-		return false, nil
+	var job []*model.Job
+	if err := jss.GetMaster().SelectBuilder(&job, builder); err != nil {
+		if errors.Is(err, sql.ErrNoRows) {
+			return nil, store.NewErrNotFound("Job", id)
+		}
+		return nil, errors.Wrapf(err, "failed to update Job with id=%s", id)
 	}

-	return true, nil
+	// we are updating by id, so we should only ever update 1 job
+	if len(job) != 1 {
+		// no row was updated, but no error above, so to remain consistent we return nil, nil
+		return nil, nil
+	}
+
+	return job[0], nil
 }

 func (jss SqlJobStore) Get(c request.CTX, id string) (*model.Job, error) {
@@ -213,7 +265,7 @@ func (jss SqlJobStore) Get(c request.CTX, id string) (*model.Job, error) {

 	var status model.Job
 	if err = jss.GetReplica().Get(&status, query, args...); err != nil {
-		if err == sql.ErrNoRows {
+		if errors.Is(err, sql.ErrNoRows) {
 			return nil, store.NewErrNotFound("Job", id)
 		}
 		return nil, errors.Wrapf(err, "failed to get Job with id=%s", id)
--- a/server/channels/store/sqlstore/utils.go
+++ b/server/channels/store/sqlstore/utils.go
@@ -5,6 +5,7 @@ package sqlstore

 import (
 	"database/sql"
+	"errors"
 	"fmt"
 	"io"
 	"net/url"
@@ -54,7 +55,7 @@ func MapStringsToQueryParams(list []string, paramPrefix string) (string, map[str
 // finalizeTransactionX ensures a transaction is closed after use, rolling back if not already committed.
 func finalizeTransactionX(transaction *sqlxTxWrapper, perr *error) {
 	// Rollback returns sql.ErrTxDone if the transaction was already closed.
-	if err := transaction.Rollback(); err != nil && err != sql.ErrTxDone {
+	if err := transaction.Rollback(); err != nil && !errors.Is(err, sql.ErrTxDone) {
 		*perr = merror.Append(*perr, err)
 	}
 }
--- a/server/channels/store/store.go
+++ b/server/channels/store/store.go
@@ -783,7 +783,7 @@ type JobStore interface {
 	SaveOnce(job *model.Job) (*model.Job, error)
 	UpdateOptimistically(job *model.Job, currentStatus string) (bool, error)
 	UpdateStatus(id string, status string) (*model.Job, error)
-	UpdateStatusOptimistically(id string, currentStatus string, newStatus string) (bool, error)
+	UpdateStatusOptimistically(id string, currentStatus string, newStatus string) (*model.Job, error)
 	Get(c request.CTX, id string) (*model.Job, error)
 	GetAllByType(c request.CTX, jobType string) ([]*model.Job, error)
 	GetAllByTypeAndStatus(c request.CTX, jobType string, status string) ([]*model.Job, error)
--- a/server/channels/store/storetest/job_store.go
+++ b/server/channels/store/storetest/job_store.go
@@ -600,9 +600,9 @@ func testJobUpdateStatusUpdateStatusOptimistically(t *testing.T, rctx request.CT

 	time.Sleep(2 * time.Millisecond)

-	updated, err := ss.Job().UpdateStatusOptimistically(job.Id, model.JobStatusInProgress, model.JobStatusSuccess)
+	updatedJob, err := ss.Job().UpdateStatusOptimistically(job.Id, model.JobStatusInProgress, model.JobStatusSuccess)
 	require.NoError(t, err)
-	require.False(t, updated)
+	require.Nil(t, updatedJob)

 	received, err = ss.Job().Get(rctx, job.Id)
 	require.NoError(t, err)
@@ -612,30 +612,26 @@ func testJobUpdateStatusUpdateStatusOptimistically(t *testing.T, rctx request.CT

 	time.Sleep(2 * time.Millisecond)

-	updated, err = ss.Job().UpdateStatusOptimistically(job.Id, model.JobStatusPending, model.JobStatusInProgress)
+	updatedJob, err = ss.Job().UpdateStatusOptimistically(job.Id, model.JobStatusPending, model.JobStatusInProgress)
 	require.NoError(t, err)
-	require.True(t, updated, "should have succeeded")
+	require.NotNil(t, updatedJob, "should have succeeded")

 	var startAtSet int64
-	received, err = ss.Job().Get(rctx, job.Id)
-	require.NoError(t, err)
-	require.Equal(t, model.JobStatusInProgress, received.Status)
-	require.NotEqual(t, 0, received.StartAt)
-	require.Greater(t, received.LastActivityAt, lastUpdateAt)
-	lastUpdateAt = received.LastActivityAt
-	startAtSet = received.StartAt
+	require.Equal(t, model.JobStatusInProgress, updatedJob.Status)
+	require.NotEqual(t, 0, updatedJob.StartAt)
+	require.Greater(t, updatedJob.LastActivityAt, lastUpdateAt)
+	lastUpdateAt = updatedJob.LastActivityAt
+	startAtSet = updatedJob.StartAt

 	time.Sleep(2 * time.Millisecond)

-	updated, err = ss.Job().UpdateStatusOptimistically(job.Id, model.JobStatusInProgress, model.JobStatusSuccess)
+	updatedJob, err = ss.Job().UpdateStatusOptimistically(job.Id, model.JobStatusInProgress, model.JobStatusSuccess)
 	require.NoError(t, err)
-	require.True(t, updated, "should have succeeded")
+	require.NotNil(t, updatedJob, "should have succeeded")

-	received, err = ss.Job().Get(rctx, job.Id)
-	require.NoError(t, err)
-	require.Equal(t, model.JobStatusSuccess, received.Status)
-	require.Equal(t, startAtSet, received.StartAt)
-	require.Greater(t, received.LastActivityAt, lastUpdateAt)
+	require.Equal(t, model.JobStatusSuccess, updatedJob.Status)
+	require.Equal(t, startAtSet, updatedJob.StartAt)
+	require.Greater(t, updatedJob.LastActivityAt, lastUpdateAt)
 }

 func testJobDelete(t *testing.T, rctx request.CTX, ss store.Store) {
--- a/server/channels/store/storetest/mocks/JobStore.go
+++ b/server/channels/store/storetest/mocks/JobStore.go
@@ -478,22 +478,24 @@ func (_m *JobStore) UpdateStatus(id string, status string) (*model.Job, error) {
 }

 // UpdateStatusOptimistically provides a mock function with given fields: id, currentStatus, newStatus
-func (_m *JobStore) UpdateStatusOptimistically(id string, currentStatus string, newStatus string) (bool, error) {
+func (_m *JobStore) UpdateStatusOptimistically(id string, currentStatus string, newStatus string) (*model.Job, error) {
 	ret := _m.Called(id, currentStatus, newStatus)

 	if len(ret) == 0 {
 		panic("no return value specified for UpdateStatusOptimistically")
 	}

-	var r0 bool
+	var r0 *model.Job
 	var r1 error
-	if rf, ok := ret.Get(0).(func(string, string, string) (bool, error)); ok {
+	if rf, ok := ret.Get(0).(func(string, string, string) (*model.Job, error)); ok {
 		return rf(id, currentStatus, newStatus)
 	}
-	if rf, ok := ret.Get(0).(func(string, string, string) bool); ok {
+	if rf, ok := ret.Get(0).(func(string, string, string) *model.Job); ok {
 		r0 = rf(id, currentStatus, newStatus)
 	} else {
-		r0 = ret.Get(0).(bool)
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*model.Job)
+		}
 	}

 	if rf, ok := ret.Get(1).(func(string, string, string) error); ok {
--- a/server/channels/store/timerlayer/timerlayer.go
+++ b/server/channels/store/timerlayer/timerlayer.go
@@ -5107,7 +5107,7 @@ func (s *TimerLayerJobStore) UpdateStatus(id string, status string) (*model.Job,
 	return result, err
 }

-func (s *TimerLayerJobStore) UpdateStatusOptimistically(id string, currentStatus string, newStatus string) (bool, error) {
+func (s *TimerLayerJobStore) UpdateStatusOptimistically(id string, currentStatus string, newStatus string) (*model.Job, error) {
 	start := time.Now()

 	result, err := s.JobStore.UpdateStatusOptimistically(id, currentStatus, newStatus)