From 921d4ab784bde24c849f0d93b4c2e615fc4a590c Mon Sep 17 00:00:00 2001
From: Lyle Franklin <lfranklin@pivotal.io>
Date: Mon, 3 Dec 2018 17:38:07 -0800
Subject: [PATCH] Allow Start to be called immediately after Stop

- Prior to this change the library would panic if you called Start a
  second time immediately after calling Stop
- This change updates Stop to block until the background Listen method is
  finished

Signed-off-by: Genevieve Lesperance <glesperance@pivotal.io>
---
 writer.go      |  6 ++++--
 writer_test.go | 11 +++++++++++
 2 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/writer.go b/writer.go
index 2639887..c58f231 100644
--- a/writer.go
+++ b/writer.go
@@ -86,7 +86,7 @@ func (w *Writer) Flush() error {
 func (w *Writer) Start() {
 	if w.ticker == nil {
 		w.ticker = time.NewTicker(w.RefreshInterval)
-		w.tdone = make(chan bool, 1)
+		w.tdone = make(chan bool)
 	}
 
 	go w.Listen()
@@ -95,7 +95,8 @@ func (w *Writer) Start() {
 // Stop stops the listener that updates the terminal
 func (w *Writer) Stop() {
 	w.Flush()
-	close(w.tdone)
+	w.tdone <- true
+	<-w.tdone
 }
 
 // Listen listens for updates to the writer's buffer and flushes to the out provided. It blocks the runtime.
@@ -111,6 +112,7 @@ func (w *Writer) Listen() {
 			w.ticker.Stop()
 			w.ticker = nil
 			w.mtx.Unlock()
+			close(w.tdone)
 			return
 		}
 	}
diff --git a/writer_test.go b/writer_test.go
index 5ed200f..dd12617 100644
--- a/writer_test.go
+++ b/writer_test.go
@@ -22,3 +22,14 @@ func TestWriter(t *testing.T) {
 		t.Fatalf("want %q, got %q", want, b.String())
 	}
 }
+
+func TestStartCalledTwice(t *testing.T) {
+	w := New()
+	b := &bytes.Buffer{}
+	w.Out = b
+
+	w.Start()
+	w.Stop()
+	w.Start()
+	w.Stop()
+}