Propagate cancel watcher deadline errors.

fbconn/cancel_operation_test.go

@@ -0,0 +1,121 @@
+package fbconn
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/firebirdsql/fbx/fbproto"
+	"github.com/stretchr/testify/require"
+)
+
+type cancelDeadlineErrorConn struct {
+	noopConn
+	cancel      context.CancelFunc
+	deadlineSet chan struct{}
+	setErr      error
+	resetErr    error
+	writeErr    error
+}
+
+func (c *cancelDeadlineErrorConn) Write(b []byte) (int, error) {
+	c.cancel()
+	<-c.deadlineSet
+	if c.writeErr != nil {
+		return 0, c.writeErr
+	}
+	return len(b), nil
+}
+
+func (c *cancelDeadlineErrorConn) SetDeadline(deadline time.Time) error {
+	if deadline.IsZero() {
+		return c.resetErr
+	}
+	close(c.deadlineSet)
+	return c.setErr
+}
+
+func TestSendCancelOperationDeadlineErrors(t *testing.T) {
+	setErr := errors.New("set deadline failed")
+	resetErr := errors.New("reset deadline failed")
+	writeErr := errors.New("write failed")
+
+	for _, tt := range []struct {
+		name     string
+		setErr   error
+		resetErr error
+		writeErr error
+	}{
+		{name: "success"},
+		{name: "set deadline", setErr: setErr},
+		{name: "reset deadline", resetErr: resetErr},
+		{name: "set and reset deadline", setErr: setErr, resetErr: resetErr},
+		{name: "write", writeErr: writeErr},
+		{name: "write and reset deadline", writeErr: writeErr, resetErr: resetErr},
+		{name: "write and both deadlines", writeErr: writeErr, setErr: setErr, resetErr: resetErr},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+			conn := &cancelDeadlineErrorConn{
+				cancel:      cancel,
+				deadlineSet: make(chan struct{}),
+				setErr:      tt.setErr,
+				resetErr:    tt.resetErr,
+				writeErr:    tt.writeErr,
+			}
+			fbConn := newCancelOperationTestFbConn(conn)
+			t.Cleanup(func() {
+				fbConn.asyncClose()
+				waitCleanupDone(t, fbConn)
+			})
+
+			err := fbConn.sendCancelOperation(ctx, fbproto.Fb_cancel_raise)
+
+			if tt.setErr == nil && tt.resetErr == nil && tt.writeErr == nil {
+				require.NoError(t, err)
+			}
+			for _, wantErr := range []error{tt.setErr, tt.resetErr, tt.writeErr} {
+				if wantErr != nil {
+					require.ErrorIs(t, err, wantErr)
+				}
+			}
+			wantClosed := tt.setErr != nil || tt.resetErr != nil
+			require.Equal(t, wantClosed, fbConn.IsClosed())
+			require.Equal(t, wantClosed, conn.closed.Load())
+			if wantClosed {
+				require.Error(t, fbConn.lock())
+				waitCleanupDone(t, fbConn)
+			}
+		})
+	}
+}
+
+func BenchmarkSendCancelOperation(b *testing.B) {
+	for _, cancellable := range []bool{false, true} {
+		name := "Background"
+		if cancellable {
+			name = "Cancellable"
+		}
+		b.Run(name, func(b *testing.B) {
+			ctx := context.Background()
+			if cancellable {
+				var cancel context.CancelFunc
+				ctx, cancel = context.WithCancel(ctx)
+				defer cancel()
+			}
+			fbConn := newCancelOperationTestFbConn(&noopConn{})
+			b.Cleanup(func() {
+				fbConn.asyncClose()
+				<-fbConn.CleanupDone()
+			})
+			b.ReportAllocs()
+			for b.Loop() {
+				if err := fbConn.sendCancelOperation(ctx, fbproto.Fb_cancel_raise); err != nil {
+					b.Fatal(err)
+				}
+			}
+		})
+	}
+}

fbconn/fbconn.go

@@ -1843,26 +1843,30 @@ func (fbConn *FbConn) internalPrepareTx(
 }
 
 // CancelRequest sends a cancel request to the server. It returns an error if unable to deliver the cancel
-// request, but lack of an error does not ensure that the query was canceled. Delivery errors do not close the
-// connection; callers that consider a failed cancel request fatal must close the connection themselves.
+// request, but lack of an error does not ensure that the query was canceled. Errors setting or clearing the
+// context's network deadline close the connection. Other delivery errors leave it open; callers that consider
+// a failed cancel request fatal must close the connection themselves.
 func (fbConn *FbConn) CancelRequest(ctx context.Context) error {
 	return fbConn.sendCancelOperation(ctx, fbproto.Fb_cancel_raise)
 }
 
-func (fbConn *FbConn) sendCancelOperation(ctx context.Context, kind int) error {
+func (fbConn *FbConn) sendCancelOperation(ctx context.Context, kind int) (err error) {
 	if ctx != context.Background() {
-		contextWatcher := ctxwatch.NewContextWatcher(&DeadlineContextWatcherHandler{Conn: fbConn.conn})
+		handler := &DeadlineContextWatcherHandler{Conn: fbConn.conn}
+		contextWatcher := ctxwatch.NewContextWatcher(handler)
 		contextWatcher.Watch(ctx)
-		defer contextWatcher.Unwatch()
+		defer func() {
+			contextWatcher.Unwatch()
+			if handler.err != nil {
+				fbConn.asyncClose()
+				err = operationerror.Join(err, handler.err)
+			}
+		}()
 	}
 
-	if err := fbConn.withWriteLockAndFlush(func() error {
+	return fbConn.withWriteLockAndFlush(func() error {
 		return fbConn.client.WriteMessage(&fbproto.OpCancel{Kind: kind})
-	}); err != nil {
-		return err
-	}
-
-	return nil
+	})
 }
 
 // ExecImmediate executes statement immediate inside internal transaction
@@ -3268,16 +3272,21 @@ type DeadlineContextWatcherHandler struct {
 
 	// DeadlineDelay is the delay to set on the deadline set on net.Conn when the context is canceled.
 	DeadlineDelay time.Duration
+
+	// err is safe to read after ContextWatcher.Unwatch returns.
+	err error
 }
 
 func (h *DeadlineContextWatcherHandler) HandleCancel(ctx context.Context) {
 	if err := h.Conn.SetDeadline(time.Now().Add(h.DeadlineDelay)); err != nil {
+		h.err = err
 		_ = h.Conn.Close()
 	}
 }
 
 func (h *DeadlineContextWatcherHandler) HandleUnwatchAfterCancel() {
 	if err := h.Conn.SetDeadline(time.Time{}); err != nil {
+		h.err = errors.Join(h.err, err)
 		_ = h.Conn.Close()
 	}
 }