From 659525c96185d0dbe60b39e89c86100328bc9a5d Mon Sep 17 00:00:00 2001
From: ngavinsir <ngavinsir@gmail.com>
Date: Sun, 12 May 2024 21:59:19 +0700
Subject: [PATCH] trace release

---
 pgxpool/conn.go        |  5 +----
 pgxpool/tracer.go      | 12 ++++--------
 pgxpool/tracer_test.go | 35 ++++++++++-------------------------
 3 files changed, 15 insertions(+), 37 deletions(-)

diff --git a/pgxpool/conn.go b/pgxpool/conn.go
index 30777be4..e46dc3f8 100644
--- a/pgxpool/conn.go
+++ b/pgxpool/conn.go
@@ -28,10 +28,7 @@ func (c *Conn) Release() {
 	c.res = nil
 
 	if c.p.releaseTracer != nil {
-		ctx := c.p.releaseTracer.TraceReleaseStart(c.ctx, c.p, TraceReleaseStartData{Conn: conn})
-		defer func() {
-			c.p.releaseTracer.TraceReleaseEnd(ctx, c.p, TraceReleaseEndData{})
-		}()
+		c.p.releaseTracer.TraceRelease(c.ctx, c.p, TraceReleaseData{Conn: conn})
 	}
 
 	if conn.IsClosed() || conn.PgConn().IsBusy() || conn.PgConn().TxStatus() != 'I' {
diff --git a/pgxpool/tracer.go b/pgxpool/tracer.go
index b8bf2362..cca0fa5b 100644
--- a/pgxpool/tracer.go
+++ b/pgxpool/tracer.go
@@ -24,15 +24,11 @@ type TraceAcquireEndData struct {
 
 // ReleaseTracer traces Release.
 type ReleaseTracer interface {
-	// TraceReleaseStart is called at the beginning of Release.
-	// The returned context will be passed to the TraceReleaseEnd.
-	TraceReleaseStart(ctx context.Context, pool *Pool, data TraceReleaseStartData) context.Context
-	// TraceReleaseEnd is called when a connection has been released.
-	TraceReleaseEnd(ctx context.Context, pool *Pool, data TraceReleaseEndData)
+	// TraceRelease is called at the beginning of Release.
+	// The passed context is coming from Acquire.
+	TraceRelease(ctx context.Context, pool *Pool, data TraceReleaseData)
 }
 
-type TraceReleaseStartData struct {
+type TraceReleaseData struct {
 	Conn *pgx.Conn
 }
-
-type TraceReleaseEndData struct{}
diff --git a/pgxpool/tracer_test.go b/pgxpool/tracer_test.go
index 72b08b29..4cc77e7c 100644
--- a/pgxpool/tracer_test.go
+++ b/pgxpool/tracer_test.go
@@ -14,8 +14,7 @@ import (
 type testTracer struct {
 	traceAcquireStart func(ctx context.Context, pool *pgxpool.Pool, data pgxpool.TraceAcquireStartData) context.Context
 	traceAcquireEnd   func(ctx context.Context, pool *pgxpool.Pool, data pgxpool.TraceAcquireEndData)
-	traceReleaseStart func(ctx context.Context, pool *pgxpool.Pool, data pgxpool.TraceReleaseStartData) context.Context
-	traceReleaseEnd   func(ctx context.Context, pool *pgxpool.Pool, data pgxpool.TraceReleaseEndData)
+	traceRelease      func(ctx context.Context, pool *pgxpool.Pool, data pgxpool.TraceReleaseData)
 }
 
 type ctxKey string
@@ -33,16 +32,9 @@ func (tt *testTracer) TraceAcquireEnd(ctx context.Context, pool *pgxpool.Pool, d
 	}
 }
 
-func (tt *testTracer) TraceReleaseStart(ctx context.Context, pool *pgxpool.Pool, data pgxpool.TraceReleaseStartData) context.Context {
-	if tt.traceReleaseStart != nil {
-		return tt.traceReleaseStart(ctx, pool, data)
-	}
-	return ctx
-}
-
-func (tt *testTracer) TraceReleaseEnd(ctx context.Context, pool *pgxpool.Pool, data pgxpool.TraceReleaseEndData) {
-	if tt.traceReleaseEnd != nil {
-		tt.traceReleaseEnd(ctx, pool, data)
+func (tt *testTracer) TraceRelease(ctx context.Context, pool *pgxpool.Pool, data pgxpool.TraceReleaseData) {
+	if tt.traceRelease != nil {
+		tt.traceRelease(ctx, pool, data)
 	}
 }
 
@@ -124,24 +116,17 @@ func TestTraceRelease(t *testing.T) {
 	require.NoError(t, err)
 	defer pool.Close()
 
-	traceReleaseStartCalled := false
-	tracer.traceReleaseStart = func(ctx context.Context, pool *pgxpool.Pool, data pgxpool.TraceReleaseStartData) context.Context {
-		traceReleaseStartCalled = true
+	traceReleaseCalled := false
+	tracer.traceRelease = func(ctx context.Context, pool *pgxpool.Pool, data pgxpool.TraceReleaseData) {
+		traceReleaseCalled = true
 		require.NotNil(t, pool)
 		require.NotNil(t, data.Conn)
-		return context.WithValue(ctx, ctxKey("fromTraceReleaseStart"), "foo")
-	}
-
-	traceReleaseEndCalled := false
-	tracer.traceReleaseEnd = func(ctx context.Context, pool *pgxpool.Pool, data pgxpool.TraceReleaseEndData) {
-		traceReleaseEndCalled = true
-		require.Equal(t, "foo", ctx.Value(ctxKey("fromTraceReleaseStart")))
-		require.NotNil(t, pool)
+		require.Equal(t, "foo", ctx.Value(ctxKey("fromAcquire")))
 	}
 
+	ctx = context.WithValue(ctx, ctxKey("fromAcquire"), "foo")
 	c, err := pool.Acquire(ctx)
 	require.NoError(t, err)
 	c.Release()
-	require.True(t, traceReleaseStartCalled)
-	require.True(t, traceReleaseEndCalled)
+	require.True(t, traceReleaseCalled)
 }