From 9309fe431f417dc9ceb93cac99f7a4d49a16f67d Mon Sep 17 00:00:00 2001
From: Daniel <dhaavi@users.noreply.github.com>
Date: Fri, 29 May 2020 15:30:28 +0200
Subject: [PATCH] Improve -print-stack-on-exit

The stack is now also printed when there is an error during startup
---
 run/main.go | 31 ++++++++++++++++++++-----------
 1 file changed, 20 insertions(+), 11 deletions(-)

diff --git a/run/main.go b/run/main.go
index 903519e..9050df5 100644
--- a/run/main.go
+++ b/run/main.go
@@ -4,6 +4,7 @@ import (
 	"bufio"
 	"flag"
 	"fmt"
+	"io"
 	"os"
 	"os/signal"
 	"runtime/pprof"
@@ -36,6 +37,10 @@ func Run() int {
 			return 0
 		}
 
+		if printStackOnExit {
+			printStackTo(os.Stdout)
+		}
+
 		_ = modules.Shutdown()
 		return modules.GetExitStatusCode()
 	}
@@ -78,20 +83,13 @@ signalLoop:
 			}()
 
 			if printStackOnExit {
-				fmt.Println("=== PRINTING TRACES ===")
-				fmt.Println("=== GOROUTINES ===")
-				_ = pprof.Lookup("goroutine").WriteTo(os.Stdout, 1)
-				fmt.Println("=== BLOCKING ===")
-				_ = pprof.Lookup("block").WriteTo(os.Stdout, 1)
-				fmt.Println("=== MUTEXES ===")
-				_ = pprof.Lookup("mutex").WriteTo(os.Stdout, 1)
-				fmt.Println("=== END TRACES ===")
+				printStackTo(os.Stdout)
 			}
 
 			go func() {
-				time.Sleep(60 * time.Second)
-				fmt.Fprintln(os.Stderr, "===== TAKING TOO LONG FOR SHUTDOWN - PRINTING STACK TRACES =====")
-				_ = pprof.Lookup("goroutine").WriteTo(os.Stderr, 1)
+				time.Sleep(3 * time.Minute)
+				fmt.Fprintln(os.Stderr, "===== TAKING TOO LONG FOR SHUTDOWN =====")
+				printStackTo(os.Stderr)
 				os.Exit(1)
 			}()
 
@@ -124,3 +122,14 @@ func inputSignals(signalCh chan os.Signal) {
 		}
 	}
 }
+
+func printStackTo(writer io.Writer) {
+	fmt.Fprintln(writer, "=== PRINTING TRACES ===")
+	fmt.Fprintln(writer, "=== GOROUTINES ===")
+	_ = pprof.Lookup("goroutine").WriteTo(writer, 1)
+	fmt.Fprintln(writer, "=== BLOCKING ===")
+	_ = pprof.Lookup("block").WriteTo(writer, 1)
+	fmt.Fprintln(writer, "=== MUTEXES ===")
+	_ = pprof.Lookup("mutex").WriteTo(writer, 1)
+	fmt.Fprintln(writer, "=== END TRACES ===")
+}