simplfy ssh tunnel logic

Author: jgao54

flow/connectors/mysql/mysql.go

@@ -70,24 +70,21 @@ func NewMySqlConnector(ctx context.Context, config *protos.MySqlConfig) (*MySqlC
 		rdsAuth:          rdsAuth,
 	}
 	c.contexts.Store(&contexts)
+	if ssh != nil {
+		go func() {
+			sshErr := <-ssh.Watch()
+			c.logger.Info("SSH tunnel closed, closing MySQL connection", slog.Any("error", sshErr))
+			if conn := c.conn.Swap(nil); conn != nil {
+				if err := conn.Close(); err != nil {
+					c.logger.Error("Failed to close MySQL connection", slog.Any("error", err))
+				}
+			}
+		}()
+	}
 	go func() { //nolint:gosec // G118: long-lived goroutine, not request-scoped
-		ctx := context.Background()
 		for {
 			var ok bool
 			select {
-			case <-ssh.GetKeepaliveChan(ctx):
-				c.logger.Info("SSH keepalive failed, closing connection")
-				ctx = context.Background()
-				// close the SSH client so that BinlogSyncer notices too
-				if err := ssh.Client.Close(); err != nil {
-					c.logger.Error("Failed to close SSH client", slog.Any("error", err))
-				}
-				if conn := c.conn.Swap(nil); conn != nil {
-					c.logger.Info("Closing connection due to SSH keepalive failure")
-					if err := conn.Close(); err != nil {
-						c.logger.Error("Failed to close MySQL connection", slog.Any("error", err))
-					}
-				}
 			case <-ctx.Done():
 				c.logger.Info("ctx canceled, closing connection")
 				ctx = context.Background()

flow/connectors/mysql/ssh_keepalive_test.go

@@ -110,7 +110,7 @@ func setupMySQLSSHKeepaliveHarness(ctx context.Context, t *testing.T, proxyName
 	t.Helper()
 
 	connector, sshProxy := setupMySQLConnectorWithSSHProxy(ctx, t, proxyName, proxyPort)
-	keepaliveChan := connector.ssh.GetKeepaliveChan(ctx)
+	keepaliveChan := connector.ssh.Watch()
 
 	return connector, utils.SSHKeepaliveTestConfig{
 		SSHProxy:      sshProxy,
@@ -198,7 +198,7 @@ func TestMySQLSSHKeepaliveCDCHang(t *testing.T) {
 	connector, sshProxy := setupMySQLConnectorWithSSHProxy(ctx, t, "my-ssh-cdc-down-test", toxiproxyCDCHangProxyPort)
 	defer connector.Close()
 
-	keepaliveChan := connector.ssh.GetKeepaliveChan(ctx)
+	keepaliveChan := connector.ssh.Watch()
 	require.NotNil(t, keepaliveChan, "SSH keepalive channel should exist")
 
 	req, otelManager := setupCDCPullRecords(ctx, t, connector, "test_ssh_cdc_hang")
@@ -252,7 +252,7 @@ func TestMySQLSSHKeepaliveCDCCloseHang(t *testing.T) {
 	connector, sshProxy := setupMySQLConnectorWithSSHProxy(ctx, t, "my-ssh-cdc-latency-test", toxiproxyCDCCloseHangProxyPort)
 	defer connector.Close()
 
-	keepaliveChan := connector.ssh.GetKeepaliveChan(ctx)
+	keepaliveChan := connector.ssh.Watch()
 	require.NotNil(t, keepaliveChan, "SSH keepalive channel should exist")
 
 	req, otelManager := setupCDCPullRecords(ctx, t, connector, "test_ssh_cdc_close_hang")

flow/connectors/postgres/postgres.go

@@ -129,24 +129,24 @@ func NewPostgresConnector(ctx context.Context, env map[string]string, pgConfig *
 		rdsAuth:                rdsAuth,
 	}
 
-	tunnel.StartKeepalive(context.Background(), func() {
-		connector.logger.Info("SSH keepalive failed, closing connections")
-		bgCtx := context.Background()
-		timeout, cancel := context.WithTimeout(bgCtx, 5*time.Second)
-		defer cancel()
-		if err := connector.conn.Close(timeout); err != nil {
-			connector.logger.Error("Failed to close Postgres connection on SSH keepalive failure",
-				slog.Any("error", err))
-		}
-		if connector.replConn != nil {
-			timeout2, cancel2 := context.WithTimeout(bgCtx, 5*time.Second)
-			defer cancel2()
-			if err := connector.replConn.Close(timeout2); err != nil {
-				connector.logger.Error("Failed to close Postgres replication connection on SSH keepalive failure",
-					slog.Any("error", err))
+	if tunnel != nil {
+		go func() {
+			sshErr := <-tunnel.Watch()
+			connector.logger.Warn("SSH tunnel closed, closing Postgres connections", slog.Any("error", sshErr))
+			timeout, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)
+			defer cancel()
+			if err := connector.conn.Close(timeout); err != nil {
+				connector.logger.Error("failed to close Postgres connection", slog.Any("error", err))
 			}
-		}
-	})
+			if connector.replConn != nil {
+				replTimeout, replCancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)
+				defer replCancel()
+				if err := connector.replConn.Close(replTimeout); err != nil {
+					connector.logger.Error("failed to close Postgres replication connection", slog.Any("error", err))
+				}
+			}
+		}()
+	}
 
 	return connector, nil
 }

flow/connectors/postgres/ssh_keepalive_test.go

@@ -44,7 +44,7 @@ func setupPostgresConnectorWithSSH(ctx context.Context, t *testing.T, proxyName
 	err = connector.ConnectionActive(ctx)
 	require.NoError(t, err, "Initial connection should work")
 
-	keepaliveChan := connector.ssh.GetKeepaliveChan(ctx)
+	keepaliveChan := connector.ssh.Watch()
 
 	return connector, utils.SSHKeepaliveTestConfig{
 		SSHProxy:      sshProxy,

flow/connectors/postgres/ssh_wrapped_conn.go

@@ -18,11 +18,11 @@ func NewPostgresConnFromConfig(
 	connConfig *pgx.ConnConfig,
 	tlsHost string,
 	rdsAuth *utils.RDSAuth,
-	tunnel *utils.SSHTunnel,
+	ssh *utils.SSHTunnel,
 ) (*pgx.Conn, error) {
-	if tunnel != nil && tunnel.Client != nil {
+	if ssh != nil && ssh.Client != nil {
 		connConfig.DialFunc = func(ctx context.Context, network, addr string) (net.Conn, error) {
-			conn, err := tunnel.Client.DialContext(ctx, network, addr)
+			conn, err := ssh.Client.DialContext(ctx, network, addr)
 			if err != nil {
 				return nil, err
 			}

flow/connectors/utils/ssh.go

@@ -3,11 +3,13 @@ package utils
 import (
 	"context"
 	"encoding/base64"
+	"errors"
 	"fmt"
 	"log/slog"
-	"sync/atomic"
+	"sync"
 	"time"
 
+	"go.temporal.io/sdk/log"
 	"golang.org/x/crypto/ssh"
 
 	"github.com/PeerDB-io/peerdb/flow/generated/protos"
@@ -16,13 +18,17 @@ import (
 	"github.com/PeerDB-io/peerdb/flow/shared/exceptions"
 )
 
-const SSHKeepaliveInterval = 15 * time.Second
+const (
+	SSHKeepaliveInterval = 15 * time.Second
+	SSHRequestTimeout    = 10 * time.Second
+)
 
 type SSHTunnel struct {
 	*ssh.Client
-	logger        *slog.Logger
-	keepaliveChan atomic.Pointer[chan struct{}]
-	badTunnel     bool
+	logger    log.Logger
+	watch     chan struct{}
+	err       error
+	closeOnce sync.Once
 }
 
 // GetSSHClientConfig returns an *ssh.ClientConfig based on provided credentials.
@@ -76,128 +82,98 @@ func NewSSHTunnel(
 	ctx context.Context,
 	sshConfig *protos.SSHConfig,
 ) (*SSHTunnel, error) {
-	if sshConfig != nil {
-		logger := internal.LoggerFromCtx(ctx)
-		sshServer := shared.JoinHostPort(sshConfig.Host, sshConfig.Port)
-		clientConfig, err := GetSSHClientConfig(sshConfig)
-		if err != nil {
-			logger.Error("Failed to get SSH client config", slog.Any("error", err))
-			return nil, err
-		}
+	if sshConfig == nil {
+		return nil, nil
+	}
 
-		logger.Info("Setting up SSH connection", slog.String("Server", sshServer))
-		client, err := ssh.Dial("tcp", sshServer, clientConfig)
-		if err != nil {
-			return nil, exceptions.NewSSHTunnelSetupError(err)
-		}
+	logger := internal.LoggerFromCtx(ctx)
+	sshServer := shared.JoinHostPort(sshConfig.Host, sshConfig.Port)
+	clientConfig, err := GetSSHClientConfig(sshConfig)
+	if err != nil {
+		logger.Error("Failed to get SSH client config", slog.Any("error", err))
+		return nil, err
+	}
 
-		return &SSHTunnel{
-			Client: client,
-			logger: internal.SlogLoggerFromCtx(ctx),
-		}, nil
+	logger.Info("Setting up SSH connection", slog.String("Server", sshServer))
+	client, err := ssh.Dial("tcp", sshServer, clientConfig)
+	if err != nil {
+		return nil, exceptions.NewSSHTunnelSetupError(err)
 	}
 
-	return nil, nil
+	tunnel := &SSHTunnel{
+		Client: client,
+		logger: internal.SlogLoggerFromCtx(ctx),
+		watch:  make(chan struct{}),
+	}
+
+	go tunnel.runKeepaliveLoop()
+
+	return tunnel, nil
+}
+
+func (tunnel *SSHTunnel) Watch() <-chan struct{} {
+	if tunnel == nil || tunnel.Client == nil {
+		return nil
+	}
+	return tunnel.watch
 }
 
 func (tunnel *SSHTunnel) Close() error {
-	if tunnel != nil && tunnel.Client != nil {
-		if keepaliveChan := tunnel.keepaliveChan.Swap(nil); keepaliveChan != nil {
-			close(*keepaliveChan)
-		}
-		tunnel.badTunnel = true
-		return tunnel.Client.Close()
+	if tunnel == nil || tunnel.Client == nil {
+		return nil
 	}
-	return nil
+	tunnel.shutdown(nil)
+	return tunnel.err
+}
+
+func (tunnel *SSHTunnel) shutdown(err error) {
+	tunnel.closeOnce.Do(func() {
+		tunnel.err = err
+		close(tunnel.watch)
+		if closeErr := tunnel.Client.Close(); closeErr != nil && tunnel.err == nil {
+			tunnel.err = closeErr
+		}
+	})
 }
 
-func (tunnel *SSHTunnel) runKeepaliveLoop(
-	ctx context.Context, stopChan <-chan struct{}, onFailure func(),
-) {
+func (tunnel *SSHTunnel) runKeepaliveLoop() {
 	ticker := time.NewTicker(SSHKeepaliveInterval)
 	defer ticker.Stop()
-	logger := tunnel.logger
-	// in case request hangs, we want to detect that and not send another request
-	requestSent := atomic.Bool{}
-	var keepaliveErr error
-	// closed by request making goroutine to signal error, keepaliveErr
-	errChan := make(chan struct{})
-
 	for {
 		select {
+		case <-tunnel.watch:
+			// tunnel closed
+			return
 		case <-ticker.C:
-			if requestSent.Load() {
-				// Previous keepalive request didn't return yet, something's wrong
-				logger.ErrorContext(ctx, "Previous keepalive request still pending, marking tunnel as bad")
-				if keepaliveChan := tunnel.keepaliveChan.Swap(nil); keepaliveChan != nil {
-					close(*keepaliveChan)
-				}
-				tunnel.badTunnel = true
-				if onFailure != nil {
-					onFailure()
-				}
+			success := tunnel.sendKeepAlive()
+			if !success {
 				return
 			}
-			go func() {
-				requestSent.Store(true)
-				_, _, err := tunnel.Client.SendRequest("keepalive@openssh.com", true, nil)
-				requestSent.Store(false)
-				if err != nil {
-					keepaliveErr = err
-					close(errChan)
-				}
-			}()
-		case <-ctx.Done():
-			if keepaliveChan := tunnel.keepaliveChan.Swap(nil); keepaliveChan != nil {
-				close(*keepaliveChan)
-			}
-			return
-		case <-stopChan:
-			// channel closed from outside
-			return
-		case <-errChan:
-			logger.ErrorContext(ctx, "Keepalive request failed, marking tunnel as bad", slog.Any("error", keepaliveErr))
-			if keepaliveChan := tunnel.keepaliveChan.Swap(nil); keepaliveChan != nil {
-				close(*keepaliveChan)
-			}
-			tunnel.badTunnel = true
-			if onFailure != nil {
-				onFailure()
-			}
-			return
 		}
 	}
 }
 
-func (tunnel *SSHTunnel) StartKeepalive(ctx context.Context, onFailure func()) {
-	if tunnel == nil || tunnel.Client == nil || tunnel.badTunnel {
-		return
-	}
-	if tunnel.keepaliveChan.Load() != nil {
-		// Already started
-		return
-	}
-	stopChan := make(chan struct{})
-	tunnel.keepaliveChan.Store(&stopChan)
-
-	go tunnel.runKeepaliveLoop(ctx, stopChan, onFailure)
-}
-
-// returns a channel that is closed if the SSH keepalive fails,
-// or nil if no SSH tunnel is configured
-func (tunnel *SSHTunnel) GetKeepaliveChan(ctx context.Context) <-chan struct{} {
-	if tunnel == nil || tunnel.Client == nil || tunnel.badTunnel {
-		// nil channel would be of no consequence in a select
-		// UNLESS it's the only branch in a select, in which case it would block forever
-		return nil
-	}
-	if keepaliveChan := tunnel.keepaliveChan.Load(); keepaliveChan != nil {
-		// Already started
-		return *keepaliveChan
+func (tunnel *SSHTunnel) sendKeepAlive() bool {
+	sent := make(chan error, 1)
+	go func() {
+		_, _, err := tunnel.SendRequest("keepalive@openssh.com", true, nil)
+		sent <- err
+	}()
+	timer := time.NewTimer(SSHRequestTimeout)
+	defer timer.Stop()
+	select {
+	case err := <-sent:
+		if err != nil {
+			tunnel.logger.Error("SSH keepalive failed, tearing down tunnel", slog.Any("error", err))
+			tunnel.shutdown(err)
+			return false
+		}
+		return true
+	case <-timer.C:
+		tunnel.logger.Error("SSH keepalive request timed out, tearing down tunnel")
+		tunnel.shutdown(errors.New("SSH keepalive request timed out"))
+		return false
+	case <-tunnel.watch:
+		return false
 	}
-	keepaliveChan := make(chan struct{})
-	tunnel.keepaliveChan.Store(&keepaliveChan)
-
-	go tunnel.runKeepaliveLoop(ctx, keepaliveChan, nil)
-	return keepaliveChan
 }