From f9e830705d02f43284888f21547bd184399a0ebf Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Tue, 10 Jun 2014 10:28:47 -0700
Subject: [PATCH] Prefixed IO to make sure stdout/stderr match up

---
 commands.go | 11 +++++++++-
 main.go     | 63 +++++++++++++++++++++++++++++------------------------
 2 files changed, 44 insertions(+), 30 deletions(-)

diff --git a/commands.go b/commands.go
index d8739a22b..fa48dfaf8 100644
--- a/commands.go
+++ b/commands.go
@@ -10,8 +10,17 @@ import (
 // Commands is the mapping of all the available Terraform commands.
 var Commands map[string]cli.CommandFactory
 
+const ErrorPrefix = "e:"
+const OutputPrefix = "o:"
+
 func init() {
-	ui := &cli.BasicUi{Writer: os.Stdout}
+	ui := &cli.PrefixedUi{
+		AskPrefix:    OutputPrefix,
+		OutputPrefix: OutputPrefix,
+		InfoPrefix:   OutputPrefix,
+		ErrorPrefix:  ErrorPrefix,
+		Ui:           &cli.BasicUi{Writer: os.Stdout},
+	}
 
 	Commands = map[string]cli.CommandFactory{
 		"apply": func() (cli.Command, error) {
diff --git a/main.go b/main.go
index f0781bbc3..89fccc73a 100644
--- a/main.go
+++ b/main.go
@@ -2,14 +2,15 @@ package main
 
 import (
 	"fmt"
+	"io"
 	"io/ioutil"
 	"log"
 	"os"
 
-	"github.com/ActiveState/tail"
 	"github.com/hashicorp/terraform/plugin"
 	"github.com/mitchellh/cli"
 	"github.com/mitchellh/panicwrap"
+	"github.com/mitchellh/prefixedio"
 )
 
 func main() {
@@ -26,6 +27,9 @@ func realMain() int {
 			fmt.Fprintf(os.Stderr, "Couldn't setup log output: %s", err)
 			return 1
 		}
+		if logWriter == nil {
+			logWriter = ioutil.Discard
+		}
 
 		// We always send logs to a temporary file that we use in case
 		// there is a panic. Otherwise, we delete it.
@@ -34,33 +38,22 @@ func realMain() int {
 			fmt.Fprintf(os.Stderr, "Couldn't setup logging tempfile: %s", err)
 			return 1
 		}
-		logTempFile.Close()
 		defer os.Remove(logTempFile.Name())
+		defer logTempFile.Close()
 
 		// Tell the logger to log to this file
-		os.Setenv(EnvLog, "1")
-		os.Setenv(EnvLogFile, logTempFile.Name())
+		os.Setenv(EnvLog, "")
+		os.Setenv(EnvLogFile, "")
 
-		if logWriter != nil {
-			// Start tailing the file beforehand to get the data
-			t, err := tail.TailFile(logTempFile.Name(), tail.Config{
-				Follow:    true,
-				Logger:    tail.DiscardingLogger,
-				MustExist: true,
-			})
-			if err != nil {
-				fmt.Fprintf(os.Stderr, "Couldn't setup logging tempfile: %s", err)
-				return 1
-			}
-			go func() {
-				for line := range t.Lines {
-					logWriter.Write([]byte(line.Text + "\n"))
-				}
-			}()
-		}
+		// Setup the prefixed readers that send data properly to
+		// stdout/stderr.
+		outR, outW := io.Pipe()
+		go copyOutput(outR)
 
 		// Create the configuration for panicwrap and wrap our executable
 		wrapConfig.Handler = panicHandler(logTempFile)
+		wrapConfig.Writer = io.MultiWriter(logTempFile, logWriter)
+		wrapConfig.Stdout = outW
 		exitStatus, err := panicwrap.Wrap(&wrapConfig)
 		if err != nil {
 			fmt.Fprintf(os.Stderr, "Couldn't start Terraform: %s", err)
@@ -82,14 +75,7 @@ func realMain() int {
 }
 
 func wrappedMain() int {
-	logOutput, err := logOutput()
-	if err != nil {
-		fmt.Fprintf(os.Stderr, "Error starting Terraform: %s", err)
-		return 1
-	}
-	if logOutput != nil {
-		log.SetOutput(logOutput)
-	}
+	log.SetOutput(os.Stderr)
 
 	// Load the configuration
 	config := BuiltinConfig
@@ -127,3 +113,22 @@ func wrappedMain() int {
 
 	return exitCode
 }
+
+func copyOutput(r io.Reader) {
+	pr, err := prefixedio.NewReader(r)
+	if err != nil {
+		panic(err)
+	}
+
+	stderrR, err := pr.Prefix(ErrorPrefix)
+	if err != nil {
+		panic(err)
+	}
+	stdoutR, err := pr.Prefix(OutputPrefix)
+	if err != nil {
+		panic(err)
+	}
+
+	go io.Copy(os.Stderr, stderrR)
+	go io.Copy(os.Stdout, stdoutR)
+}