trace release

pull/2017/head
ngavinsir 2024-05-12 21:59:19 +07:00 committed by Jack Christensen
parent 4dd1810d8b
commit 659525c961
3 changed files with 15 additions and 37 deletions

View File

@ -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' {

View File

@ -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{}

View File

@ -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)
}