From 511a8cbe046d5344c7a6e9dffbe1d555d094ea95 Mon Sep 17 00:00:00 2001
From: Harshavardhana <harsha@minio.io>
Date: Thu, 1 Sep 2016 03:50:31 -0700
Subject: [PATCH] profiler: Fix how we generate profiles. (#2596)

Fixes #2594
---
 cmd/main.go  | 11 ++++++++---
 cmd/utils.go | 10 ++++++++++
 main.go      | 16 ++++++++++++++++
 3 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/cmd/main.go b/cmd/main.go
index 8f7ced10e..503f7d42c 100644
--- a/cmd/main.go
+++ b/cmd/main.go
@@ -140,6 +140,11 @@ func checkMainSyntax(c *cli.Context) {
 	}
 }
 
+// Global profiler to be cleanly set and saved inside graceful shutdown.
+var globalProfiler interface {
+	Stop()
+}
+
 // Main - main for minio server.
 func Main() {
 	app := registerApp()
@@ -181,11 +186,11 @@ func Main() {
 	// Enable profiler if ``MINIO_PROFILER`` is set. Supported options are [cpu, mem, block].
 	switch os.Getenv("MINIO_PROFILER") {
 	case "cpu":
-		defer profile.Start(profile.CPUProfile, profile.ProfilePath(profileDir)).Stop()
+		globalProfiler = profile.Start(profile.CPUProfile, profile.ProfilePath(profileDir))
 	case "mem":
-		defer profile.Start(profile.MemProfile, profile.ProfilePath(profileDir)).Stop()
+		globalProfiler = profile.Start(profile.MemProfile, profile.ProfilePath(profileDir))
 	case "block":
-		defer profile.Start(profile.BlockProfile, profile.ProfilePath(profileDir)).Stop()
+		globalProfiler = profile.Start(profile.BlockProfile, profile.ProfilePath(profileDir))
 	}
 
 	// Run the app - exit on error.
diff --git a/cmd/utils.go b/cmd/utils.go
index d79577acd..2067e977f 100644
--- a/cmd/utils.go
+++ b/cmd/utils.go
@@ -201,8 +201,18 @@ func startMonitorShutdownSignal(onExitFn onExitFunc) error {
 					if err != nil {
 						errorIf(errors.New("Unable to reboot."), err.Error())
 					}
+					// Successfully forked.
 					onExitFn(int(exitSuccess))
 				}
+
+				// Enable profiler if ``MINIO_PROFILER`` is set.
+				switch os.Getenv("MINIO_PROFILER") {
+				case "cpu", "mem", "block":
+					// Stop any running profiler.
+					globalProfiler.Stop()
+				}
+
+				// Exit as success if no errors.
 				onExitFn(int(exitSuccess))
 			}
 		}
diff --git a/main.go b/main.go
index 03fd0d366..532cb448b 100644
--- a/main.go
+++ b/main.go
@@ -1,3 +1,19 @@
+/*
+ * Minio Cloud Storage, (C) 2016 Minio, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 package main
 
 import minio "github.com/minio/minio/cmd"