From 17f22e77a324047d928bc3e142f05fc73e4ded13 Mon Sep 17 00:00:00 2001
From: Jimmy Zelinskie <jimmy.zelinskie+git@gmail.com>
Date: Tue, 3 Jul 2018 23:14:33 -0400
Subject: [PATCH] cmd/chihaya: add tracing flag

---
 cmd/chihaya/main.go | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/cmd/chihaya/main.go b/cmd/chihaya/main.go
index 6958dc3..0b5f27d 100644
--- a/cmd/chihaya/main.go
+++ b/cmd/chihaya/main.go
@@ -6,6 +6,7 @@ import (
 	"os/signal"
 	"runtime"
 	"runtime/pprof"
+	"runtime/trace"
 	"strings"
 	"syscall"
 
@@ -200,6 +201,19 @@ func PreRunCmdFunc(cmd *cobra.Command, args []string) error {
 		log.Info("enabled debug logging")
 	}
 
+	tracePath, err := cmd.Flags().GetString("trace")
+	if err != nil {
+		return err
+	}
+	if tracePath != "" {
+		f, err := os.Create(tracePath)
+		if err != nil {
+			return err
+		}
+		trace.Start(f)
+		log.Info("enabled tracing", log.Fields{"path": tracePath})
+	}
+
 	cpuProfilePath, err := cmd.Flags().GetString("cpuprofile")
 	if err != nil {
 		return err
@@ -219,8 +233,9 @@ func PreRunCmdFunc(cmd *cobra.Command, args []string) error {
 // PostRunCmdFunc handles clean up of any state initialized by command line
 // flags.
 func PostRunCmdFunc(cmd *cobra.Command, args []string) error {
-	// This can be called regardless because it noops when not profiling.
+	// These can be called regardless because it noops when not profiling.
 	pprof.StopCPUProfile()
+	trace.Stop()
 
 	return nil
 }
@@ -237,6 +252,7 @@ func main() {
 
 	rootCmd.Flags().String("config", "/etc/chihaya.yaml", "location of configuration file")
 	rootCmd.Flags().String("cpuprofile", "", "location to save a CPU profile")
+	rootCmd.Flags().String("trace", "", "location to save a trace")
 	rootCmd.Flags().Bool("debug", false, "enable debug logging")
 	rootCmd.Flags().Bool("json", false, "enable json logging")
 	if runtime.GOOS == "windows" {