Fix race conditions and leaks in connection pool

GODRIVER-388
GODRIVER-389
GODRIVER-390

Change-Id: I1390c184beb379d4cc775ba07b5d222342dc957e

Author: skriptble

core/connection/pool.go

@@ -123,7 +123,7 @@ loop:
 		select {
 		case pc := <-p.conns:
 			// This error would be overwritten by the semaphore
-			_ = pc.Close()
+			_ = p.closeConnection(pc)
 		default:
 			break loop
 		}
@@ -155,24 +155,30 @@ func (p *pool) Get(ctx context.Context) (Connection, *description.Server, error)
 		return nil, nil, ErrPoolClosed
 	}
 
+	err := p.sem.Acquire(ctx, 1)
+	if err != nil {
+		return nil, nil, err
+	}
+	return p.get(ctx)
+}
+
+func (p *pool) get(ctx context.Context) (Connection, *description.Server, error) {
 	g := atomic.LoadUint64(&p.generation)
 	select {
 	case c := <-p.conns:
 		if c.Expired() {
 			go p.closeConnection(c)
-			return p.Get(ctx)
+			return p.get(ctx)
 		}
 
-		return &acquired{Connection: c}, nil, nil
+		return &acquired{Connection: c, sem: p.sem}, nil, nil
 	case <-ctx.Done():
+		p.sem.Release(1)
 		return nil, nil, ctx.Err()
 	default:
-		err := p.sem.Acquire(ctx, 1)
-		if err != nil {
-			return nil, nil, err
-		}
 		c, desc, err := New(ctx, p.address, p.opts...)
 		if err != nil {
+			p.sem.Release(1)
 			return nil, nil, err
 		}
 
@@ -185,20 +191,20 @@ func (p *pool) Get(ctx context.Context) (Connection, *description.Server, error)
 		p.Lock()
 		if atomic.LoadInt32(&p.connected) != connected {
 			p.Unlock()
+			p.sem.Release(1)
 			p.closeConnection(pc)
 			return nil, nil, ErrPoolClosed
 		}
 		defer p.Unlock()
 		p.inflight[pc.id] = pc
-		return &acquired{Connection: pc}, desc, nil
+		return &acquired{Connection: pc, sem: p.sem}, desc, nil
 	}
 }
 
 func (p *pool) closeConnection(pc *pooledConnection) error {
 	if !atomic.CompareAndSwapInt32(&pc.closed, 0, 1) {
 		return nil
 	}
-	pc.p.sem.Release(1)
 	p.Lock()
 	delete(p.inflight, pc.id)
 	p.Unlock()
@@ -241,6 +247,7 @@ func (pc *pooledConnection) Expired() bool {
 type acquired struct {
 	Connection
 
+	sem *semaphore.Weighted
 	sync.Mutex
 }
 
@@ -269,6 +276,7 @@ func (a *acquired) Close() error {
 		return nil
 	}
 	err := a.Connection.Close()
+	a.sem.Release(1)
 	a.Connection = nil
 	return err
 }

core/connection/pool_test.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"errors"
 	"net"
+	"runtime"
+	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -119,6 +121,8 @@ func TestPool(t *testing.T) {
 				t.Errorf("Should have closed 3 connections, but didn't. got %d; want %d", d.lenclosed(), 3)
 			}
 			close(cleanup)
+			err = conns[2].Close()
+			noerr(t, err)
 			ok := p.(*pool).sem.TryAcquire(int64(p.(*pool).capacity))
 			if !ok {
 				t.Errorf("clean shutdown should acquire and release semaphore, but semaphore still held")
@@ -478,6 +482,9 @@ func TestPool(t *testing.T) {
 			err = p.Drain()
 			noerr(t, err)
 
+			err = conns[1].Close()
+			noerr(t, err)
+
 			ctx, cancel = context.WithTimeout(context.Background(), 10*time.Millisecond)
 			defer cancel()
 			c, _, err = p.Get(ctx)
@@ -487,6 +494,129 @@ func TestPool(t *testing.T) {
 			}
 			close(cleanup)
 		})
+		t.Run("Cannot starve connection request", func(t *testing.T) {
+			cleanup := make(chan struct{})
+			address := bootstrapConnections(t, 3, func(nc net.Conn) {
+				<-cleanup
+				nc.Close()
+			})
+			d := newdialer(&net.Dialer{})
+			p, err := NewPool(addr.Addr(address.String()), 1, 1, WithDialer(func(Dialer) Dialer { return d }))
+			noerr(t, err)
+			err = p.Connect(context.Background())
+			noerr(t, err)
+			conn, _, err := p.Get(context.Background())
+			if d.lenopened() != 1 {
+				t.Errorf("Should have opened 1 connections, but didn't. got %d; want %d", d.lenopened(), 1)
+			}
+
+			var wg sync.WaitGroup
+
+			wg.Add(1)
+			ch := make(chan struct{})
+			go func() {
+				ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+				defer cancel()
+				ch <- struct{}{}
+				_, _, err := p.Get(ctx)
+				if err != nil {
+					t.Errorf("Should not be able to starve connection request, but got error: %v", err)
+				}
+				wg.Done()
+			}()
+			<-ch
+			runtime.Gosched()
+			err = conn.Close()
+			noerr(t, err)
+			wg.Wait()
+			close(cleanup)
+		})
+		t.Run("Does not leak permit from failure to dial connection", func(t *testing.T) {
+			cleanup := make(chan struct{})
+			address := bootstrapConnections(t, 0, func(nc net.Conn) {
+				<-cleanup
+				nc.Close()
+			})
+			close(cleanup)
+			want := errors.New("dialing error")
+			p, err := NewPool(
+				addr.Addr(address.String()), 1, 2,
+				WithDialer(
+					func(Dialer) Dialer {
+						return DialerFunc(func(ctx context.Context, network, address string) (net.Conn, error) {
+							return nil, want
+						})
+					}),
+			)
+			noerr(t, err)
+			err = p.Connect(context.Background())
+			noerr(t, err)
+			_, _, err = p.Get(context.Background())
+			if err != want {
+				t.Errorf("Expected dial failure but got: %v", err)
+			}
+			ok := p.(*pool).sem.TryAcquire(int64(p.(*pool).capacity))
+			if !ok {
+				t.Errorf("Dial failure should not leak semaphore permit")
+			} else {
+				p.(*pool).sem.Release(int64(p.(*pool).capacity))
+			}
+		})
+		t.Run("Does not leak permit from cancelled context", func(t *testing.T) {
+			cleanup := make(chan struct{})
+			address := bootstrapConnections(t, 1, func(nc net.Conn) {
+				<-cleanup
+				nc.Close()
+			})
+			close(cleanup)
+			d := newdialer(&net.Dialer{})
+			p, err := NewPool(addr.Addr(address.String()), 1, 2, WithDialer(func(Dialer) Dialer { return d }))
+			noerr(t, err)
+			err = p.Connect(context.Background())
+			noerr(t, err)
+			ctx, cancel := context.WithCancel(context.Background())
+			cancel()
+			_, _, err = p.Get(ctx)
+			if err != context.Canceled {
+				t.Errorf("Expected context cancelled error. got %v; want %v", err, context.Canceled)
+			}
+			ok := p.(*pool).sem.TryAcquire(int64(p.(*pool).capacity))
+			if !ok {
+				t.Errorf("Canceled context should not leak semaphore permit")
+			} else {
+				p.(*pool).sem.Release(int64(p.(*pool).capacity))
+			}
+		})
+		t.Run("Get does not acquire multiple permits", func(t *testing.T) {
+			cleanup := make(chan struct{})
+			address := bootstrapConnections(t, 2, func(nc net.Conn) {
+				<-cleanup
+				nc.Close()
+			})
+			close(cleanup)
+			d := newdialer(&net.Dialer{})
+			p, err := NewPool(addr.Addr(address.String()), 1, 2, WithDialer(func(Dialer) Dialer { return d }))
+			noerr(t, err)
+			err = p.Connect(context.Background())
+			noerr(t, err)
+			c, _, err := p.Get(context.Background())
+			noerr(t, err)
+			err = c.Close()
+			noerr(t, err)
+
+			p.Drain()
+
+			c, _, err = p.Get(context.Background())
+			noerr(t, err)
+			err = c.Close()
+			noerr(t, err)
+			ok := p.(*pool).sem.TryAcquire(int64(p.(*pool).capacity))
+			if !ok {
+				t.Errorf("Get should not acquire multiple permits (when expired conn in idle pool)")
+			} else {
+				p.(*pool).sem.Release(int64(p.(*pool).capacity))
+			}
+		})
 	})
 	t.Run("Connection", func(t *testing.T) {
 		t.Run("Connection Close Does Not Error After Pool Is Disconnected", func(t *testing.T) {