diff --git a/internal/pool/pool_single.go b/internal/pool/pool_single.go index 5a3fde19..6c0cdffa 100644 --- a/internal/pool/pool_single.go +++ b/internal/pool/pool_single.go @@ -37,6 +37,11 @@ func (p *SingleConnPool) Put(ctx context.Context, cn *Conn) {} func (p *SingleConnPool) Remove(ctx context.Context, cn *Conn, reason error) { p.cn = nil p.stickyErr = reason + // If ctx is cancelled without a reason(error) value, + // then the ctx.Error is used as the reason for why the p.cn is assigned nil. + if reason == nil && ctx != nil { + p.stickyErr = ctx.Err() + } } func (p *SingleConnPool) Close() error { diff --git a/internal/pool/pool_single_test.go b/internal/pool/pool_single_test.go new file mode 100644 index 00000000..b52fdf18 --- /dev/null +++ b/internal/pool/pool_single_test.go @@ -0,0 +1,25 @@ +package pool_test + +import ( + "context" + . "github.com/onsi/ginkgo" + . "github.com/onsi/gomega" + + "github.com/go-pg/pg/v10/internal/pool" +) + +var _ = Describe("SingleConnPool", func() { + It("remove a conn due to context is cancelled", func() { + p := pool.NewSingleConnPool(nil, &pool.Conn{}) + ctx, cancel := context.WithCancel(context.TODO()) + cn, err := p.Get(nil) + Expect(err).To(BeNil()) + Expect(cn).ToNot(BeNil()) + + cancel() + p.Remove(ctx, cn, nil) + cn, err = p.Get(nil) + Expect(cn).To(BeNil()) + Expect(err).ToNot(BeNil()) + }) +})