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 c.res = nil
if c.p.releaseTracer != nil { if c.p.releaseTracer != nil {
ctx := c.p.releaseTracer.TraceReleaseStart(c.ctx, c.p, TraceReleaseStartData{Conn: conn}) c.p.releaseTracer.TraceRelease(c.ctx, c.p, TraceReleaseData{Conn: conn})
defer func() {
c.p.releaseTracer.TraceReleaseEnd(ctx, c.p, TraceReleaseEndData{})
}()
} }
if conn.IsClosed() || conn.PgConn().IsBusy() || conn.PgConn().TxStatus() != 'I' { if conn.IsClosed() || conn.PgConn().IsBusy() || conn.PgConn().TxStatus() != 'I' {

View File

@ -24,15 +24,11 @@ type TraceAcquireEndData struct {
// ReleaseTracer traces Release. // ReleaseTracer traces Release.
type ReleaseTracer interface { type ReleaseTracer interface {
// TraceReleaseStart is called at the beginning of Release. // TraceRelease is called at the beginning of Release.
// The returned context will be passed to the TraceReleaseEnd. // The passed context is coming from Acquire.
TraceReleaseStart(ctx context.Context, pool *Pool, data TraceReleaseStartData) context.Context TraceRelease(ctx context.Context, pool *Pool, data TraceReleaseData)
// TraceReleaseEnd is called when a connection has been released.
TraceReleaseEnd(ctx context.Context, pool *Pool, data TraceReleaseEndData)
} }
type TraceReleaseStartData struct { type TraceReleaseData struct {
Conn *pgx.Conn Conn *pgx.Conn
} }
type TraceReleaseEndData struct{}

View File

@ -14,8 +14,7 @@ import (
type testTracer struct { type testTracer struct {
traceAcquireStart func(ctx context.Context, pool *pgxpool.Pool, data pgxpool.TraceAcquireStartData) context.Context traceAcquireStart func(ctx context.Context, pool *pgxpool.Pool, data pgxpool.TraceAcquireStartData) context.Context
traceAcquireEnd func(ctx context.Context, pool *pgxpool.Pool, data pgxpool.TraceAcquireEndData) traceAcquireEnd func(ctx context.Context, pool *pgxpool.Pool, data pgxpool.TraceAcquireEndData)
traceReleaseStart func(ctx context.Context, pool *pgxpool.Pool, data pgxpool.TraceReleaseStartData) context.Context traceRelease func(ctx context.Context, pool *pgxpool.Pool, data pgxpool.TraceReleaseData)
traceReleaseEnd func(ctx context.Context, pool *pgxpool.Pool, data pgxpool.TraceReleaseEndData)
} }
type ctxKey string 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 { func (tt *testTracer) TraceRelease(ctx context.Context, pool *pgxpool.Pool, data pgxpool.TraceReleaseData) {
if tt.traceReleaseStart != nil { if tt.traceRelease != nil {
return tt.traceReleaseStart(ctx, pool, data) tt.traceRelease(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)
} }
} }
@ -124,24 +116,17 @@ func TestTraceRelease(t *testing.T) {
require.NoError(t, err) require.NoError(t, err)
defer pool.Close() defer pool.Close()
traceReleaseStartCalled := false traceReleaseCalled := false
tracer.traceReleaseStart = func(ctx context.Context, pool *pgxpool.Pool, data pgxpool.TraceReleaseStartData) context.Context { tracer.traceRelease = func(ctx context.Context, pool *pgxpool.Pool, data pgxpool.TraceReleaseData) {
traceReleaseStartCalled = true traceReleaseCalled = true
require.NotNil(t, pool) require.NotNil(t, pool)
require.NotNil(t, data.Conn) require.NotNil(t, data.Conn)
return context.WithValue(ctx, ctxKey("fromTraceReleaseStart"), "foo") require.Equal(t, "foo", ctx.Value(ctxKey("fromAcquire")))
}
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)
} }
ctx = context.WithValue(ctx, ctxKey("fromAcquire"), "foo")
c, err := pool.Acquire(ctx) c, err := pool.Acquire(ctx)
require.NoError(t, err) require.NoError(t, err)
c.Release() c.Release()
require.True(t, traceReleaseStartCalled) require.True(t, traceReleaseCalled)
require.True(t, traceReleaseEndCalled)
} }