Fix racy test issues (#24971)

server/channels/jobs/batch_migration_worker.go

@@ -5,6 +5,7 @@ package jobs
 
 import (
 	"net/http"
+	"sync/atomic"
 	"time"
 
 	"github.com/mattermost/mattermost/server/public/model"
@@ -33,8 +34,9 @@ type BatchMigrationWorker struct {
 	store     store.Store
 	app       BatchMigrationWorkerAppIFace
 
-	stop    chan bool
+	stop    chan struct{}
 	stopped chan bool
+	closed  atomic.Bool
 	jobs    chan model.Job
 
 	migrationKey       string
@@ -49,7 +51,7 @@ func MakeBatchMigrationWorker(jobServer *JobServer, store store.Store, app Batch
 		logger:             jobServer.Logger().With(mlog.String("worker_name", migrationKey)),
 		store:              store,
 		app:                app,
-		stop:               make(chan bool, 1),
+		stop:               make(chan struct{}),
 		stopped:            make(chan bool, 1),
 		jobs:               make(chan model.Job),
 		migrationKey:       migrationKey,
@@ -64,7 +66,9 @@ func (worker *BatchMigrationWorker) Run() {
 	worker.logger.Debug("Worker started")
 	// We have to re-assign the stop channel again, because
 	// it might happen that the job was restarted due to a config change.
-	worker.stop = make(chan bool, 1)
+	if worker.closed.CompareAndSwap(true, false) {
+		worker.stop = make(chan struct{})
+	}
 
 	defer func() {
 		worker.logger.Debug("Worker finished")
@@ -84,6 +88,11 @@ func (worker *BatchMigrationWorker) Run() {
 
 // Stop interrupts the worker even if the migration has not yet completed.
 func (worker *BatchMigrationWorker) Stop() {
+	// Set to close, and if already closed before, then return.
+	if !worker.closed.CompareAndSwap(false, true) {
+		return
+	}
+
 	worker.logger.Debug("Worker stopping")
 	close(worker.stop)
 	<-worker.stopped

server/channels/jobs/helper_test.go

@@ -88,7 +88,6 @@ func setupTestHelper(dbStore store.Store, enterprise bool, includeCacheLayer boo
 		IncludeCacheLayer: includeCacheLayer,
 		ConfigStore:       configStore,
 	}
-	th.Context.SetLogger(testLogger)
 
 	prevListenAddress := *th.App.Config().ServiceSettings.ListenAddress
 	th.App.UpdateConfig(func(cfg *model.Config) { *cfg.ServiceSettings.ListenAddress = "localhost:0" })

server/channels/jobs/migrations/worker.go

@@ -100,10 +100,10 @@ func (worker *Worker) DoJob(job *model.Job) {
 		return
 	}
 
-	cancelContext := request.EmptyContext(worker.logger)
+	var cancelContext request.CTX = request.EmptyContext(worker.logger)
 	cancelCtx, cancelCancelWatcher := context.WithCancel(context.Background())
 	cancelWatcherChan := make(chan struct{}, 1)
-	cancelContext.SetContext(cancelCtx)
+	cancelContext = cancelContext.WithContext(cancelCtx)
 	go worker.jobServer.CancellationWatcher(cancelContext, job.Id, cancelWatcherChan)
 	defer cancelCancelWatcher()
 

server/channels/jobs/s3_path_migration/s3_path_migration.go

@@ -29,7 +29,7 @@ type S3PathMigrationWorker struct {
 	store       store.Store
 	fileBackend *filestore.S3FileBackend
 
-	stop    chan bool
+	stop    chan struct{}
 	stopped chan bool
 	jobs    chan model.Job
 }
@@ -45,7 +45,7 @@ func MakeWorker(jobServer *jobs.JobServer, store store.Store, fileBackend filest
 		logger:      jobServer.Logger().With(mlog.String("worker_name", workerName)),
 		store:       store,
 		fileBackend: s3Backend,
-		stop:        make(chan bool, 1),
+		stop:        make(chan struct{}),
 		stopped:     make(chan bool, 1),
 		jobs:        make(chan model.Job),
 	}
@@ -56,7 +56,7 @@ func (worker *S3PathMigrationWorker) Run() {
 	worker.logger.Debug("Worker started")
 	// We have to re-assign the stop channel again, because
 	// it might happen that the job was restarted due to a config change.
-	worker.stop = make(chan bool, 1)
+	worker.stop = make(chan struct{}, 1)
 
 	defer func() {
 		worker.logger.Debug("Worker finished")