persist ctx in pgxpool conn

This commit is contained in:
ngavinsir 2024-05-12 11:44:23 +07:00 committed by Jack Christensen
parent 25914e21f3
commit 4dd1810d8b
2 changed files with 6 additions and 4 deletions

View File

@ -13,6 +13,7 @@ import (
type Conn struct { type Conn struct {
res *puddle.Resource[*connResource] res *puddle.Resource[*connResource]
p *Pool p *Pool
ctx context.Context
} }
// Release returns c to the pool it was acquired from. Once Release has been called, other methods must not be called. // Release returns c to the pool it was acquired from. Once Release has been called, other methods must not be called.
@ -27,7 +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(context.Background(), c.p, TraceReleaseStartData{Conn: conn}) ctx := c.p.releaseTracer.TraceReleaseStart(c.ctx, c.p, TraceReleaseStartData{Conn: conn})
defer func() { defer func() {
c.p.releaseTracer.TraceReleaseEnd(ctx, c.p, TraceReleaseEndData{}) c.p.releaseTracer.TraceReleaseEnd(ctx, c.p, TraceReleaseEndData{})
}() }()

View File

@ -29,7 +29,7 @@ type connResource struct {
maxAgeTime time.Time maxAgeTime time.Time
} }
func (cr *connResource) getConn(p *Pool, res *puddle.Resource[*connResource]) *Conn { func (cr *connResource) getConn(ctx context.Context, p *Pool, res *puddle.Resource[*connResource]) *Conn {
if len(cr.conns) == 0 { if len(cr.conns) == 0 {
cr.conns = make([]Conn, 128) cr.conns = make([]Conn, 128)
} }
@ -39,6 +39,7 @@ func (cr *connResource) getConn(p *Pool, res *puddle.Resource[*connResource]) *C
c.res = res c.res = res
c.p = p c.p = p
c.ctx = ctx
return c return c
} }
@ -538,7 +539,7 @@ func (p *Pool) Acquire(ctx context.Context) (c *Conn, err error) {
} }
if p.beforeAcquire == nil || p.beforeAcquire(ctx, cr.conn) { if p.beforeAcquire == nil || p.beforeAcquire(ctx, cr.conn) {
return cr.getConn(p, res), nil return cr.getConn(ctx, p, res), nil
} }
res.Destroy() res.Destroy()
@ -566,7 +567,7 @@ func (p *Pool) AcquireAllIdle(ctx context.Context) []*Conn {
for _, res := range resources { for _, res := range resources {
cr := res.Value() cr := res.Value()
if p.beforeAcquire == nil || p.beforeAcquire(ctx, cr.conn) { if p.beforeAcquire == nil || p.beforeAcquire(ctx, cr.conn) {
conns = append(conns, cr.getConn(p, res)) conns = append(conns, cr.getConn(ctx, p, res))
} else { } else {
res.Destroy() res.Destroy()
} }