Merge pull request #21139 from hashicorp/svh/b-stdin

Make sure UIInput keeps working after being canceled

command/ui_input.go

@@ -12,6 +12,7 @@ import (
 	"os/signal"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"unicode"
 
 	"github.com/hashicorp/terraform/terraform"
@@ -30,10 +31,13 @@ type UIInput struct {
 	Colorize *colorstring.Colorize
 
 	// Reader and Writer for IO. If these aren't set, they will default to
-	// Stdout and Stderr respectively.
+	// Stdin and Stdout respectively.
 	Reader io.Reader
 	Writer io.Writer
 
+	listening int32
+	result    chan string
+
 	interrupted bool
 	l           sync.Mutex
 	once        sync.Once
@@ -117,20 +121,24 @@ func (i *UIInput) Input(ctx context.Context, opts *terraform.InputOpts) (string,
 	}
 
 	// Listen for the input in a goroutine. This will allow us to
-	// interrupt this if we are interrupted (SIGINT)
-	result := make(chan string, 1)
+	// interrupt this if we are interrupted (SIGINT).
 	go func() {
+		if !atomic.CompareAndSwapInt32(&i.listening, 0, 1) {
+			return // We are already listening for input.
+		}
+		defer atomic.CompareAndSwapInt32(&i.listening, 1, 0)
+
 		buf := bufio.NewReader(r)
 		line, err := buf.ReadString('\n')
 		if err != nil {
 			log.Printf("[ERR] UIInput scan err: %s", err)
 		}
 
-		result <- strings.TrimRightFunc(line, unicode.IsSpace)
+		i.result <- strings.TrimRightFunc(line, unicode.IsSpace)
 	}()
 
 	select {
-	case line := <-result:
+	case line := <-i.result:
 		fmt.Fprint(w, "\n")
 
 		if line == "" {
@@ -157,6 +165,8 @@ func (i *UIInput) Input(ctx context.Context, opts *terraform.InputOpts) (string,
 }
 
 func (i *UIInput) init() {
+	i.result = make(chan string)
+
 	if i.Colorize == nil {
 		i.Colorize = &colorstring.Colorize{
 			Colors:  colorstring.DefaultColors,

command/ui_input_test.go

@@ -3,7 +3,11 @@ package command
 import (
 	"bytes"
 	"context"
+	"fmt"
+	"io"
+	"sync/atomic"
 	"testing"
+	"time"
 
 	"github.com/hashicorp/terraform/terraform"
 )
@@ -20,11 +24,61 @@ func TestUIInputInput(t *testing.T) {
 
 	v, err := i.Input(context.Background(), &terraform.InputOpts{})
 	if err != nil {
-		t.Fatalf("err: %s", err)
+		t.Fatalf("unexpected error: %v", err)
 	}
 
 	if v != "foo" {
-		t.Fatalf("bad: %#v", v)
+		t.Fatalf("unexpected input: %s", v)
+	}
+}
+
+func TestUIInputInput_canceled(t *testing.T) {
+	r, w := io.Pipe()
+	i := &UIInput{
+		Reader: r,
+		Writer: bytes.NewBuffer(nil),
+	}
+
+	// Make a context that can be canceled.
+	ctx, cancel := context.WithCancel(context.Background())
+
+	go func() {
+		// Cancel the context after 2 seconds.
+		time.Sleep(2 * time.Second)
+		cancel()
+	}()
+
+	// Get input until the context is canceled.
+	v, err := i.Input(ctx, &terraform.InputOpts{})
+	if err != context.Canceled {
+		t.Fatalf("expected a context.Canceled error, got: %v", err)
+	}
+
+	// As the context was canceled v should be empty.
+	if v != "" {
+		t.Fatalf("unexpected input: %s", v)
+	}
+
+	// As the context was canceled we should still be listening.
+	listening := atomic.LoadInt32(&i.listening)
+	if listening != 1 {
+		t.Fatalf("expected listening to be 1, got: %d", listening)
+	}
+
+	go func() {
+		// Fake input is given after 1 second.
+		time.Sleep(time.Second)
+		fmt.Fprint(w, "foo\n")
+		w.Close()
+	}()
+
+	v, err = i.Input(context.Background(), &terraform.InputOpts{})
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	if v != "foo" {
+		t.Fatalf("unexpected input: %s", v)
 	}
 }
 
@@ -36,10 +90,10 @@ func TestUIInputInput_spaces(t *testing.T) {
 
 	v, err := i.Input(context.Background(), &terraform.InputOpts{})
 	if err != nil {
-		t.Fatalf("err: %s", err)
+		t.Fatalf("unexpected error: %v", err)
 	}
 
 	if v != "foo bar" {
-		t.Fatalf("bad: %#v", v)
+		t.Fatalf("unexpected input: %s", v)
 	}
 }