From c57c37f80a260cef0ad766e3be269aba39dcae8a Mon Sep 17 00:00:00 2001
From: Pavel Gabriel <alovak@gmail.com>
Date: Tue, 28 May 2024 17:03:02 +0200
Subject: [PATCH] report err from OnClose and close connection regardless of
 the error

---
 connection.go      |  2 +-
 connection_test.go | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/connection.go b/connection.go
index 29dd4dd..c08c121 100644
--- a/connection.go
+++ b/connection.go
@@ -326,7 +326,7 @@ func (c *Connection) CloseCtx(ctx context.Context) error {
 
 	if onClose != nil {
 		if err := onClose(ctx, c); err != nil {
-			return fmt.Errorf("on close callback: %w", err)
+			c.handleError(fmt.Errorf("on close callback: %w", err))
 		}
 	}
 
diff --git a/connection_test.go b/connection_test.go
index bacfb9c..df78981 100644
--- a/connection_test.go
+++ b/connection_test.go
@@ -220,6 +220,43 @@ func TestClient_Connect(t *testing.T) {
 			return atomic.LoadInt32(&onClosedCalled) == 1
 		}, 100*time.Millisecond, 20*time.Millisecond, "onClose should be called")
 	})
+
+	t.Run("when OnCloseCtx returns error, we still close the connection", func(t *testing.T) {
+		server, err := NewTestServer()
+		require.NoError(t, err)
+		defer server.Close()
+
+		var onClosedCalled int32
+		onCloseCtx := func(ctx context.Context, c *connection.Connection) error {
+			// increase the counter
+			atomic.AddInt32(&onClosedCalled, 1)
+			return errors.New("error from on close handler")
+		}
+
+		var onErrCalled int32
+		errHandler := func(err error) {
+			atomic.AddInt32(&onErrCalled, 1)
+			require.Contains(t, err.Error(), "error from on close handler")
+		}
+
+		c, err := connection.New(
+			server.Addr,
+			testSpec,
+			readMessageLength,
+			writeMessageLength,
+			connection.ErrorHandler(errHandler),
+			connection.OnCloseCtx(onCloseCtx),
+		)
+		require.NoError(t, err)
+
+		err = c.CloseCtx(context.Background())
+		require.NoError(t, err)
+
+		// eventually the onClosedCalled should be 1
+		require.Eventually(t, func() bool {
+			return atomic.LoadInt32(&onClosedCalled) == 1
+		}, 100*time.Millisecond, 20*time.Millisecond, "onClose should be called")
+	})
 }
 
 func TestClient_Write(t *testing.T) {