From 17e3f3dc6a6f2f98780ca7d3c1c2c92ea9abe96d Mon Sep 17 00:00:00 2001
From: U Cirello <ulderico.cirello@gmail.com>
Date: Thu, 7 Mar 2024 10:38:16 +0100
Subject: [PATCH] pglock: not log context cancelation on heartbeats (#94)

---
 client.go               | 10 ++++--
 client_internal_test.go | 68 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 75 insertions(+), 3 deletions(-)

diff --git a/client.go b/client.go
index bdb7cd8..3de1d33 100644
--- a/client.go
+++ b/client.go
@@ -359,11 +359,11 @@ func (c *Client) heartbeat(ctx context.Context, l *Lock) {
 	c.log.Debug("heartbeat started: %v", l.name)
 	defer c.log.Debug("heartbeat stopped: %v", l.name)
 	for {
-		if err := ctx.Err(); err != nil {
-			return
-		} else if err := c.SendHeartbeat(ctx, l); err != nil {
+		if err := c.SendHeartbeat(ctx, l); err != nil && !isContextError(err) {
 			defer c.log.Error("heartbeat missed: %v", err)
 			return
+		} else if err := ctx.Err(); err != nil {
+			return
 		}
 		waitFor(ctx, c.heartbeatFrequency)
 	}
@@ -618,3 +618,7 @@ func waitFor(ctx context.Context, d time.Duration) {
 	case <-ctx.Done():
 	}
 }
+
+func isContextError(err error) bool {
+	return errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded)
+}
diff --git a/client_internal_test.go b/client_internal_test.go
index e263ed1..5a4996d 100644
--- a/client_internal_test.go
+++ b/client_internal_test.go
@@ -17,6 +17,7 @@ limitations under the License.
 package pglock
 
 import (
+	"bytes"
 	"context"
 	"database/sql"
 	"errors"
@@ -24,6 +25,7 @@ import (
 	"io"
 	"log"
 	"net"
+	"strings"
 	"testing"
 	"time"
 
@@ -331,3 +333,69 @@ func Test_waitFor(t *testing.T) {
 		}
 	})
 }
+
+func Test_heartbeatLogging(t *testing.T) {
+	t.Run("cancelled", func(t *testing.T) {
+		db, err := sql.Open("postgres", "")
+		if err != nil {
+			t.Fatal("cannot connect to test database server:", err)
+		}
+		logger := &testBufferLogger{}
+		client, _ := New(db, WithHeartbeatFrequency(0), WithLogger(logger))
+		db, mock, err := sqlmock.New()
+		if err != nil {
+			t.Fatal("cannot create mock:", err)
+		}
+		client.db = db
+		ctx, cancel := context.WithCancel(context.Background())
+		fakeLock := &Lock{
+			heartbeatContext: ctx,
+			heartbeatCancel:  cancel,
+			leaseDuration:    time.Minute,
+		}
+		hbCtx, hbCancel := context.WithCancel(context.Background())
+		hbCancel()
+		mock.ExpectQuery(`SELECT nextval\('locks_rvn'\)`).WillReturnError(hbCtx.Err())
+		fakeLock.heartbeatWG.Add(1)
+		client.heartbeat(hbCtx, fakeLock)
+		t.Log(logger.buf.String())
+		if strings.Contains(logger.buf.String(), context.Canceled.Error()) {
+			t.Fatal("must not log context cancellations")
+		}
+	})
+	t.Run("errored", func(t *testing.T) {
+		db, err := sql.Open("postgres", "")
+		if err != nil {
+			t.Fatal("cannot connect to test database server:", err)
+		}
+		logger := &testBufferLogger{}
+		client, _ := New(db, WithHeartbeatFrequency(0), WithLogger(logger))
+		db, mock, err := sqlmock.New()
+		if err != nil {
+			t.Fatal("cannot create mock:", err)
+		}
+		client.db = db
+		ctx, cancel := context.WithCancel(context.Background())
+		fakeLock := &Lock{
+			heartbeatContext: ctx,
+			heartbeatCancel:  cancel,
+			leaseDuration:    time.Minute,
+		}
+		errExpected := errors.New("expected error")
+		mock.ExpectQuery(`SELECT nextval\('locks_rvn'\)`).WillReturnError(errExpected)
+		fakeLock.heartbeatWG.Add(1)
+		client.heartbeat(context.Background(), fakeLock)
+		t.Log(logger.buf.String())
+		if !strings.Contains(logger.buf.String(), errExpected.Error()) {
+			t.Fatal("expected error missing")
+		}
+	})
+}
+
+type testBufferLogger struct {
+	buf bytes.Buffer
+}
+
+func (t *testBufferLogger) Println(v ...interface{}) {
+	fmt.Fprintln(&t.buf, v...)
+}