review

Author: jgao54

flow/connectors/mysql/cdc.go

@@ -29,6 +29,7 @@ import (
 	"github.com/PeerDB-io/peerdb/flow/shared"
 	"github.com/PeerDB-io/peerdb/flow/shared/datatypes"
 	"github.com/PeerDB-io/peerdb/flow/shared/exceptions"
+	"github.com/PeerDB-io/peerdb/flow/shared/timeout"
 	"github.com/PeerDB-io/peerdb/flow/shared/types"
 )
 
@@ -324,18 +325,11 @@ func (c *MySqlConnector) startCdcStreamingGtid(
 // for ssh tunnel, and killing a connection that has already been reaped by the server
 // but is not propagated to the client). This can lead to syncer.Close() stuck indefinitely.
 // TODO: better to fix properly upstream
-func (c *MySqlConnector) closeSyncerWithTimeout(syncer *replication.BinlogSyncer, timeout time.Duration) {
-	done := make(chan struct{})
-	go func() {
-		syncer.Close()
-		close(done)
-	}()
-	select {
-	case <-done:
-	case <-time.After(timeout):
-		c.logger.Error("[mysql] syncer.Close hung, force-closing SSH tunnel to unblock")
+func (c *MySqlConnector) closeSyncer(syncer *replication.BinlogSyncer, deadline time.Duration) {
+	timeout.CloseWithTimeout(syncer.Close, deadline, func() {
+		c.logger.Error("[mysql] syncer.Close hung, force-closing SSH tunnel")
 		_ = c.ssh.Close()
-	}
+	})
 }
 
 func (c *MySqlConnector) ReplPing(context.Context) error {
@@ -373,7 +367,7 @@ func (c *MySqlConnector) PullRecords(
 	if err != nil {
 		return err
 	}
-	defer c.closeSyncerWithTimeout(syncer, 10*time.Second)
+	defer c.closeSyncer(syncer, 10*time.Second)
 
 	c.logger.Info("[mysql] PullRecords started streaming")
 

flow/connectors/mysql/ssh_keepalive_test.go

@@ -21,12 +21,11 @@ import (
 )
 
 const (
-	toxiproxyDownProxyPort              = 42001
-	toxiproxyLatencyProxyPort           = 42002
-	toxiproxyResetProxyPort             = 42003
-	toxiproxyCDCHangProxyPort           = 42004
-	toxiproxyCDCCloseHangProxyPort      = 42005
-	toxiproxyCloseSyncerWithTimeoutPort = 42006
+	toxiproxyDownProxyPort         = 42001
+	toxiproxyLatencyProxyPort      = 42002
+	toxiproxyResetProxyPort        = 42003
+	toxiproxyCDCHangProxyPort      = 42004
+	toxiproxyCDCCloseHangProxyPort = 42005
 )
 
 func setupMySQLConnectorWithProxy(ctx context.Context, t *testing.T, proxyName string, proxyPort int,
@@ -271,49 +270,3 @@ func TestMySQLSSHKeepaliveCDCCloseHang(t *testing.T) {
 		t.Fatal("PullRecords did not return after SSH keepalive closed the connection")
 	}
 }
-
-func TestMySQLCloseSyncerWithTimeout(t *testing.T) {
-	t.Parallel()
-	if os.Getenv("CI_MYSQL_VERSION") == "maria" {
-		t.Skip("Skipping for MariaDB")
-	}
-
-	ctx := t.Context()
-	connector, sshProxy := setupMySQLConnectorWithProxy(
-		ctx, t, "my-close-syncer-timeout-test", toxiproxyCloseSyncerWithTimeoutPort)
-	defer connector.Close()
-
-	pos, err := connector.GetMasterPos(ctx)
-	require.NoError(t, err)
-	syncer, _, _, _, err := connector.startCdcStreamingFilePos(ctx, pos)
-	require.NoError(t, err)
-
-	// Let CDC streaming establish before blocking the tunnel.
-	time.Sleep(2 * time.Second)
-
-	// Black-hole the SSH transport so syncer.Close() hangs
-	_, err = sshProxy.AddToxic("latency", "latency", "", 1.0, toxiproxy.Attributes{
-		"latency": 120000,
-	})
-	require.NoError(t, err)
-
-	// Must be < SSHKeepaliveInterval so the timing bounds below prove the syncCloseTimeout
-	// (not SSH keepalive's own force-close) is what unblocked syncer.Close.
-	syncerCloseTimeout := 2 * time.Second
-	require.Less(t, syncerCloseTimeout, utils.SSHKeepaliveInterval)
-	done := make(chan struct{})
-	start := time.Now()
-	go func() {
-		connector.closeSyncerWithTimeout(syncer, syncerCloseTimeout)
-		close(done)
-	}()
-
-	select {
-	case <-done:
-		elapsed := time.Since(start)
-		require.GreaterOrEqual(t, elapsed, syncerCloseTimeout)
-		require.Less(t, elapsed, syncerCloseTimeout+time.Second)
-	case <-time.After(10 * time.Second):
-		t.Fatal("closeSyncerWithTimeout did not return on timeout")
-	}
-}

flow/shared/timeout/close.go

@@ -0,0 +1,19 @@
+package timeout
+
+import "time"
+
+// CloseWithTimeout runs closer in a goroutine and waits up to timeout for it
+// to return. If it doesn't, onTimeout is invoked and CloseWithTimeout returns
+// without waiting further.
+func CloseWithTimeout(closer func(), timeout time.Duration, onTimeout func()) {
+	done := make(chan struct{})
+	go func() {
+		closer()
+		close(done)
+	}()
+	select {
+	case <-done:
+	case <-time.After(timeout):
+		onTimeout()
+	}
+}

flow/shared/timeout/close_test.go

@@ -0,0 +1,38 @@
+package timeout
+
+import (
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestCloseWithTimeout_FastClose(t *testing.T) {
+	var onTimeoutCalls atomic.Int32
+	CloseWithTimeout(
+		func() {}, // returns immediately
+		time.Second,
+		func() { onTimeoutCalls.Add(1) },
+	)
+	require.Zero(t, onTimeoutCalls.Load(), "onTimeout should not fire when closer returns quickly")
+}
+
+func TestCloseWithTimeout_ClosureHangs(t *testing.T) {
+	var onTimeoutCalls atomic.Int32
+	release := make(chan struct{})
+	defer close(release)
+
+	timeout := 50 * time.Millisecond
+	start := time.Now()
+	CloseWithTimeout(
+		func() { <-release }, // blocks
+		timeout,
+		func() { onTimeoutCalls.Add(1) },
+	)
+	elapsed := time.Since(start)
+
+	require.Equal(t, int32(1), onTimeoutCalls.Load())
+	require.GreaterOrEqual(t, elapsed, timeout)
+	require.Less(t, elapsed, timeout+time.Second)
+}