Fixed infinite loop in ExpectProcess.ExpectFunc

Signed-off-by: Vitalii Levitskii <vitalii@uber.com>

pkg/expect/expect.go

@@ -18,6 +18,7 @@ package expect
 
 import (
 	"bufio"
+	"context"
 	"fmt"
 	"io"
 	"os"
@@ -96,7 +97,7 @@ func (ep *ExpectProcess) read() {
 }
 
 // ExpectFunc returns the first line satisfying the function f.
-func (ep *ExpectProcess) ExpectFunc(f func(string) bool) (string, error) {
+func (ep *ExpectProcess) ExpectFunc(ctx context.Context, f func(string) bool) (string, error) {
 	i := 0
 
 	for {
@@ -114,7 +115,13 @@ func (ep *ExpectProcess) ExpectFunc(f func(string) bool) (string, error) {
 			break
 		}
 		ep.mu.Unlock()
-		time.Sleep(time.Millisecond * 10)
+
+		select {
+		case <-ctx.Done():
+			return "", fmt.Errorf("failed to find match string: %w", ctx.Err())
+		case <-time.After(time.Millisecond * 10):
+			// continue loop
+		}
 	}
 	ep.mu.Lock()
 	lastLinesIndex := len(ep.lines) - DEBUG_LINES_TAIL
@@ -128,9 +135,15 @@ func (ep *ExpectProcess) ExpectFunc(f func(string) bool) (string, error) {
 		ep.err, lastLines)
 }
 
+// ExpectWithContext returns the first line containing the given string.
+func (ep *ExpectProcess) ExpectWithContext(ctx context.Context, s string) (string, error) {
+	return ep.ExpectFunc(ctx, func(txt string) bool { return strings.Contains(txt, s) })
+}
+
 // Expect returns the first line containing the given string.
+// Deprecated: please use ExpectWithContext instead.
 func (ep *ExpectProcess) Expect(s string) (string, error) {
-	return ep.ExpectFunc(func(txt string) bool { return strings.Contains(txt, s) })
+	return ep.ExpectWithContext(context.Background(), s)
 }
 
 // LineCount returns the number of recorded lines since

pkg/expect/expect_test.go

@@ -17,9 +17,12 @@
 package expect
 
 import (
+	"context"
 	"os"
 	"testing"
 	"time"
+
+	"github.com/stretchr/testify/require"
 )
 
 func TestExpectFunc(t *testing.T) {
@@ -28,7 +31,7 @@ func TestExpectFunc(t *testing.T) {
 		t.Fatal(err)
 	}
 	wstr := "hello world\r\n"
-	l, eerr := ep.ExpectFunc(func(a string) bool { return len(a) > 10 })
+	l, eerr := ep.ExpectFunc(context.Background(), func(a string) bool { return len(a) > 10 })
 	if eerr != nil {
 		t.Fatal(eerr)
 	}
@@ -40,6 +43,33 @@ func TestExpectFunc(t *testing.T) {
 	}
 }
 
+func TestExpectFuncTimeout(t *testing.T) {
+	ep, err := NewExpect("tail", "-f", "/dev/null")
+	if err != nil {
+		t.Fatal(err)
+	}
+	go func() {
+		// It's enough to have "talkative" process to stuck in the infinite loop of reading
+		for {
+			err := ep.Send("new line\n")
+			if err != nil {
+				return
+			}
+		}
+	}()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
+	defer cancel()
+
+	_, err = ep.ExpectFunc(ctx, func(a string) bool { return false })
+
+	require.ErrorAs(t, err, &context.DeadlineExceeded)
+
+	if err = ep.Stop(); err != nil {
+		t.Fatal(err)
+	}
+}
+
 func TestEcho(t *testing.T) {
 	ep, err := NewExpect("echo", "hello world")
 	if err != nil {

tests/e2e/ctl_v3_elect_test.go

@@ -15,6 +15,7 @@
 package e2e
 
 import (
+	"context"
 	"os"
 	"strings"
 	"testing"
@@ -106,7 +107,7 @@ func ctlV3Elect(cx ctlCtx, name, proposal string) (*expect.ExpectProcess, <-chan
 		return proc, outc, err
 	}
 	go func() {
-		s, xerr := proc.ExpectFunc(func(string) bool { return true })
+		s, xerr := proc.ExpectFunc(context.TODO(), func(string) bool { return true })
 		if xerr != nil {
 			cx.t.Errorf("expect failed (%v)", xerr)
 		}

tests/e2e/ctl_v3_lock_test.go

@@ -15,6 +15,7 @@
 package e2e
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"strings"
@@ -127,7 +128,7 @@ func ctlV3Lock(cx ctlCtx, name string) (*expect.ExpectProcess, <-chan string, er
 		return proc, outc, err
 	}
 	go func() {
-		s, xerr := proc.ExpectFunc(func(string) bool { return true })
+		s, xerr := proc.ExpectFunc(context.TODO(), func(string) bool { return true })
 		if xerr != nil {
 			cx.t.Errorf("expect failed (%v)", xerr)
 		}

tests/e2e/v3_curl_test.go

@@ -15,6 +15,7 @@
 package e2e
 
 import (
+	"context"
 	"encoding/base64"
 	"encoding/json"
 	"fmt"
@@ -249,7 +250,7 @@ func testV3CurlAuth(cx ctlCtx) {
 		testutil.AssertNil(cx.t, err)
 		defer proc.Close()
 
-		cURLRes, err := proc.ExpectFunc(lineFunc)
+		cURLRes, err := proc.ExpectFunc(context.Background(), lineFunc)
 		testutil.AssertNil(cx.t, err)
 
 		authRes := make(map[string]interface{})

tests/framework/e2e/util.go

@@ -15,6 +15,7 @@
 package e2e
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"math/rand"
@@ -36,7 +37,7 @@ func WaitReadyExpectProc(exproc *expect.ExpectProcess, readyStrs []string) error
 		}
 		return false
 	}
-	_, err := exproc.ExpectFunc(matchSet)
+	_, err := exproc.ExpectFunc(context.Background(), matchSet)
 	return err
 }