diff --git a/flow/connectors/mysql/mysql.go b/flow/connectors/mysql/mysql.go index 3f1dc0db5..951e289b2 100644 --- a/flow/connectors/mysql/mysql.go +++ b/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() diff --git a/flow/connectors/mysql/ssh_keepalive_test.go b/flow/connectors/mysql/ssh_keepalive_test.go index d69a41ea0..41b5b5701 100644 --- a/flow/connectors/mysql/ssh_keepalive_test.go +++ b/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") diff --git a/flow/connectors/postgres/postgres.go b/flow/connectors/postgres/postgres.go index e60c2f94d..a013e5a15 100644 --- a/flow/connectors/postgres/postgres.go +++ b/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 } diff --git a/flow/connectors/postgres/ssh_keepalive_test.go b/flow/connectors/postgres/ssh_keepalive_test.go index 2794618bc..58059179a 100644 --- a/flow/connectors/postgres/ssh_keepalive_test.go +++ b/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, diff --git a/flow/connectors/postgres/ssh_wrapped_conn.go b/flow/connectors/postgres/ssh_wrapped_conn.go index 291c2f4f0..4f087ef7c 100644 --- a/flow/connectors/postgres/ssh_wrapped_conn.go +++ b/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 } diff --git a/flow/connectors/utils/ssh.go b/flow/connectors/utils/ssh.go index 430117de9..c74e75cd0 100644 --- a/flow/connectors/utils/ssh.go +++ b/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 } diff --git a/flow/connectors/utils/ssh_test.go b/flow/connectors/utils/ssh_test.go index 191d29f99..6dd73cfb9 100644 --- a/flow/connectors/utils/ssh_test.go +++ b/flow/connectors/utils/ssh_test.go @@ -1,41 +1,19 @@ package utils import ( - "context" "testing" "github.com/stretchr/testify/require" ) -func TestSSHTunnel_GetKeepaliveChan_NilCases(t *testing.T) { +func TestSSHTunnel_Watch_NilCases(t *testing.T) { // Nil tunnel var tunnel *SSHTunnel = nil - require.Nil(t, tunnel.GetKeepaliveChan(context.Background()), "Nil tunnel should return nil channel") + require.Nil(t, tunnel.Watch(), "Nil tunnel should return nil channel") // Nil client tunnel = &SSHTunnel{Client: nil} - require.Nil(t, tunnel.GetKeepaliveChan(context.Background()), "Tunnel with nil client should return nil channel") - - // Bad tunnel - tunnel = &SSHTunnel{Client: nil, badTunnel: true} - require.Nil(t, tunnel.GetKeepaliveChan(context.Background()), "Bad tunnel should return nil channel") -} - -func TestSSHTunnel_StartKeepalive_NilCases(t *testing.T) { - called := false - onFailure := func() { called = true } - - var tunnel *SSHTunnel = nil - tunnel.StartKeepalive(context.Background(), onFailure) - require.False(t, called, "Nil tunnel should not call onFailure") - - tunnel = &SSHTunnel{Client: nil} - tunnel.StartKeepalive(context.Background(), onFailure) - require.False(t, called, "Tunnel with nil client should not call onFailure") - - tunnel = &SSHTunnel{Client: nil, badTunnel: true} - tunnel.StartKeepalive(context.Background(), onFailure) - require.False(t, called, "Bad tunnel should not call onFailure") + require.Nil(t, tunnel.Watch(), "Tunnel with nil client should return nil channel") } func TestSSHTunnel_Close_NilCases(t *testing.T) {