Retry BadConnError

internal/pool/pool_single.go

@@ -12,7 +12,17 @@ const (
 	stateClosed  = 2
 )
 
-var ErrBadConn = fmt.Errorf("pg: Conn is in a bad state")
+type BadConnError struct {
+	wrapped error
+}
+
+func (e BadConnError) Error() string {
+	return "pg: Conn is in a bad state"
+}
+
+func (e BadConnError) Unwrap() error {
+	return e.wrapped
+}
 
 type SingleConnPool struct {
 	pool Pooler
@@ -20,8 +30,8 @@ type SingleConnPool struct {
 	state uint32 // atomic
 	ch    chan *Conn
 
-	level       int32  // atomic
-	_hasBadConn uint32 // atomic
+	level         int32 // atomic
+	_badConnError atomic.Value
 }
 
 var _ Pooler = (*SingleConnPool)(nil)
@@ -66,10 +76,10 @@ func (p *SingleConnPool) Get(c context.Context) (*Conn, error) {
 			if atomic.CompareAndSwapUint32(&p.state, stateDefault, stateInited) {
 				return cn, nil
 			}
-			p.pool.Remove(cn)
+			p.pool.Remove(cn, ErrClosed)
 		case stateInited:
-			if p.hasBadConn() {
-				return nil, ErrBadConn
+			if err := p.badConnError(); err != nil {
+				return nil, err
 			}
 			cn, ok := <-p.ch
 			if !ok {
@@ -95,20 +105,20 @@ func (p *SingleConnPool) Put(cn *Conn) {
 }
 
 func (p *SingleConnPool) freeConn(cn *Conn) {
-	if p.hasBadConn() {
-		p.pool.Remove(cn)
+	if err := p.badConnError(); err != nil {
+		p.pool.Remove(cn, err)
 	} else {
 		p.pool.Put(cn)
 	}
 }
 
-func (p *SingleConnPool) Remove(cn *Conn) {
+func (p *SingleConnPool) Remove(cn *Conn, reason error) {
 	defer func() {
 		if recover() != nil {
-			p.pool.Remove(cn)
+			p.pool.Remove(cn, ErrClosed)
 		}
 	}()
-	atomic.StoreUint32(&p._hasBadConn, 1)
+	p._badConnError.Store(BadConnError{wrapped: reason})
 	p.ch <- cn
 }
 
@@ -158,7 +168,7 @@ func (p *SingleConnPool) Close() error {
 }
 
 func (p *SingleConnPool) Reset() error {
-	if !atomic.CompareAndSwapUint32(&p._hasBadConn, 1, 0) {
+	if p.badConnError() == nil {
 		return nil
 	}
 
@@ -167,7 +177,8 @@ func (p *SingleConnPool) Reset() error {
 		if !ok {
 			return ErrClosed
 		}
-		p.pool.Remove(cn)
+		p.pool.Remove(cn, ErrClosed)
+		p._badConnError.Store(nil)
 	default:
 		return fmt.Errorf("pg: SingleConnPool does not have a Conn")
 	}
@@ -180,6 +191,9 @@ func (p *SingleConnPool) Reset() error {
 	return nil
 }
 
-func (p *SingleConnPool) hasBadConn() bool {
-	return atomic.LoadUint32(&p._hasBadConn) == 1
+func (p *SingleConnPool) badConnError() error {
+	if v := p._badConnError.Load(); v != nil {
+		return v.(BadConnError)
+	}
+	return nil
 }