1 files changed, 23 insertions, 3 deletions
diff --git a/modules/migrations/update.go b/modules/migrations/update.go
index d1465b2baf..3d0962657c 100644
--- a/modules/migrations/update.go
+++ b/modules/migrations/update.go
@@ -5,21 +5,28 @@
 package migrations
 
 import (
+	"context"
+
 	"code.gitea.io/gitea/models"
 	"code.gitea.io/gitea/modules/log"
 	"code.gitea.io/gitea/modules/structs"
 )
 
 // UpdateMigrationPosterID updates all migrated repositories' issues and comments posterID
-func UpdateMigrationPosterID() {
+func UpdateMigrationPosterID(ctx context.Context) {
 	for _, gitService := range structs.SupportedFullGitService {
-		if err := updateMigrationPosterIDByGitService(gitService); err != nil {
+		select {
+		case <-ctx.Done():
+			log.Warn("UpdateMigrationPosterID aborted due to shutdown before %s", gitService.Name())
+		default:
+		}
+		if err := updateMigrationPosterIDByGitService(ctx, gitService); err != nil {
 			log.Error("updateMigrationPosterIDByGitService failed: %v", err)
 		}
 	}
 }
 
-func updateMigrationPosterIDByGitService(tp structs.GitServiceType) error {
+func updateMigrationPosterIDByGitService(ctx context.Context, tp structs.GitServiceType) error {
 	provider := tp.Name()
 	if len(provider) == 0 {
 		return nil
@@ -28,6 +35,13 @@ func updateMigrationPosterIDByGitService(tp structs.GitServiceType) error {
 	const batchSize = 100
 	var start int
 	for {
+		select {
+		case <-ctx.Done():
+			log.Warn("UpdateMigrationPosterIDByGitService(%s) aborted due to shutdown", tp.Name())
+			return nil
+		default:
+		}
+
 		users, err := models.FindExternalUsersByProvider(models.FindExternalUserOptions{
 			Provider: provider,
 			Start:    start,
@@ -38,6 +52,12 @@ func updateMigrationPosterIDByGitService(tp structs.GitServiceType) error {
 		}
 
 		for _, user := range users {
+			select {
+			case <-ctx.Done():
+				log.Warn("UpdateMigrationPosterIDByGitService(%s) aborted due to shutdown", tp.Name())
+				return nil
+			default:
+			}
 			externalUserID := user.ExternalID
 			if err := models.UpdateMigrationsByType(tp, externalUserID, user.UserID); err != nil {
 				log.Error("UpdateMigrationsByType type %s external user id %v to local user id %v failed: %v", tp.Name(), user.ExternalID, user.UserID, err)