Second attempt at preventing zombies

cmd/cmd.go

@@ -7,9 +7,13 @@
 package cmd
 
 import (
+	"context"
 	"errors"
 	"fmt"
+	"os"
+	"os/signal"
 	"strings"
+	"syscall"
 
 	"code.gitea.io/gitea/models"
 	"code.gitea.io/gitea/modules/setting"
@@ -66,3 +70,25 @@ func initDBDisableConsole(disableConsole bool) error {
 	}
 	return nil
 }
+
+func installSignals() (context.Context, context.CancelFunc) {
+	ctx, cancel := context.WithCancel(context.Background())
+	go func() {
+		// install notify
+		signalChannel := make(chan os.Signal, 1)
+
+		signal.Notify(
+			signalChannel,
+			syscall.SIGINT,
+			syscall.SIGTERM,
+		)
+		select {
+		case <-signalChannel:
+		case <-ctx.Done():
+		}
+		cancel()
+		signal.Reset()
+	}()
+
+	return ctx, cancel
+}

cmd/hook.go

@@ -152,17 +152,18 @@ func runHookPreReceive(c *cli.Context) error {
 	if os.Getenv(models.EnvIsInternal) == "true" {
 		return nil
 	}
+	ctx, cancel := installSignals()
+	defer cancel()
 
 	setup("hooks/pre-receive.log", c.Bool("debug"))
 
 	if len(os.Getenv("SSH_ORIGINAL_COMMAND")) == 0 {
 		if setting.OnlyAllowPushIfGiteaEnvironmentSet {
-			fail(`Rejecting changes as Gitea environment not set.
+			return fail(`Rejecting changes as Gitea environment not set.
 If you are pushing over SSH you must push with a key managed by
 Gitea or set your environment appropriately.`, "")
-		} else {
-			return nil
 		}
+		return nil
 	}
 
 	// the environment is set by serv command
@@ -235,14 +236,14 @@ Gitea or set your environment appropriately.`, "")
 				hookOptions.OldCommitIDs = oldCommitIDs
 				hookOptions.NewCommitIDs = newCommitIDs
 				hookOptions.RefFullNames = refFullNames
-				statusCode, msg := private.HookPreReceive(username, reponame, hookOptions)
+				statusCode, msg := private.HookPreReceive(ctx, username, reponame, hookOptions)
 				switch statusCode {
 				case http.StatusOK:
 					// no-op
 				case http.StatusInternalServerError:
-					fail("Internal Server Error", msg)
+					return fail("Internal Server Error", msg)
 				default:
-					fail(msg, "")
+					return fail(msg, "")
 				}
 				count = 0
 				lastline = 0
@@ -263,12 +264,12 @@ Gitea or set your environment appropriately.`, "")
 
 		fmt.Fprintf(out, " Checking %d references\n", count)
 
-		statusCode, msg := private.HookPreReceive(username, reponame, hookOptions)
+		statusCode, msg := private.HookPreReceive(ctx, username, reponame, hookOptions)
 		switch statusCode {
 		case http.StatusInternalServerError:
-			fail("Internal Server Error", msg)
+			return fail("Internal Server Error", msg)
 		case http.StatusForbidden:
-			fail(msg, "")
+			return fail(msg, "")
 		}
 	} else if lastline > 0 {
 		fmt.Fprintf(out, "\n")
@@ -285,8 +286,11 @@ func runHookUpdate(c *cli.Context) error {
 }
 
 func runHookPostReceive(c *cli.Context) error {
+	ctx, cancel := installSignals()
+	defer cancel()
+
 	// First of all run update-server-info no matter what
-	if _, err := git.NewCommand("update-server-info").Run(); err != nil {
+	if _, err := git.NewCommand("update-server-info").SetParentContext(ctx).Run(); err != nil {
 		return fmt.Errorf("Failed to call 'git update-server-info': %v", err)
 	}
 
@@ -299,12 +303,11 @@ func runHookPostReceive(c *cli.Context) error {
 
 	if len(os.Getenv("SSH_ORIGINAL_COMMAND")) == 0 {
 		if setting.OnlyAllowPushIfGiteaEnvironmentSet {
-			fail(`Rejecting changes as Gitea environment not set.
+			return fail(`Rejecting changes as Gitea environment not set.
 If you are pushing over SSH you must push with a key managed by
 Gitea or set your environment appropriately.`, "")
-		} else {
-			return nil
 		}
+		return nil
 	}
 
 	var out io.Writer
@@ -371,11 +374,11 @@ Gitea or set your environment appropriately.`, "")
 			hookOptions.OldCommitIDs = oldCommitIDs
 			hookOptions.NewCommitIDs = newCommitIDs
 			hookOptions.RefFullNames = refFullNames
-			resp, err := private.HookPostReceive(repoUser, repoName, hookOptions)
+			resp, err := private.HookPostReceive(ctx, repoUser, repoName, hookOptions)
 			if resp == nil {
 				_ = dWriter.Close()
 				hookPrintResults(results)
-				fail("Internal Server Error", err)
+				return fail("Internal Server Error", err)
 			}
 			wasEmpty = wasEmpty || resp.RepoWasEmpty
 			results = append(results, resp.Results...)
@@ -386,9 +389,9 @@ Gitea or set your environment appropriately.`, "")
 	if count == 0 {
 		if wasEmpty && masterPushed {
 			// We need to tell the repo to reset the default branch to master
-			err := private.SetDefaultBranch(repoUser, repoName, "master")
+			err := private.SetDefaultBranch(ctx, repoUser, repoName, "master")
 			if err != nil {
-				fail("Internal Server Error", "SetDefaultBranch failed with Error: %v", err)
+				return fail("Internal Server Error", "SetDefaultBranch failed with Error: %v", err)
 			}
 		}
 		fmt.Fprintf(out, "Processed %d references in total\n", total)
@@ -404,11 +407,11 @@ Gitea or set your environment appropriately.`, "")
 
 	fmt.Fprintf(out, " Processing %d references\n", count)
 
-	resp, err := private.HookPostReceive(repoUser, repoName, hookOptions)
+	resp, err := private.HookPostReceive(ctx, repoUser, repoName, hookOptions)
 	if resp == nil {
 		_ = dWriter.Close()
 		hookPrintResults(results)
-		fail("Internal Server Error", err)
+		return fail("Internal Server Error", err)
 	}
 	wasEmpty = wasEmpty || resp.RepoWasEmpty
 	results = append(results, resp.Results...)
@@ -417,9 +420,9 @@ Gitea or set your environment appropriately.`, "")
 
 	if wasEmpty && masterPushed {
 		// We need to tell the repo to reset the default branch to master
-		err := private.SetDefaultBranch(repoUser, repoName, "master")
+		err := private.SetDefaultBranch(ctx, repoUser, repoName, "master")
 		if err != nil {
-			fail("Internal Server Error", "SetDefaultBranch failed with Error: %v", err)
+			return fail("Internal Server Error", "SetDefaultBranch failed with Error: %v", err)
 		}
 	}
 	_ = dWriter.Close()

cmd/keys.go

@@ -62,9 +62,12 @@ func runKeys(c *cli.Context) error {
 		return errors.New("No key type and content provided")
 	}
 
+	ctx, cancel := installSignals()
+	defer cancel()
+
 	setup("keys.log", false)
 
-	authorizedString, err := private.AuthorizedPublicKeyByContent(content)
+	authorizedString, err := private.AuthorizedPublicKeyByContent(ctx, content)
 	if err != nil {
 		return err
 	}

cmd/mailer.go

@@ -14,6 +14,9 @@ import (
 )
 
 func runSendMail(c *cli.Context) error {
+	ctx, cancel := installSignals()
+	defer cancel()
+
 	setting.NewContext()
 
 	if err := argsSet(c, "title"); err != nil {
@@ -39,7 +42,7 @@ func runSendMail(c *cli.Context) error {
 		}
 	}
 
-	status, message := private.SendEmail(subject, body, nil)
+	status, message := private.SendEmail(ctx, subject, body, nil)
 	if status != http.StatusOK {
 		fmt.Printf("error: %s\n", message)
 		return nil

cmd/manager.go

@@ -236,10 +236,13 @@ func runRemoveLogger(c *cli.Context) error {
 		group = log.DEFAULT
 	}
 	name := c.Args().First()
-	statusCode, msg := private.RemoveLogger(group, name)
+	ctx, cancel := installSignals()
+	defer cancel()
+
+	statusCode, msg := private.RemoveLogger(ctx, group, name)
 	switch statusCode {
 	case http.StatusInternalServerError:
-		fail("InternalServerError", msg)
+		return fail("InternalServerError", msg)
 	}
 
 	fmt.Fprintln(os.Stdout, msg)
@@ -371,82 +374,103 @@ func commonAddLogger(c *cli.Context, mode string, vals map[string]interface{}) e
 	if c.IsSet("name") {
 		name = c.String("name")
 	}
-	statusCode, msg := private.AddLogger(group, name, mode, vals)
+	ctx, cancel := installSignals()
+	defer cancel()
+
+	statusCode, msg := private.AddLogger(ctx, group, name, mode, vals)
 	switch statusCode {
 	case http.StatusInternalServerError:
-		fail("InternalServerError", msg)
+		return fail("InternalServerError", msg)
 	}
 
 	fmt.Fprintln(os.Stdout, msg)
 	return nil
 }
 
 func runShutdown(c *cli.Context) error {
+	ctx, cancel := installSignals()
+	defer cancel()
+
 	setup("manager", c.Bool("debug"))
-	statusCode, msg := private.Shutdown()
+	statusCode, msg := private.Shutdown(ctx)
 	switch statusCode {
 	case http.StatusInternalServerError:
-		fail("InternalServerError", msg)
+		return fail("InternalServerError", msg)
 	}
 
 	fmt.Fprintln(os.Stdout, msg)
 	return nil
 }
 
 func runRestart(c *cli.Context) error {
+	ctx, cancel := installSignals()
+	defer cancel()
+
 	setup("manager", c.Bool("debug"))
-	statusCode, msg := private.Restart()
+	statusCode, msg := private.Restart(ctx)
 	switch statusCode {
 	case http.StatusInternalServerError:
-		fail("InternalServerError", msg)
+		return fail("InternalServerError", msg)
 	}
 
 	fmt.Fprintln(os.Stdout, msg)
 	return nil
 }
 
 func runFlushQueues(c *cli.Context) error {
+	ctx, cancel := installSignals()
+	defer cancel()
+
 	setup("manager", c.Bool("debug"))
-	statusCode, msg := private.FlushQueues(c.Duration("timeout"), c.Bool("non-blocking"))
+	statusCode, msg := private.FlushQueues(ctx, c.Duration("timeout"), c.Bool("non-blocking"))
 	switch statusCode {
 	case http.StatusInternalServerError:
-		fail("InternalServerError", msg)
+		return fail("InternalServerError", msg)
 	}
 
 	fmt.Fprintln(os.Stdout, msg)
 	return nil
 }
 
 func runPauseLogging(c *cli.Context) error {
+	ctx, cancel := installSignals()
+	defer cancel()
+
 	setup("manager", c.Bool("debug"))
-	statusCode, msg := private.PauseLogging()
+	statusCode, msg := private.PauseLogging(ctx)
 	switch statusCode {
 	case http.StatusInternalServerError:
-		fail("InternalServerError", msg)
+		return fail("InternalServerError", msg)
 	}
 
 	fmt.Fprintln(os.Stdout, msg)
 	return nil
 }
 
 func runResumeLogging(c *cli.Context) error {
+	ctx, cancel := installSignals()
+	defer cancel()
+
 	setup("manager", c.Bool("debug"))
-	statusCode, msg := private.ResumeLogging()
+	statusCode, msg := private.ResumeLogging(ctx)
 	switch statusCode {
 	case http.StatusInternalServerError:
-		fail("InternalServerError", msg)
+		return fail("InternalServerError", msg)
 	}
 
 	fmt.Fprintln(os.Stdout, msg)
 	return nil
 }
 
 func runReleaseReopenLogging(c *cli.Context) error {
+	ctx, cancel := installSignals()
+	defer cancel()
+
 	setup("manager", c.Bool("debug"))
-	statusCode, msg := private.ReleaseReopenLogging()
+	statusCode, msg := private.ReleaseReopenLogging(ctx)
 	switch statusCode {
 	case http.StatusInternalServerError:
-		fail("InternalServerError", msg)
+		return fail("InternalServerError", msg)
 	}
 
 	fmt.Fprintln(os.Stdout, msg)

cmd/restore_repo.go

@@ -40,20 +40,24 @@ var CmdRestoreRepository = cli.Command{
 		cli.StringFlag{
 			Name:  "units",
 			Value: "",
-			Usage: `Which items will be restored, one or more units should be separated as comma. 
+			Usage: `Which items will be restored, one or more units should be separated as comma.
 wiki, issues, labels, releases, release_assets, milestones, pull_requests, comments are allowed. Empty means all units.`,
 		},
 	},
 }
 
-func runRestoreRepository(ctx *cli.Context) error {
+func runRestoreRepository(c *cli.Context) error {
+	ctx, cancel := installSignals()
+	defer cancel()
+
 	setting.NewContext()
 
 	statusCode, errStr := private.RestoreRepo(
-		ctx.String("repo_dir"),
-		ctx.String("owner_name"),
-		ctx.String("repo_name"),
-		ctx.StringSlice("units"),
+		ctx,
+		c.String("repo_dir"),
+		c.String("owner_name"),
+		c.String("repo_name"),
+		c.StringSlice("units"),
 	)
 	if statusCode == http.StatusOK {
 		return nil