From 238ec05f2f9a1e418f0ad0b8471695e20b8f090e Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Fri, 3 Oct 2014 13:02:16 -0700
Subject: [PATCH] Fix issue where output could be truncated

---
 main.go | 33 ++++++++++++++++++++++++++++-----
 1 file changed, 28 insertions(+), 5 deletions(-)

diff --git a/main.go b/main.go
index 7f1dde2ce..29ee63743 100644
--- a/main.go
+++ b/main.go
@@ -6,6 +6,7 @@ import (
 	"io/ioutil"
 	"log"
 	"os"
+	"sync"
 
 	"github.com/hashicorp/terraform/plugin"
 	"github.com/mitchellh/cli"
@@ -47,8 +48,9 @@ func realMain() int {
 
 		// Setup the prefixed readers that send data properly to
 		// stdout/stderr.
+		doneCh := make(chan struct{})
 		outR, outW := io.Pipe()
-		go copyOutput(outR)
+		go copyOutput(outR, doneCh)
 
 		// Create the configuration for panicwrap and wrap our executable
 		wrapConfig.Handler = panicHandler(logTempFile)
@@ -62,6 +64,12 @@ func realMain() int {
 
 		// If >= 0, we're the parent, so just exit
 		if exitStatus >= 0 {
+			// Close the stdout writer so that our copy process can finish
+			outW.Close()
+
+			// Wait for the output copying to finish
+			<-doneCh
+
 			return exitStatus
 		}
 
@@ -175,7 +183,9 @@ func cliConfigFile() (string, error) {
 // copyOutput uses output prefixes to determine whether data on stdout
 // should go to stdout or stderr. This is due to panicwrap using stderr
 // as the log and error channel.
-func copyOutput(r io.Reader) {
+func copyOutput(r io.Reader, doneCh chan<- struct{}) {
+	defer close(doneCh)
+
 	pr, err := prefixedio.NewReader(r)
 	if err != nil {
 		panic(err)
@@ -194,7 +204,20 @@ func copyOutput(r io.Reader) {
 		panic(err)
 	}
 
-	go io.Copy(os.Stderr, stderrR)
-	go io.Copy(os.Stdout, stdoutR)
-	go io.Copy(os.Stdout, defaultR)
+	var wg sync.WaitGroup
+	wg.Add(3)
+	go func() {
+		defer wg.Done()
+		io.Copy(os.Stderr, stderrR)
+	}()
+	go func() {
+		defer wg.Done()
+		io.Copy(os.Stdout, stdoutR)
+	}()
+	go func() {
+		defer wg.Done()
+		io.Copy(os.Stdout, defaultR)
+	}()
+
+	wg.Wait()
 }