Graceful: Xorm, RepoIndexer, Cron and Others

cmd/migrate.go

@@ -5,6 +5,8 @@
 package cmd
 
 import (
+	"context"
+
 	"code.gitea.io/gitea/models"
 	"code.gitea.io/gitea/models/migrations"
 	"code.gitea.io/gitea/modules/log"
@@ -32,7 +34,7 @@ func runMigrate(ctx *cli.Context) error {
 	log.Trace("Log path: %s", setting.LogRootPath)
 	setting.InitDBConfig()
 
-	if err := models.NewEngine(migrations.Migrate); err != nil {
+	if err := models.NewEngine(context.Background(), migrations.Migrate); err != nil {
 		log.Fatal("Failed to initialize ORM engine: %v", err)
 		return err
 	}

cmd/web.go

@@ -5,6 +5,7 @@
 package cmd
 
 import (
+	"context"
 	"fmt"
 	"net/http"
 	_ "net/http/pprof" // Used for debugging if enabled and a web server is running
@@ -96,6 +97,10 @@ func runLetsEncryptFallbackHandler(w http.ResponseWriter, r *http.Request) {
 }
 
 func runWeb(ctx *cli.Context) error {
+	managerCtx, cancel := context.WithCancel(context.Background())
+	graceful.InitManager(managerCtx)
+	defer cancel()
+
 	if os.Getppid() > 1 && len(os.Getenv("LISTEN_FDS")) > 0 {
 		log.Info("Restarting Gitea on PID: %d from parent PID: %d", os.Getpid(), os.Getppid())
 	} else {
@@ -108,7 +113,7 @@ func runWeb(ctx *cli.Context) error {
 	}
 
 	// Perform global initialization
-	routers.GlobalInit()
+	routers.GlobalInit(graceful.GetManager().HammerContext())
 
 	// Set up Macaron
 	m := routes.NewMacaron()
@@ -199,8 +204,7 @@ func runWeb(ctx *cli.Context) error {
 		log.Critical("Failed to start server: %v", err)
 	}
 	log.Info("HTTP Listener: %s Closed", listenAddr)
-	graceful.Manager.WaitForServers()
-	graceful.Manager.WaitForTerminate()
+	<-graceful.GetManager().Done()
 	log.Info("PID: %d Gitea Web Finished", os.Getpid())
 	log.Close()
 	return nil

cmd/web_graceful.go

@@ -28,13 +28,13 @@ func runHTTPSWithTLSConfig(network, listenAddr string, tlsConfig *tls.Config, m
 
 // NoHTTPRedirector tells our cleanup routine that we will not be using a fallback http redirector
 func NoHTTPRedirector() {
-	graceful.Manager.InformCleanup()
+	graceful.GetManager().InformCleanup()
 }
 
 // NoMainListener tells our cleanup routine that we will not be using a possibly provided listener
 // for our main HTTP/HTTPS service
 func NoMainListener() {
-	graceful.Manager.InformCleanup()
+	graceful.GetManager().InformCleanup()
 }
 
 func runFCGI(network, listenAddr string, m http.Handler) error {

contrib/pr/checkout.go

@@ -5,6 +5,7 @@ Checkout a PR and load the tests data into sqlite database
 */
 
 import (
+	"context"
 	"flag"
 	"fmt"
 	"io/ioutil"
@@ -92,7 +93,7 @@ func runPR() {
 	//x, err = xorm.NewEngine("sqlite3", "file::memory:?cache=shared")
 
 	var helper testfixtures.Helper = &testfixtures.SQLite{}
-	models.NewEngine(func(_ *xorm.Engine) error {
+	models.NewEngine(context.Background(), func(_ *xorm.Engine) error {
 		return nil
 	})
 	models.HasEngine = true

integrations/auth_ldap_test.go

@@ -5,6 +5,7 @@
 package integrations
 
 import (
+	"context"
 	"net/http"
 	"os"
 	"strings"
@@ -147,7 +148,7 @@ func TestLDAPUserSync(t *testing.T) {
 	}
 	defer prepareTestEnv(t)()
 	addAuthSourceLDAP(t, "")
-	models.SyncExternalUsers()
+	models.SyncExternalUsers(context.Background())
 
 	session := loginUser(t, "user1")
 	// Check if users exists
@@ -206,7 +207,8 @@ func TestLDAPUserSSHKeySync(t *testing.T) {
 	}
 	defer prepareTestEnv(t)()
 	addAuthSourceLDAP(t, "sshPublicKey")
-	models.SyncExternalUsers()
+
+	models.SyncExternalUsers(context.Background())
 
 	// Check if users has SSH keys synced
 	for _, u := range gitLDAPUsers {

integrations/integration_test.go

@@ -6,6 +6,7 @@ package integrations
 
 import (
 	"bytes"
+	"context"
 	"database/sql"
 	"encoding/json"
 	"fmt"
@@ -24,6 +25,7 @@ import (
 
 	"code.gitea.io/gitea/models"
 	"code.gitea.io/gitea/modules/base"
+	"code.gitea.io/gitea/modules/graceful"
 	"code.gitea.io/gitea/modules/setting"
 	"code.gitea.io/gitea/routers"
 	"code.gitea.io/gitea/routers/routes"
@@ -55,6 +57,10 @@ func NewNilResponseRecorder() *NilResponseRecorder {
 }
 
 func TestMain(m *testing.M) {
+	managerCtx, cancel := context.WithCancel(context.Background())
+	graceful.InitManager(managerCtx)
+	defer cancel()
+
 	initIntegrationTest()
 	mac = routes.NewMacaron()
 	routes.RegisterRoutes(mac)
@@ -171,7 +177,7 @@ func initIntegrationTest() {
 		}
 		defer db.Close()
 	}
-	routers.GlobalInit()
+	routers.GlobalInit(graceful.GetManager().HammerContext())
 }
 
 func prepareTestEnv(t testing.TB, skip ...int) func() {

integrations/migration-test/migration_test.go

@@ -6,6 +6,7 @@ package migrations
 
 import (
 	"compress/gzip"
+	"context"
 	"database/sql"
 	"fmt"
 	"io/ioutil"
@@ -220,7 +221,7 @@ func doMigrationTest(t *testing.T, version string) {
 	err := models.SetEngine()
 	assert.NoError(t, err)
 
-	err = models.NewEngine(wrappedMigrate)
+	err = models.NewEngine(context.Background(), wrappedMigrate)
 	assert.NoError(t, err)
 	currentEngine.Close()
 }

models/branches.go

@@ -5,6 +5,7 @@
 package models
 
 import (
+	"context"
 	"fmt"
 	"time"
 
@@ -525,7 +526,8 @@ func (deletedBranch *DeletedBranch) LoadUser() {
 }
 
 // RemoveOldDeletedBranches removes old deleted branches
-func RemoveOldDeletedBranches() {
+func RemoveOldDeletedBranches(ctx context.Context) {
+	// Nothing to do for shutdown or terminate
 	log.Trace("Doing: DeletedBranchesCleanup")
 
 	deleteBefore := time.Now().Add(-setting.Cron.DeletedBranchesCleanup.OlderThan)

models/models.go

@@ -6,6 +6,7 @@
 package models
 
 import (
+	"context"
 	"database/sql"
 	"errors"
 	"fmt"
@@ -164,11 +165,13 @@ func SetEngine() (err error) {
 }
 
 // NewEngine initializes a new xorm.Engine
-func NewEngine(migrateFunc func(*xorm.Engine) error) (err error) {
+func NewEngine(ctx context.Context, migrateFunc func(*xorm.Engine) error) (err error) {
 	if err = SetEngine(); err != nil {
 		return err
 	}
 
+	x.SetDefaultContext(ctx)
+
 	if err = x.Ping(); err != nil {
 		return err
 	}

models/pull_list.go

@@ -68,11 +68,12 @@ func GetUnmergedPullRequestsByBaseInfo(repoID int64, branch string) ([]*PullRequ
 		Find(&prs)
 }
 
-// GetPullRequestsByCheckStatus returns all pull requests according the special checking status.
-func GetPullRequestsByCheckStatus(status PullRequestStatus) ([]*PullRequest, error) {
-	prs := make([]*PullRequest, 0, 10)
-	return prs, x.
+// GetPullRequestIDsByCheckStatus returns all pull requests according the special checking status.
+func GetPullRequestIDsByCheckStatus(status PullRequestStatus) ([]int64, error) {
+	prs := make([]int64, 0, 10)
+	return prs, x.Table("pull_request").
 		Where("status=?", status).
+		Cols("pull_request.id").
 		Find(&prs)
 }
 

models/repo.go

@@ -7,6 +7,7 @@ package models
 
 import (
 	"bytes"
+	"context"
 	"crypto/md5"
 	"errors"
 	"fmt"
@@ -2098,19 +2099,27 @@ func DeleteRepositoryArchives() error {
 }
 
 // DeleteOldRepositoryArchives deletes old repository archives.
-func DeleteOldRepositoryArchives() {
+func DeleteOldRepositoryArchives(ctx context.Context) {
 	log.Trace("Doing: ArchiveCleanup")
 
-	if err := x.Where("id > 0").Iterate(new(Repository), deleteOldRepositoryArchives); err != nil {
+	if err := x.Where("id > 0").Iterate(new(Repository), func(idx int, bean interface{}) error {
+		return deleteOldRepositoryArchives(ctx, idx, bean)
+	}); err != nil {
 		log.Error("ArchiveClean: %v", err)
 	}
 }
 
-func deleteOldRepositoryArchives(idx int, bean interface{}) error {
+func deleteOldRepositoryArchives(ctx context.Context, idx int, bean interface{}) error {
 	repo := bean.(*Repository)
 	basePath := filepath.Join(repo.RepoPath(), "archives")
 
 	for _, ty := range []string{"zip", "targz"} {
+		select {
+		case <-ctx.Done():
+			return fmt.Errorf("Aborted due to shutdown:\nin delete of old repository archives %v\nat delete file %s", repo, ty)
+		default:
+		}
+
 		path := filepath.Join(basePath, ty)
 		file, err := os.Open(path)
 		if err != nil {
@@ -2133,6 +2142,11 @@ func deleteOldRepositoryArchives(idx int, bean interface{}) error {
 		minimumOldestTime := time.Now().Add(-setting.Cron.ArchiveCleanup.OlderThan)
 		for _, info := range files {
 			if info.ModTime().Before(minimumOldestTime) && !info.IsDir() {
+				select {
+				case <-ctx.Done():
+					return fmt.Errorf("Aborted due to shutdown:\nin delete of old repository archives %v\nat delete file %s - %s", repo, ty, info.Name())
+				default:
+				}
 				toDelete := filepath.Join(path, info.Name())
 				// This is a best-effort purge, so we do not check error codes to confirm removal.
 				if err = os.Remove(toDelete); err != nil {
@@ -2226,13 +2240,17 @@ func SyncRepositoryHooks() error {
 }
 
 // GitFsck calls 'git fsck' to check repository health.
-func GitFsck() {
+func GitFsck(ctx context.Context) {
 	log.Trace("Doing: GitFsck")
-
 	if err := x.
 		Where("id>0 AND is_fsck_enabled=?", true).BufferSize(setting.Database.IterateBufferSize).
 		Iterate(new(Repository),
 			func(idx int, bean interface{}) error {
+				select {
+				case <-ctx.Done():
+					return fmt.Errorf("Aborted due to shutdown")
+				default:
+				}
 				repo := bean.(*Repository)
 				repoPath := repo.RepoPath()
 				log.Trace("Running health check on repository %s", repoPath)
@@ -2278,13 +2296,19 @@ type repoChecker struct {
 	desc                 string
 }
 
-func repoStatsCheck(checker *repoChecker) {
+func repoStatsCheck(ctx context.Context, checker *repoChecker) {
 	results, err := x.Query(checker.querySQL)
 	if err != nil {
 		log.Error("Select %s: %v", checker.desc, err)
 		return
 	}
 	for _, result := range results {
+		select {
+		case <-ctx.Done():
+			log.Warn("CheckRepoStats: Aborting due to shutdown")
+			return
+		default:
+		}
 		id := com.StrTo(result["id"]).MustInt64()
 		log.Trace("Updating %s: %d", checker.desc, id)
 		_, err = x.Exec(checker.correctSQL, id, id)
@@ -2295,7 +2319,7 @@ func repoStatsCheck(checker *repoChecker) {
 }
 
 // CheckRepoStats checks the repository stats
-func CheckRepoStats() {
+func CheckRepoStats(ctx context.Context) {
 	log.Trace("Doing: CheckRepoStats")
 
 	checkers := []*repoChecker{
@@ -2331,7 +2355,13 @@ func CheckRepoStats() {
 		},
 	}
 	for i := range checkers {
-		repoStatsCheck(checkers[i])
+		select {
+		case <-ctx.Done():
+			log.Warn("CheckRepoStats: Aborting due to shutdown")
+			return
+		default:
+			repoStatsCheck(ctx, checkers[i])
+		}
 	}
 
 	// ***** START: Repository.NumClosedIssues *****
@@ -2341,6 +2371,12 @@ func CheckRepoStats() {
 		log.Error("Select %s: %v", desc, err)
 	} else {
 		for _, result := range results {
+			select {
+			case <-ctx.Done():
+				log.Warn("CheckRepoStats: Aborting due to shutdown")
+				return
+			default:
+			}
 			id := com.StrTo(result["id"]).MustInt64()
 			log.Trace("Updating %s: %d", desc, id)
 			_, err = x.Exec("UPDATE `repository` SET num_closed_issues=(SELECT COUNT(*) FROM `issue` WHERE repo_id=? AND is_closed=? AND is_pull=?) WHERE id=?", id, true, false, id)
@@ -2358,6 +2394,12 @@ func CheckRepoStats() {
 		log.Error("Select %s: %v", desc, err)
 	} else {
 		for _, result := range results {
+			select {
+			case <-ctx.Done():
+				log.Warn("CheckRepoStats: Aborting due to shutdown")
+				return
+			default:
+			}
 			id := com.StrTo(result["id"]).MustInt64()
 			log.Trace("Updating %s: %d", desc, id)
 			_, err = x.Exec("UPDATE `repository` SET num_closed_pulls=(SELECT COUNT(*) FROM `issue` WHERE repo_id=? AND is_closed=? AND is_pull=?) WHERE id=?", id, true, true, id)
@@ -2375,6 +2417,12 @@ func CheckRepoStats() {
 		log.Error("Select repository count 'num_forks': %v", err)
 	} else {
 		for _, result := range results {
+			select {
+			case <-ctx.Done():
+				log.Warn("CheckRepoStats: Aborting due to shutdown")
+				return
+			default:
+			}
 			id := com.StrTo(result["id"]).MustInt64()
 			log.Trace("Updating repository count 'num_forks': %d", id)