Improve ssh handover logic and tests (#3757)

## Changes
With this, instead of doing a handover every N seconds flat, we first
wait N seconds, then do a handover for however long it takes to
complete, then wait N seconds again, etc.

In tests we can now better control when handover should be performed,
and the actual logic behaves better with small handover timeouts. Tests
now don't use any hardcoded timeouts for initiating handover and should
be more reliable.

## Tests
Tested with `-count 1000 -parallel 100` and `-race` flags.

Author: ilia-db

experimental/ssh/internal/client/client.go

@@ -264,7 +264,10 @@ func runSSHProxy(ctx context.Context, client *databricks.WorkspaceClient, server
 	createConn := func(ctx context.Context, connID string) (*websocket.Conn, error) {
 		return createWebsocketConnection(ctx, client, connID, opts.ClusterID, serverPort)
 	}
-	return proxy.RunClientProxy(ctx, os.Stdin, os.Stdout, opts.HandoverTimeout, createConn)
+	requestHandoverTick := func() <-chan time.Time {
+		return time.After(opts.HandoverTimeout)
+	}
+	return proxy.RunClientProxy(ctx, os.Stdin, os.Stdout, requestHandoverTick, createConn)
 }
 
 func checkClusterState(ctx context.Context, client *databricks.WorkspaceClient, clusterID string, autoStart bool) error {

experimental/ssh/internal/proxy/client.go

@@ -10,7 +10,7 @@ import (
 	"golang.org/x/sync/errgroup"
 )
 
-func RunClientProxy(ctx context.Context, src io.Reader, dst io.Writer, handoverTimeout time.Duration, createConn createWebsocketConnectionFunc) error {
+func RunClientProxy(ctx context.Context, src io.Reader, dst io.Writer, requestHandoverTick func() <-chan time.Time, createConn createWebsocketConnectionFunc) error {
 	proxy := newProxyConnection(createConn)
 	cmdio.LogString(ctx, "Establishing SSH proxy connection...")
 	g, gCtx := errgroup.WithContext(ctx)
@@ -20,16 +20,12 @@ func RunClientProxy(ctx context.Context, src io.Reader, dst io.Writer, handoverT
 	defer proxy.close()
 	cmdio.LogString(ctx, "SSH proxy connection established")
 
-	cmdio.LogString(ctx, fmt.Sprintf("Connection handover timeout: %v", handoverTimeout))
-	handoverTicker := time.NewTicker(handoverTimeout)
-	defer handoverTicker.Stop()
-
 	g.Go(func() error {
 		for {
 			select {
 			case <-gCtx.Done():
 				return gCtx.Err()
-			case <-handoverTicker.C:
+			case <-requestHandoverTick():
 				err := proxy.initiateHandover(gCtx)
 				if err != nil {
 					return err

experimental/ssh/internal/proxy/client_server_test.go

@@ -10,7 +10,6 @@ import (
 	"io"
 	"net/http/httptest"
 	"os/exec"
-	"strings"
 	"sync"
 	"testing"
 	"time"
@@ -30,21 +29,23 @@ func createTestServer(t *testing.T, maxClients int, shutdownDelay time.Duration)
 	return httptest.NewServer(proxyServer)
 }
 
-func createTestClient(t *testing.T, serverURL string, handoverTimeout time.Duration, errChan, createConnChan chan error) (io.WriteCloser, *testBuffer) {
+func createTestClient(t *testing.T, serverURL string, requestHandoverTick func() <-chan time.Time, errChan chan error) (io.WriteCloser, *testBuffer) {
 	ctx := t.Context()
 	clientInput, clientInputWriter := io.Pipe()
 	clientOutput := newTestBuffer(t)
 	wsURL := "ws" + serverURL[4:]
 	createConn := func(ctx context.Context, connID string) (*websocket.Conn, error) {
 		url := fmt.Sprintf("%s?id=%s", wsURL, connID)
 		conn, _, err := websocket.DefaultDialer.Dial(url, nil) // nolint:bodyclose
-		if createConnChan != nil {
-			createConnChan <- err
-		}
 		return conn, err
 	}
+	if requestHandoverTick == nil {
+		requestHandoverTick = func() <-chan time.Time {
+			return time.After(time.Hour)
+		}
+	}
 	go func() {
-		err := RunClientProxy(ctx, clientInput, clientOutput, handoverTimeout, createConn)
+		err := RunClientProxy(ctx, clientInput, clientOutput, requestHandoverTick, createConn)
 		if err != nil && !isNormalClosure(err) && !errors.Is(err, context.Canceled) {
 			if errChan != nil {
 				errChan <- err
@@ -59,7 +60,7 @@ func createTestClient(t *testing.T, serverURL string, handoverTimeout time.Durat
 func TestClientServerEcho(t *testing.T) {
 	server := createTestServer(t, 2, time.Hour)
 	defer server.Close()
-	clientInputWriter, clientOutput := createTestClient(t, server.URL, time.Hour, nil, nil)
+	clientInputWriter, clientOutput := createTestClient(t, server.URL, nil, nil)
 	defer clientInputWriter.Close()
 
 	testMsg1 := []byte("test message 1\n")
@@ -81,9 +82,9 @@ func TestClientServerEcho(t *testing.T) {
 func TestMultipleClients(t *testing.T) {
 	server := createTestServer(t, 2, time.Hour)
 	defer server.Close()
-	clientInputWriter1, clientOutput1 := createTestClient(t, server.URL, time.Hour, nil, nil)
+	clientInputWriter1, clientOutput1 := createTestClient(t, server.URL, nil, nil)
 	defer clientInputWriter1.Close()
-	clientInputWriter2, clientOutput2 := createTestClient(t, server.URL, time.Hour, nil, nil)
+	clientInputWriter2, clientOutput2 := createTestClient(t, server.URL, nil, nil)
 	defer clientInputWriter2.Close()
 
 	messageCount := 10
@@ -113,9 +114,9 @@ func TestMaxClients(t *testing.T) {
 	maxClients := 2
 	server := createTestServer(t, maxClients, time.Hour)
 	defer server.Close()
-	clientInputWriter1, clientOutput1 := createTestClient(t, server.URL, time.Hour, nil, nil)
+	clientInputWriter1, clientOutput1 := createTestClient(t, server.URL, nil, nil)
 	defer clientInputWriter1.Close()
-	clientInputWriter2, clientOutput2 := createTestClient(t, server.URL, time.Hour, nil, nil)
+	clientInputWriter2, clientOutput2 := createTestClient(t, server.URL, nil, nil)
 	defer clientInputWriter2.Close()
 
 	testMsg1 := []byte("test message 1\n")
@@ -129,7 +130,7 @@ func TestMaxClients(t *testing.T) {
 	require.NoError(t, err)
 
 	errChan := make(chan error, 1)
-	clientInputWriter3, _ := createTestClient(t, server.URL, time.Hour, errChan, nil)
+	clientInputWriter3, _ := createTestClient(t, server.URL, nil, errChan)
 	defer clientInputWriter3.Close()
 	select {
 	case err = <-errChan:
@@ -143,60 +144,35 @@ func TestHandover(t *testing.T) {
 	server := createTestServer(t, 2, time.Hour)
 	defer server.Close()
 
-	maxHandoverCount := 3
-	handoverTimeout := 500 * time.Millisecond
-	createConnChan := make(chan error, 1)
-	clientInputWriter, clientOutput := createTestClient(t, server.URL, handoverTimeout, nil, createConnChan)
+	handoverChan := make(chan time.Time)
+	requestHandoverTick := func() <-chan time.Time {
+		return handoverChan
+	}
+	clientInputWriter, clientOutput := createTestClient(t, server.URL, requestHandoverTick, nil)
 	defer clientInputWriter.Close()
 
-	messageCount := 0
 	expectedOutput := ""
-	sendMessage := func() {
-		message := fmt.Appendf(nil, "message %d\n", messageCount)
-		_, err := clientInputWriter.Write(message)
-		if err != nil {
-			t.Errorf("failed to write message %d: %v", messageCount, err)
-		}
-		messageCount++
-		if messageCount > TOTAL_MESSAGE_COUNT {
-			t.Errorf("exceeded total message count, test buffer won't work correctly")
-		}
-		expectedOutput += string(message)
-	}
 
 	wg := sync.WaitGroup{}
 	wg.Go(func() {
-		handoverCount := 0
-		for {
-			select {
-			case <-createConnChan:
-				sendMessage()
-				handoverCount++
-				if handoverCount >= maxHandoverCount {
-					return
-				}
-			default:
-				sendMessage()
-				time.Sleep(time.Millisecond)
+		for i := range TOTAL_MESSAGE_COUNT {
+			if i > 0 && i%MESSAGES_PER_CHUNK == 0 {
+				handoverChan <- time.Now()
 			}
+			message := fmt.Appendf(nil, "message %d\n", i)
+			_, err := clientInputWriter.Write(message)
+			if err != nil {
+				t.Errorf("failed to write message %d: %v", i, err)
+			}
+			expectedOutput += string(message)
 		}
 	})
 
-	wg.Wait()
+	err := clientOutput.WaitForWrite(fmt.Appendf(nil, "message %d\n", TOTAL_MESSAGE_COUNT-1))
+	require.NoError(t, err, "failed to receive the last message (%d)", TOTAL_MESSAGE_COUNT-1)
 
-	for i := 0; i < messageCount; {
-		// Client can receive multiple echo messages in one response,
-		// so we split them again and verify each one.
-		data, err := clientOutput.WaitForWrite()
-		require.NoError(t, err, "failed to receive message %d", i)
-		lines := strings.SplitSeq(string(data), "\n")
-		for line := range lines {
-			if line != "" {
-				assert.Equal(t, fmt.Sprintf("message %d\n", i), line+"\n")
-				i++
-			}
-		}
-	}
+	wg.Wait()
 
+	// clientOutput is created by appending incoming messages as they arrive, so we are also test correct order here
 	assert.Equal(t, expectedOutput, clientOutput.String())
 }

experimental/ssh/internal/proxy/proxy_test.go

@@ -8,6 +8,7 @@ import (
 	"io"
 	"net/http"
 	"net/http/httptest"
+	"sync"
 	"testing"
 	"time"
 
@@ -18,6 +19,7 @@ import (
 
 type testBuffer struct {
 	t       *testing.T
+	m       sync.Mutex
 	buff    *bytes.Buffer
 	OnWrite chan []byte
 }
@@ -31,6 +33,7 @@ const (
 func newTestBuffer(t *testing.T) *testBuffer {
 	return &testBuffer{
 		t:       t,
+		m:       sync.Mutex{},
 		buff:    new(bytes.Buffer),
 		OnWrite: make(chan []byte, TOTAL_MESSAGE_COUNT),
 	}
@@ -45,7 +48,9 @@ func (tb *testBuffer) Read(p []byte) (n int, err error) {
 }
 
 func (tb *testBuffer) Write(p []byte) (n int, err error) {
+	tb.m.Lock()
 	n, err = tb.buff.Write(p)
+	tb.m.Unlock()
 	require.NoError(tb.t, err)
 	tb.OnWrite <- p
 	return n, err
@@ -61,12 +66,22 @@ func (tb *testBuffer) AssertWrite(expected []byte) error {
 	}
 }
 
-func (tb *testBuffer) WaitForWrite() ([]byte, error) {
-	select {
-	case data := <-tb.OnWrite:
-		return data, nil
-	case <-time.After(2 * time.Second):
-		return nil, errors.New("timeout waiting for write")
+func (tb *testBuffer) Contains(data []byte) bool {
+	tb.m.Lock()
+	defer tb.m.Unlock()
+	return bytes.Contains(tb.buff.Bytes(), data)
+}
+
+func (tb *testBuffer) WaitForWrite(expected []byte) error {
+	for {
+		select {
+		case <-tb.OnWrite:
+			if tb.Contains(expected) {
+				return nil
+			}
+		case <-time.After(2 * time.Second):
+			return errors.New("timeout waiting for write")
+		}
 	}
 }