Refactor application initialization

app/app.go

+package app
+
+import (
+	"os"
+	"path"
+
+	"github.com/sirupsen/logrus"
+	"github.com/urfave/cli"
+
+	"gitlab.com/gitlab-org/gitlab-runner/common"
+	"gitlab.com/gitlab-org/gitlab-runner/log"
+)
+
+var (
+	authors = []cli.Author{
+		{
+			Name:  "GitLab Inc.",
+			Email: "support@gitlab.com",
+		},
+	}
+)
+
+type Handler func(cliCtx *cli.Context) error
+type Handlers []Handler
+
+func (a *Handlers) Handle(cliCtx *cli.Context) error {
+	for _, f := range *a {
+		err := f(cliCtx)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+type App struct {
+	app *cli.App
+
+	beforeFunctions Handlers
+	afterFunctions  Handlers
+}
+
+func (a *App) init(usage string) {
+	app := cli.NewApp()
+	a.app = app
+
+	app.Name = path.Base(os.Args[0])
+	app.Usage = usage
+	app.Authors = authors
+	app.Version = common.AppVersion.ShortLine()
+	cli.VersionPrinter = common.AppVersion.Printer
+
+	app.Commands = common.GetCommands()
+	app.CommandNotFound = func(cliCtx *cli.Context, command string) {
+		logrus.Fatalf("Command %s not found", command)
+	}
+
+	a.beforeFunctions = make(Handlers, 0)
+	app.Before = a.beforeFunctions.Handle
+
+	a.afterFunctions = make(Handlers, 0)
+	app.After = a.afterFunctions.Handle
+}
+
+func (a *App) Run() {
+	if err := a.app.Run(os.Args); err != nil {
+		logrus.WithError(err).Fatal("Application execution failed")
+	}
+}
+
+func (a *App) Extend(extension func(*cli.App)) {
+	extension(a.app)
+}
+
+func (a *App) AppendBeforeFunc(f Handler) {
+	a.beforeFunctions = append(a.beforeFunctions, f)
+}
+
+func (a *App) AppendAfterFunc(f Handler) {
+	a.afterFunctions = append(a.afterFunctions, f)
+}
+
+func New(usage string) *App {
+	app := new(App)
+	app.init(usage)
+	app.Extend(log.AddFlags)
+	app.AppendBeforeFunc(log.ConfigureLogging)
+
+	return app
+}
+
+func Recover() {
+	r := recover()
+	if r != nil {
+		// log panics forces exit
+		if _, ok := r.(*logrus.Entry); ok {
+			os.Exit(1)
+		}
+		panic(r)
+	}
+}

app/cpuprofile.go

+package app
+
+import (
+	"os"
+	"runtime/pprof"
+
+	"github.com/urfave/cli"
+)
+
+func CPUProfileFlags(app *cli.App) {
+	app.Flags = append(app.Flags, cli.StringFlag{
+		Name:   "cpuprofile",
+		Usage:  "write cpu profile to file",
+		EnvVar: "CPU_PROFILE",
+	})
+}
+
+func CPUProfileSetup(cliCtx *cli.Context) error {
+	if cpuProfile := cliCtx.String("cpuprofile"); cpuProfile != "" {
+		f, err := os.Create(cpuProfile)
+		if err != nil {
+			return err
+		}
+		pprof.StartCPUProfile(f)
+	}
+
+	return nil
+}
+
+func CPUProfileTeardown(cliCtx *cli.Context) error {
+	pprof.StopCPUProfile()
+	return nil
+}

app/fix_home.go

+package app
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/docker/docker/pkg/homedir"
+	"github.com/urfave/cli"
+)
+
+func FixHOME(cliCtx *cli.Context) error {
+	// Fix home
+	if key := homedir.Key(); os.Getenv(key) == "" {
+		value := homedir.Get()
+		if value == "" {
+			return fmt.Errorf("the %q is not set", key)
+		}
+		os.Setenv(key, value)
+	}
+
+	return nil
+}

app/runtime_platform.go

+package app
+
+import (
+	"os"
+	"runtime"
+
+	"github.com/sirupsen/logrus"
+	"github.com/urfave/cli"
+
+	"gitlab.com/gitlab-org/gitlab-runner/common"
+)
+
+func LogRuntimePlatform(cliCtx *cli.Context) error {
+	fields := logrus.Fields{
+		"os":       runtime.GOOS,
+		"arch":     runtime.GOARCH,
+		"version":  common.VERSION,
+		"revision": common.REVISION,
+		"pid":      os.Getpid(),
+	}
+
+	logrus.WithFields(fields).Info("Runtime platform")
+
+	return nil
+}

helpers/cli/warn_on_bool.go → app/warn_on_bool.go

-package cli_helpers
+package app
 
 import (
 	"strings"

cmd/gitlab-runner-helper/main.go

 package main
 
 import (
-	"os"
-	"path"
-
 	"github.com/sirupsen/logrus"
-	"github.com/urfave/cli"
 
-	"gitlab.com/gitlab-org/gitlab-runner/common"
+	"gitlab.com/gitlab-org/gitlab-runner/app"
 	"gitlab.com/gitlab-org/gitlab-runner/log"
 
 	_ "gitlab.com/gitlab-org/gitlab-runner/commands/helpers"
 )
 
 func main() {
-	defer func() {
-		if r := recover(); r != nil {
-			// log panics forces exit
-			if _, ok := r.(*logrus.Entry); ok {
-				os.Exit(1)
-			}
-			panic(r)
-		}
-	}()
-
-	app := cli.NewApp()
-	app.Name = path.Base(os.Args[0])
-	app.Usage = "a GitLab Runner Helper"
-	app.Version = common.AppVersion.ShortLine()
-	cli.VersionPrinter = common.AppVersion.Printer
-	app.Authors = []cli.Author{
-		{
-			Name:  "GitLab Inc.",
-			Email: "support@gitlab.com",
-		},
-	}
-	app.Commands = common.GetCommands()
-	app.CommandNotFound = func(context *cli.Context, command string) {
-		logrus.Fatalln("Command", command, "not found")
-	}
+	defer app.Recover()
 
+	a := app.New("GitLab Runner Helper")
 	log.AddSecretsCleanupLogHook(logrus.StandardLogger())
-	log.ConfigureLogging(app)
 
-	if err := app.Run(os.Args); err != nil {
-		logrus.Fatal(err)
-	}
+	a.Run()
 }

cmd/gitlab-runner/main.go

@@ -2,14 +2,8 @@ package main
 
 import (
 	"os"
-	"path"
 
-	"github.com/sirupsen/logrus"
-	"github.com/urfave/cli"
-
-	"gitlab.com/gitlab-org/gitlab-runner/common"
-	"gitlab.com/gitlab-org/gitlab-runner/helpers/cli"
-	"gitlab.com/gitlab-org/gitlab-runner/log"
+	"gitlab.com/gitlab-org/gitlab-runner/app"
 
 	_ "gitlab.com/gitlab-org/gitlab-runner/cache/gcs"
 	_ "gitlab.com/gitlab-org/gitlab-runner/cache/s3"
@@ -26,40 +20,15 @@ import (
 )
 
 func main() {
-	defer func() {
-		if r := recover(); r != nil {
-			// log panics forces exit
-			if _, ok := r.(*logrus.Entry); ok {
-				os.Exit(1)
-			}
-			panic(r)
-		}
-	}()
-
-	app := cli.NewApp()
-	app.Name = path.Base(os.Args[0])
-	app.Usage = "a GitLab Runner"
-	app.Version = common.AppVersion.ShortLine()
-	cli.VersionPrinter = common.AppVersion.Printer
-	app.Authors = []cli.Author{
-		{
-			Name:  "GitLab Inc.",
-			Email: "support@gitlab.com",
-		},
-	}
-	app.Commands = common.GetCommands()
-	app.CommandNotFound = func(context *cli.Context, command string) {
-		logrus.Fatalln("Command", command, "not found.")
-	}
-
-	cli_helpers.LogRuntimePlatform(app)
-	cli_helpers.SetupCPUProfile(app)
-	cli_helpers.FixHOME(app)
-	cli_helpers.WarnOnBool(os.Args)
+	defer app.Recover()
 
-	log.ConfigureLogging(app)
+	a := app.New("GitLab Runner")
+	a.AppendBeforeFunc(app.LogRuntimePlatform)
+	a.Extend(app.CPUProfileFlags)
+	a.AppendBeforeFunc(app.CPUProfileSetup)
+	a.AppendAfterFunc(app.CPUProfileTeardown)
+	a.AppendBeforeFunc(app.FixHOME)
+	app.WarnOnBool(os.Args)
 
-	if err := app.Run(os.Args); err != nil {
-		logrus.Fatal(err)
-	}
+	a.Run()
 }

helpers/cli/cpuprofile.go

-package cli_helpers
-
-import (
-	"os"
-	"runtime/pprof"
-
-	"github.com/urfave/cli"
-)
-
-func SetupCPUProfile(app *cli.App) {
-	app.Flags = append(app.Flags, cli.StringFlag{
-		Name:   "cpuprofile",
-		Usage:  "write cpu profile to file",
-		EnvVar: "CPU_PROFILE",
-	})
-
-	appBefore := app.Before
-	appAfter := app.After
-
-	app.Before = func(c *cli.Context) error {
-		if cpuProfile := c.String("cpuprofile"); cpuProfile != "" {
-			f, err := os.Create(cpuProfile)
-			if err != nil {
-				return err
-			}
-			pprof.StartCPUProfile(f)
-		}
-
-		if appBefore != nil {
-			return appBefore(c)
-		}
-		return nil
-	}
-
-	app.After = func(c *cli.Context) error {
-		pprof.StopCPUProfile()
-
-		if appAfter != nil {
-			return appAfter(c)
-		}
-		return nil
-	}
-}

helpers/cli/fix_home.go

-package cli_helpers
-
-import (
-	"fmt"
-	"os"
-
-	"github.com/docker/docker/pkg/homedir"
-	"github.com/urfave/cli"
-)
-
-func FixHOME(app *cli.App) {
-	appBefore := app.Before
-
-	app.Before = func(c *cli.Context) error {
-		// Fix home
-		if key := homedir.Key(); os.Getenv(key) == "" {
-			value := homedir.Get()
-			if value == "" {
-				return fmt.Errorf("the %q is not set", key)
-			}
-			os.Setenv(key, value)
-		}
-
-		if appBefore != nil {
-			return appBefore(c)
-		}
-		return nil
-	}
-}

helpers/cli/runtime_platform.go

-package cli_helpers
-
-import (
-	"os"
-	"runtime"
-
-	"github.com/sirupsen/logrus"
-	"github.com/urfave/cli"
-
-	"gitlab.com/gitlab-org/gitlab-runner/common"
-)
-
-func LogRuntimePlatform(app *cli.App) {
-	appBefore := app.Before
-	app.Before = func(c *cli.Context) error {
-		fields := logrus.Fields{
-			"os":       runtime.GOOS,
-			"arch":     runtime.GOARCH,
-			"version":  common.VERSION,
-			"revision": common.REVISION,
-			"pid":      os.Getpid(),
-		}
-
-		logrus.WithFields(fields).Info("Runtime platform")
-
-		if appBefore != nil {
-			return appBefore(c)
-		}
-		return nil
-	}
-
-}

log/configuration.go

@@ -162,21 +162,17 @@ func Configuration() *Config {
 	return configuration
 }
 
-func ConfigureLogging(app *cli.App) {
+func AddFlags(app *cli.App) {
 	app.Flags = append(app.Flags, logFlags...)
+}
 
-	appBefore := app.Before
-	app.Before = func(cliCtx *cli.Context) error {
-		Configuration().logger.SetOutput(os.Stderr)
-
-		err := Configuration().handleCliCtx(cliCtx)
-		if err != nil {
-			logrus.WithError(err).Fatal("Error while setting up logging configuration")
-		}
+func ConfigureLogging(cliCtx *cli.Context) error {
+	Configuration().logger.SetOutput(os.Stderr)
 
-		if appBefore != nil {
-			return appBefore(cliCtx)
-		}
-		return nil
+	err := Configuration().handleCliCtx(cliCtx)
+	if err != nil {
+		return fmt.Errorf("error while setting up logging configuration: %v", err)
 	}
+
+	return nil
 }

log/configuration_test.go

@@ -6,10 +6,7 @@ import (
 	"github.com/sirupsen/logrus"
 	"github.com/sirupsen/logrus/hooks/test"
 	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
 	"github.com/urfave/cli"
-
-	"gitlab.com/gitlab-org/gitlab-runner/helpers"
 )
 
 func prepareFakeConfiguration(logger *logrus.Logger) func() {
@@ -22,7 +19,7 @@ func prepareFakeConfiguration(logger *logrus.Logger) func() {
 	}
 }
 
-func testCommandRun(args ...string) {
+func testCommandRun(t *testing.T, args ...string) error {
 	app := cli.NewApp()
 	app.Commands = []cli.Command{
 		{
@@ -31,12 +28,13 @@ func testCommandRun(args ...string) {
 		},
 	}
 
-	ConfigureLogging(app)
+	AddFlags(app)
+	app.Before = ConfigureLogging
 
 	args = append([]string{"binary"}, args...)
 	args = append(args, "logtest")
 
-	app.Run(args)
+	return app.Run(args)
 }
 
 type handleCliCtxTestCase struct {
@@ -96,46 +94,25 @@ func TestHandleCliCtx(t *testing.T) {
 			logger, _ := test.NewNullLogger()
 
 			defer prepareFakeConfiguration(logger)()
-			defer helpers.MakeFatalToPanic()()
-
-			testFunc := func() {
-				testCommandRun(testCase.args...)
-				if testCase.expectedError == "" {
-					assert.Equal(t, testCase.expectedLevel, Configuration().level)
-					assert.Equal(t, testCase.expectedFormatter, Configuration().format)
-					assert.Equal(t, testCase.expectedLevelSetWithCli, Configuration().IsLevelSetWithCli())
-					assert.Equal(t, testCase.expectedFormatSetWithCli, Configuration().IsFormatSetWithCli())
-
-					if testCase.goroutinesDumpStopChExists {
-						assert.NotNil(t, Configuration().goroutinesDumpStopCh)
-					} else {
-						assert.Nil(t, Configuration().goroutinesDumpStopCh)
-					}
-				}
-			}
-
-			if testCase.expectedError != "" {
-				var message *logrus.Entry
-				var ok bool
 
-				func() {
-					defer func() {
-						message, ok = recover().(*logrus.Entry)
-					}()
+			err := testCommandRun(t, testCase.args...)
 
-					testFunc()
-				}()
+			if testCase.expectedError == "" {
+				assert.NoError(t, err)
 
-				require.True(t, ok)
-
-				panicMessage, err := message.String()
-				require.NoError(t, err)
-
-				assert.Contains(t, panicMessage, "Error while setting up logging configuration")
-				assert.Contains(t, panicMessage, testCase.expectedError)
+				assert.Equal(t, testCase.expectedLevel, Configuration().level)
+				assert.Equal(t, testCase.expectedFormatter, Configuration().format)
+				assert.Equal(t, testCase.expectedLevelSetWithCli, Configuration().IsLevelSetWithCli())
+				assert.Equal(t, testCase.expectedFormatSetWithCli, Configuration().IsFormatSetWithCli())
 
+				if testCase.goroutinesDumpStopChExists {
+					assert.NotNil(t, Configuration().goroutinesDumpStopCh)
+				} else {
+					assert.Nil(t, Configuration().goroutinesDumpStopCh)
+				}
 			} else {
-				assert.NotPanics(t, testFunc)
+				assert.Error(t, err)
+				assert.Contains(t, err.Error(), "error while setting up logging configuration")
 			}
 		})
 	}