Fix MySQL CDC hang when SSH tunnel dies

.github/workflows/flow.yml

@@ -80,6 +80,8 @@ jobs:
           - 42001:42001
           - 42002:42002
           - 42003:42003
+          - 42004:42004
+          - 42005:42005
           - 49001:49001
           - 49002:49002
           - 49003:49003

ancillary-docker-compose.yml

@@ -142,6 +142,8 @@ services:
       - 42001:42001
       - 42002:42002
       - 42003:42003
+      - 42004:42004
+      - 42005:42005
       - 49001:49001
       - 49002:49002
       - 49003:49003

flow/connectors/mysql/cdc.go

@@ -247,17 +247,18 @@ func (c *MySqlConnector) startSyncer(ctx context.Context) (*replication.BinlogSy
 
 	//nolint:gosec
 	return replication.NewBinlogSyncer(replication.BinlogSyncerConfig{
-		ServerID:   rand.Uint32(),
-		Flavor:     c.Flavor(),
-		Host:       config.Host,
-		Port:       uint16(config.Port),
-		User:       config.User,
-		Password:   config.Password,
-		Logger:     internal.SlogLoggerFromCtx(ctx),
-		Dialer:     c.Dialer(),
-		UseDecimal: true,
-		ParseTime:  true,
-		TLSConfig:  tlsConfig,
+		ServerID:         rand.Uint32(),
+		Flavor:           c.Flavor(),
+		Host:             config.Host,
+		Port:             uint16(config.Port),
+		User:             config.User,
+		Password:         config.Password,
+		Logger:           internal.SlogLoggerFromCtx(ctx),
+		Dialer:           c.Dialer(),
+		DisableRetrySync: true,
+		UseDecimal:       true,
+		ParseTime:        true,
+		TLSConfig:        tlsConfig,
 	}), nil
 }
 

flow/connectors/mysql/mysql.go

@@ -76,6 +76,10 @@ func NewMySqlConnector(ctx context.Context, config *protos.MySqlConfig) (*MySqlC
 			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 {

flow/connectors/mysql/ssh_keepalive_test.go

@@ -5,22 +5,31 @@ import (
 	"os"
 	"strconv"
 	"strings"
+	"sync/atomic"
 	"testing"
+	"time"
 
+	toxiproxy "github.com/Shopify/toxiproxy/v2/client"
 	"github.com/stretchr/testify/require"
 
 	"github.com/PeerDB-io/peerdb/flow/connectors/utils"
 	"github.com/PeerDB-io/peerdb/flow/generated/protos"
+	"github.com/PeerDB-io/peerdb/flow/model"
+	"github.com/PeerDB-io/peerdb/flow/otel_metrics"
+	"github.com/PeerDB-io/peerdb/flow/shared"
+	"github.com/PeerDB-io/peerdb/flow/shared/concurrency"
 )
 
 const (
-	toxiproxyDownProxyPort    = 42001
-	toxiproxyLatencyProxyPort = 42002
-	toxiproxyResetProxyPort   = 42003
+	toxiproxyDownProxyPort         = 42001
+	toxiproxyLatencyProxyPort      = 42002
+	toxiproxyResetProxyPort        = 42003
+	toxiproxyCDCHangProxyPort      = 42004
+	toxiproxyCDCCloseHangProxyPort = 42005
 )
 
-func setupMySQLConnectorWithSSH(ctx context.Context, t *testing.T, proxyName string, proxyPort int,
-) (*MySqlConnector, utils.SSHKeepaliveTestConfig) {
+func setupMySQLConnectorWithProxy(ctx context.Context, t *testing.T, proxyName string, proxyPort int,
+) (*MySqlConnector, *toxiproxy.Proxy) {
 	t.Helper()
 
 	toxiproxyClient := utils.NewToxiproxyClient(t)
@@ -63,6 +72,14 @@ func setupMySQLConnectorWithSSH(ctx context.Context, t *testing.T, proxyName str
 	err = connector.ConnectionActive(ctx)
 	require.NoError(t, err, "Initial connection should work")
 
+	return connector, sshProxy
+}
+
+func setupMySQLConnectorWithSSH(ctx context.Context, t *testing.T, proxyName string, proxyPort int,
+) (*MySqlConnector, utils.SSHKeepaliveTestConfig) {
+	t.Helper()
+
+	connector, sshProxy := setupMySQLConnectorWithProxy(ctx, t, proxyName, proxyPort)
 	keepaliveChan := connector.ssh.GetKeepaliveChan(ctx)
 
 	return connector, utils.SSHKeepaliveTestConfig{
@@ -76,6 +93,7 @@ func setupMySQLConnectorWithSSH(ctx context.Context, t *testing.T, proxyName str
 }
 
 func TestMySQLSSHKeepaliveWithToxiproxy(t *testing.T) {
+	t.Parallel()
 	if os.Getenv("CI_MYSQL_VERSION") == "maria" {
 		t.Skip("Skipping SSH keepalive test for MariaDB")
 	}
@@ -85,6 +103,7 @@ func TestMySQLSSHKeepaliveWithToxiproxy(t *testing.T) {
 }
 
 func TestMySQLSSHKeepaliveLatency(t *testing.T) {
+	t.Parallel()
 	if os.Getenv("CI_MYSQL_VERSION") == "maria" {
 		t.Skip("Skipping SSH keepalive test for MariaDB")
 	}
@@ -94,10 +113,160 @@ func TestMySQLSSHKeepaliveLatency(t *testing.T) {
 }
 
 func TestMySQLSSHResetPeer(t *testing.T) {
+	t.Parallel()
 	if os.Getenv("CI_MYSQL_VERSION") == "maria" {
 		t.Skip("Skipping SSH keepalive test for MariaDB")
 	}
 	connector, cfg := setupMySQLConnectorWithSSH(t.Context(), t, "my-ssh-reset-peer-test", toxiproxyResetProxyPort)
 	defer connector.Close()
 	utils.RunSSHResetPeerTest(t, cfg)
 }
+
+func setupCDCPullRecords(
+	ctx context.Context, t *testing.T, connector *MySqlConnector, flowJobName string,
+) (*model.PullRecordsRequest[model.RecordItems], *otel_metrics.OtelManager) {
+	t.Helper()
+
+	gtidOn, err := connector.GetGtidModeOn(ctx)
+	require.NoError(t, err)
+
+	var offsetText string
+	if gtidOn {
+		gset, err := connector.GetMasterGTIDSet(ctx)
+		require.NoError(t, err)
+		offsetText = gset.String()
+	} else {
+		pos, err := connector.GetMasterPos(ctx)
+		require.NoError(t, err)
+		offsetText = posToOffsetText(pos)
+	}
+
+	otelManager, err := otel_metrics.NewOtelManager(ctx, "test", false)
+	require.NoError(t, err)
+
+	req := &model.PullRecordsRequest[model.RecordItems]{
+		FlowJobName:            flowJobName,
+		RecordStream:           model.NewCDCStream[model.RecordItems](100),
+		TableNameMapping:       map[string]model.NameAndExclude{},
+		TableNameSchemaMapping: map[string]*protos.TableSchema{},
+		LastOffset:             model.CdcCheckpoint{Text: offsetText},
+		MaxBatchSize:           10000,
+		IdleTimeout:            time.Minute,
+		ConsumedOffset:         &atomic.Int64{},
+	}
+
+	return req, otelManager
+}
+
+func TestMySQLSSHKeepaliveCDCHang(t *testing.T) {
+	t.Parallel()
+	if os.Getenv("CI_MYSQL_VERSION") == "maria" {
+		t.Skip("Skipping SSH keepalive test for MariaDB")
+	}
+	ctx := t.Context()
+
+	connector, sshProxy := setupMySQLConnectorWithProxy(ctx, t, "my-ssh-cdc-down-test", toxiproxyCDCHangProxyPort)
+	defer connector.Close()
+
+	keepaliveChan := connector.ssh.GetKeepaliveChan(ctx)
+	require.NotNil(t, keepaliveChan, "SSH keepalive channel should exist")
+
+	req, otelManager := setupCDCPullRecords(ctx, t, connector, "test_ssh_cdc_hang")
+
+	pullDone := concurrency.NewLatch[error]()
+	go func() {
+		pullDone.Set(connector.PullRecords(ctx, shared.CatalogPool{}, otelManager, req))
+	}()
+	go func() {
+		for range req.RecordStream.GetRecords() {
+		}
+	}()
+
+	time.Sleep(2 * time.Second)
+
+	t.Log("Adding latency toxic to simulate network black hole during CDC streaming")
+	_, err := sshProxy.AddToxic("latency", "latency", "", 1.0, toxiproxy.Attributes{
+		"latency": 120000,
+	})
+	require.NoError(t, err)
+
+	// Wait for keepalive to detect the failure first — if PullRecords exits before
+	// the keepalive fires, something else closed the connection (driver timeout, etc.)
+	select {
+	case <-keepaliveChan:
+		t.Log("SSH keepalive detected failure")
+	case <-pullDone.Chan():
+		t.Fatal("PullRecords exited before SSH keepalive fired — connection closed by something other than keepalive")
+	case <-time.After(3 * utils.SSHKeepaliveInterval):
+		t.Fatal("SSH keepalive did not fire in time")
+	}
+
+	select {
+	case <-pullDone.Chan():
+		pullErr := pullDone.Wait()
+		require.Error(t, pullErr, "PullRecords should fail after SSH keepalive detects black hole")
+		t.Logf("PullRecords returned: %v", pullErr)
+	case <-time.After(10 * time.Second):
+		t.Fatal("PullRecords did not return after SSH keepalive closed the connection")
+	}
+}
+
+func TestMySQLSSHKeepaliveCDCCloseHang(t *testing.T) {
+	t.Parallel()
+	if os.Getenv("CI_MYSQL_VERSION") == "maria" {
+		t.Skip("Skipping SSH keepalive test for MariaDB")
+	}
+
+	ctx := t.Context()
+
+	connector, sshProxy := setupMySQLConnectorWithProxy(ctx, t, "my-ssh-cdc-latency-test", toxiproxyCDCCloseHangProxyPort)
+	defer connector.Close()
+
+	keepaliveChan := connector.ssh.GetKeepaliveChan(ctx)
+	require.NotNil(t, keepaliveChan, "SSH keepalive channel should exist")
+
+	req, otelManager := setupCDCPullRecords(ctx, t, connector, "test_ssh_cdc_close_hang")
+
+	// Use a short-lived context so PullRecords starts shutting down while
+	// the connection is blocked by latency
+	cancelCtx, cancel := context.WithTimeout(ctx, 7*time.Second)
+	defer cancel()
+
+	pullDone := concurrency.NewLatch[error]()
+	go func() {
+		pullDone.Set(connector.PullRecords(cancelCtx, shared.CatalogPool{}, otelManager, req))
+	}()
+	go func() {
+		for range req.RecordStream.GetRecords() {
+		}
+	}()
+
+	// Wait for CDC streaming to establish before adding latency
+	time.Sleep(2 * time.Second)
+
+	// Add high latency so the connection becomes unresponsive without erroring
+	t.Log("Adding latency toxic to block the connection during CDC streaming and cleanup")
+	_, err := sshProxy.AddToxic("latency", "latency", "", 1.0, toxiproxy.Attributes{
+		"latency": 120000,
+	})
+	require.NoError(t, err)
+
+	t.Logf("Waiting for context cancel (~5s), then PullRecords should hang until keepalive fires (within %s)", 3*utils.SSHKeepaliveInterval)
+
+	// Wait for keepalive to detect the failure — PullRecords should be hanging
+	// until keepalive closes the underlying connection
+	select {
+	case <-keepaliveChan:
+		t.Log("SSH keepalive detected failure")
+	case <-time.After(3 * utils.SSHKeepaliveInterval):
+		t.Fatal("SSH keepalive did not fire in time")
+	}
+
+	select {
+	case <-pullDone.Chan():
+		pullErr := pullDone.Wait()
+		require.ErrorIs(t, pullErr, context.DeadlineExceeded)
+	case <-time.After(10 * time.Second):
+		t.Fatal("PullRecords did not return after SSH keepalive closed the connection")
+	}
+}

flow/connectors/utils/ssh.go

@@ -20,6 +20,7 @@ const SSHKeepaliveInterval = 15 * time.Second
 
 type SSHTunnel struct {
 	*ssh.Client
+	logger        *slog.Logger
 	keepaliveChan atomic.Pointer[chan struct{}]
 	badTunnel     bool
 }
@@ -90,7 +91,10 @@ func NewSSHTunnel(
 			return nil, exceptions.NewSSHTunnelSetupError(err)
 		}
 
-		return &SSHTunnel{Client: client, badTunnel: false}, nil
+		return &SSHTunnel{
+			Client: client,
+			logger: internal.SlogLoggerFromCtx(ctx),
+		}, nil
 	}
 
 	return nil, nil
@@ -112,7 +116,7 @@ func (tunnel *SSHTunnel) runKeepaliveLoop(
 ) {
 	ticker := time.NewTicker(SSHKeepaliveInterval)
 	defer ticker.Stop()
-	logger := internal.LoggerFromCtx(ctx)
+	logger := tunnel.logger
 	// in case request hangs, we want to detect that and not send another request
 	requestSent := atomic.Bool{}
 	var keepaliveErr error
@@ -124,7 +128,7 @@ func (tunnel *SSHTunnel) runKeepaliveLoop(
 		case <-ticker.C:
 			if requestSent.Load() {
 				// Previous keepalive request didn't return yet, something's wrong
-				logger.Error("Previous keepalive request still pending, marking tunnel as bad")
+				logger.ErrorContext(ctx, "Previous keepalive request still pending, marking tunnel as bad")
 				if keepaliveChan := tunnel.keepaliveChan.Swap(nil); keepaliveChan != nil {
 					close(*keepaliveChan)
 				}
@@ -152,7 +156,7 @@ func (tunnel *SSHTunnel) runKeepaliveLoop(
 			// channel closed from outside
 			return
 		case <-errChan:
-			logger.Error("Keepalive request failed, marking tunnel as bad", slog.Any("error", keepaliveErr))
+			logger.ErrorContext(ctx, "Keepalive request failed, marking tunnel as bad", slog.Any("error", keepaliveErr))
 			if keepaliveChan := tunnel.keepaliveChan.Swap(nil); keepaliveChan != nil {
 				close(*keepaliveChan)
 			}