classify error as retryable when pipe closed in deadlineCappableConn

flow/alerting/classifier.go

@@ -428,6 +428,13 @@ func GetErrorClass(ctx context.Context, err error) (ErrorClass, ErrorInfo) {
 		}
 	}
 
+	if _, ok := errors.AsType[*exceptions.SSHTunnelClosedError](err); ok {
+		return ErrorRetryRecoverable, ErrorInfo{
+			Source: ErrorSourceSSH,
+			Code:   "TUNNEL_CONNECTION_CLOSED",
+		}
+	}
+
 	// An SSH-tunneled connection that goes bad (e.g. keepalive failure) is wrapped explicitly.
 	if _, ok := errors.AsType[*exceptions.SSHTunnelConnectionError](err); ok {
 		return ErrorNotifyConnectivity, ErrorInfo{

flow/alerting/classifier_test.go

@@ -5,6 +5,7 @@ import (
 	"crypto/tls"
 	"errors"
 	"fmt"
+	"io"
 	"net"
 	"strconv"
 	"syscall"
@@ -72,6 +73,20 @@ func TestSSHTunnelConnectionErrorShouldBeConnectivity(t *testing.T) {
 	}, errInfo, "Unexpected error info")
 }
 
+func TestSSHTunnelClosedPipeErrorShouldBeRetryable(t *testing.T) {
+	t.Parallel()
+
+	// Mirrors how deadlineCapableConn wraps net.Pipe teardown errors when an
+	// SSH-tunneled connection is torn down while a query is in flight.
+	err := fmt.Errorf("receive message failed: %w", exceptions.NewSSHTunnelClosedError(io.ErrClosedPipe))
+	errorClass, errInfo := GetErrorClass(t.Context(), err)
+	assert.Equal(t, ErrorRetryRecoverable, errorClass)
+	assert.Equal(t, ErrorInfo{
+		Source: ErrorSourceSSH,
+		Code:   "TUNNEL_CONNECTION_CLOSED",
+	}, errInfo)
+}
+
 func TestNeonConnectivityErrorShouldBeConnectivity(t *testing.T) {
 	t.Skip("Not a good idea to run this test in CI as it goes to Neon, maybe we need a better mock")
 	config, err := pgx.ParseConfig("postgres://random-endpoint-id-here.us-east-2.aws.neon.tech:5432/db?options=endpoint%3Dtest_endpoint")

flow/connectors/utils/deadline_capable_conn.go

@@ -1,9 +1,12 @@
 package utils
 
 import (
+	"errors"
 	"io"
 	"net"
 	"sync"
+
+	"github.com/PeerDB-io/peerdb/flow/shared/exceptions"
 )
 
 type deadlineCapableConn struct {
@@ -49,6 +52,23 @@ func NewDeadlineCapableConn(sshConn net.Conn) net.Conn {
 	return dcConn
 }
 
+func (c *deadlineCapableConn) Read(b []byte) (int, error) {
+	n, err := c.Conn.Read(b)
+	return n, translateErr(err)
+}
+
+func (c *deadlineCapableConn) Write(b []byte) (int, error) {
+	n, err := c.Conn.Write(b)
+	return n, translateErr(err)
+}
+
+func translateErr(err error) error {
+	if errors.Is(err, io.ErrClosedPipe) || errors.Is(err, io.EOF) {
+		return exceptions.NewSSHTunnelClosedError(err)
+	}
+	return err
+}
+
 func (c *deadlineCapableConn) Close() error {
 	err := c.Conn.Close()
 	c.closeAll()

flow/connectors/utils/deadline_capable_conn_test.go

@@ -9,6 +9,8 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/require"
+
+	"github.com/PeerDB-io/peerdb/flow/shared/exceptions"
 )
 
 func TestNewDeadlineCapableConn_ReadDeadlineCanBeCleared(t *testing.T) {
@@ -147,7 +149,8 @@ func TestNewDeadlineCapableConn_SSHCloseClosesLocalAndRemote(t *testing.T) {
 
 	select {
 	case err := <-localDone:
-		require.ErrorIs(t, err, io.EOF)
+		var tunnelErr *exceptions.SSHTunnelClosedError
+		require.ErrorAs(t, err, &tunnelErr)
 	case <-time.After(time.Second):
 		t.Fatal("SSH close did not close local side")
 	}
@@ -176,10 +179,10 @@ func TestNewDeadlineCapableConn_RemoteCloseClosesSSHAndLocalSide(t *testing.T) {
 		sshDone <- err
 	}()
 
-	localDonn := make(chan error, 1)
+	localDone := make(chan error, 1)
 	go func() {
 		_, err := conn.Read(make([]byte, 1))
-		localDonn <- err
+		localDone <- err
 	}()
 
 	select {
@@ -196,8 +199,9 @@ func TestNewDeadlineCapableConn_RemoteCloseClosesSSHAndLocalSide(t *testing.T) {
 	}
 
 	select {
-	case err := <-localDonn:
-		require.ErrorIs(t, err, io.EOF)
+	case err := <-localDone:
+		var tunnelErr *exceptions.SSHTunnelClosedError
+		require.ErrorAs(t, err, &tunnelErr)
 	case <-time.After(time.Second):
 		t.Fatal("remote close did not close server side")
 	}

flow/shared/exceptions/ssh.go

@@ -35,3 +35,19 @@ func (e *SSHTunnelConnectionError) Error() string {
 func (e *SSHTunnelConnectionError) Unwrap() error {
 	return e.error
 }
+
+type SSHTunnelClosedError struct {
+	error
+}
+
+func NewSSHTunnelClosedError(err error) *SSHTunnelClosedError {
+	return &SSHTunnelClosedError{err}
+}
+
+func (e *SSHTunnelClosedError) Error() string {
+	return "SSH Tunnel Closed: " + e.error.Error()
+}
+
+func (e *SSHTunnelClosedError) Unwrap() error {
+	return e.error
+}