reset output colors before exiting when interrupt encountered (#1842)

go/cmd/dolt/cli/stdio.go

@@ -19,6 +19,7 @@ import (
 	"io"
 	"os"
 	"path/filepath"
+	"sync/atomic"
 
 	"github.com/dolthub/dolt/go/libraries/utils/iohelp"
 
@@ -26,6 +27,19 @@ import (
 	"github.com/google/uuid"
 )
 
+var outputClosed uint64
+
+func CloseOutput() {
+	if atomic.CompareAndSwapUint64(&outputClosed, 0, 1) {
+		fmt.Fprintln(CliOut)
+	}
+}
+
+func outputIsClosed() bool {
+	isClosed := atomic.LoadUint64(&outputClosed)
+	return isClosed == 1
+}
+
 var CliOut = color.Output
 var CliErr = color.Error
 
@@ -80,30 +94,58 @@ func InitIO() (restoreIO func()) {
 }
 
 func Println(a ...interface{}) {
+	if outputIsClosed() {
+		return
+	}
+
 	fmt.Fprintln(CliOut, a...)
 }
 
 func Print(a ...interface{}) {
+	if outputIsClosed() {
+		return
+	}
+
 	fmt.Fprint(CliOut, a...)
 }
 
 func Printf(format string, a ...interface{}) {
+	if outputIsClosed() {
+		return
+	}
+
 	fmt.Fprintf(CliOut, format, a...)
 }
 
 func PrintErrln(a ...interface{}) {
+	if outputIsClosed() {
+		return
+	}
+
 	fmt.Fprintln(CliErr, a...)
 }
 
 func PrintErr(a ...interface{}) {
+	if outputIsClosed() {
+		return
+	}
+
 	fmt.Fprint(CliErr, a...)
 }
 
 func PrintErrf(format string, a ...interface{}) {
+	if outputIsClosed() {
+		return
+	}
+
 	fmt.Fprintf(CliErr, format, a...)
 }
 
 func DeleteAndPrint(prevMsgLen int, msg string) int {
+	if outputIsClosed() {
+		return 0
+	}
+
 	msgLen := len(msg)
 	backspacesAndMsg := make([]byte, prevMsgLen+msgLen, 2*prevMsgLen+msgLen)
 	for i := 0; i < prevMsgLen; i++ {

go/cmd/dolt/dolt.go

@@ -16,14 +16,15 @@ package main
 
 import (
 	"context"
-	"fmt"
 	"log"
 	"net/http"
 	_ "net/http/pprof"
 	"os"
 	"os/exec"
+	"os/signal"
 	"strconv"
 	"strings"
+	"syscall"
 
 	"github.com/fatih/color"
 	"github.com/opentracing/opentracing-go"
@@ -111,7 +112,20 @@ const traceProf = "trace"
 
 const featureVersionFlag = "--feature-version"
 
+func ResetColorHandler() {
+	c := make(chan os.Signal)
+	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
+	go func() {
+		<-c
+		cli.CloseOutput()
+		color.Unset()
+		os.Exit(0)
+	}()
+}
+
 func main() {
+	ResetColorHandler()
+
 	os.Exit(runMain())
 }
 
@@ -126,16 +140,16 @@ func runMain() int {
 			case profFlag:
 				switch args[1] {
 				case cpuProf:
-					fmt.Println("cpu profiling enabled.")
+					cli.Println("cpu profiling enabled.")
 					defer profile.Start(profile.CPUProfile).Stop()
 				case memProf:
-					fmt.Println("mem profiling enabled.")
+					cli.Println("mem profiling enabled.")
 					defer profile.Start(profile.MemProfile).Stop()
 				case blockingProf:
-					fmt.Println("block profiling enabled")
+					cli.Println("block profiling enabled")
 					defer profile.Start(profile.BlockProfile).Stop()
 				case traceProf:
-					fmt.Println("trace profiling enabled")
+					cli.Println("trace profiling enabled")
 					defer profile.Start(profile.TraceProfile).Stop()
 				default:
 					panic("Unexpected prof flag: " + args[1])
@@ -165,7 +179,7 @@ func runMain() int {
 			//    jaegertracing/all-in-one:1.21
 			// and browse to http://localhost:16686
 			case jaegerFlag:
-				fmt.Println("running with jaeger tracing reporting to localhost")
+				cli.Println("running with jaeger tracing reporting to localhost")
 				transport := transport.NewHTTPTransport("http://localhost:14268/api/traces?format=jaeger.thrift", transport.HTTPBatchSize(128000))
 				reporter := jaeger.NewRemoteReporter(transport)
 				tracer, closer := jaeger.NewTracer("dolt", jaeger.NewConstSampler(true), reporter)
@@ -187,7 +201,7 @@ func runMain() int {
 
 			case stdInFlag:
 				stdInFile := args[1]
-				fmt.Println("Using file contents as stdin:", stdInFile)
+				cli.Println("Using file contents as stdin:", stdInFile)
 
 				f, err := os.Open(stdInFile)