From 5dbe6af1d085463ff0a1414e067daef48aa3d838 Mon Sep 17 00:00:00 2001
From: ErickSkrauch <erickskrauch@yandex.ru>
Date: Fri, 1 May 2020 00:06:56 +0300
Subject: [PATCH] Added --cpuprofile flag for the dev Docker images

---
 .travis.yml           |  3 +++
 CHANGELOG.md          |  1 +
 cmd/root_profiling.go | 55 +++++++++++++++++++++++++++++++++++++++++++
 http/http.go          |  3 +--
 4 files changed, 60 insertions(+), 2 deletions(-)
 create mode 100644 cmd/root_profiling.go

diff --git a/.travis.yml b/.travis.yml
index db5c053..92fa5ee 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -48,9 +48,12 @@ jobs:
         - docker login -u="$DOCKER_USERNAME" -p="$DOCKER_PASSWORD"
         - export DOCKER_TAG="${TRAVIS_TAG:-dev}"
         - export APP_VERSION="${TRAVIS_TAG:-dev-${TRAVIS_COMMIT:0:7}}"
+        - export BUILD_TAGS=""
+        - if [ "$DOCKER_TAG" == "dev" ]; then export BUILD_TAGS="$BUILD_TAGS --tags profiling"; fi
         - >
           env CGO_ENABLED=0 GOOS=linux GOARCH=amd64
           go build
+          $BUILD_TAGS
           -o release/chrly
           -ldflags "-extldflags '-static' -X github.com/elyby/chrly/version.version=$APP_VERSION -X github.com/elyby/chrly/version.commit=$TRAVIS_COMMIT"
           main.go
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4d78627..a4a4976 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,6 +12,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - New configuration params: `MOJANG_API_BASE_URL` and `MOJANG_SESSION_SERVER_BASE_URL`, that allow you to spoof
   Mojang API base addresses.
 - New health checker, that ensures that response for textures provider from Mojang's API is valid.
+- `dev` Docker images now have the `--cpuprofile` flag, which allows you to run the program with CPU profiling.
 
 ### Fixed
 - Handle the case when there is no textures property in Mojang's response.
diff --git a/cmd/root_profiling.go b/cmd/root_profiling.go
new file mode 100644
index 0000000..bbd7e18
--- /dev/null
+++ b/cmd/root_profiling.go
@@ -0,0 +1,55 @@
+// +build profiling
+
+package cmd
+
+import (
+	"log"
+	"os"
+	"runtime/pprof"
+
+	"github.com/spf13/cobra"
+)
+
+func init() {
+	var profilePath string
+	RootCmd.PersistentFlags().StringVar(&profilePath, "cpuprofile", "", "enables pprof profiling and sets its output path")
+
+	pprofEnabled := false
+	originalPersistentPreRunE := RootCmd.PersistentPreRunE
+	RootCmd.PersistentPreRunE = func(cmd *cobra.Command, args []string) error {
+		if profilePath == "" {
+			return nil
+		}
+
+		f, err := os.Create(profilePath)
+		if err != nil {
+			return err
+		}
+
+		log.Println("enabling profiling")
+		err = pprof.StartCPUProfile(f)
+		if err != nil {
+			return err
+		}
+
+		pprofEnabled = true
+
+		if originalPersistentPreRunE != nil {
+			return originalPersistentPreRunE(cmd, args)
+		}
+
+		return nil
+	}
+
+	originalPersistentPostRun := RootCmd.PersistentPreRun
+	RootCmd.PersistentPostRun = func(cmd *cobra.Command, args []string) {
+		if pprofEnabled {
+			log.Println("shutting down profiling")
+			pprof.StopCPUProfile()
+		}
+
+		if originalPersistentPostRun != nil {
+			originalPersistentPostRun(cmd, args)
+		}
+	}
+}
diff --git a/http/http.go b/http/http.go
index baa9621..353eca1 100644
--- a/http/http.go
+++ b/http/http.go
@@ -28,9 +28,8 @@ func StartServer(server *http.Server, logger slf.Logger) {
 		logger.Info("Starting the server, HTTP on: :addr", wd.StringParam("addr", server.Addr))
 		if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
 			logger.Emergency("Error in main(): :err", wd.ErrParam(err))
+			close(done)
 		}
-
-		close(done)
 	}()
 
 	go func() {