GODRIVER-1613 fix how pool handles MaxPoolSize (#413)

Author: iwysiu

internal/testutil/assert/assert.go

@@ -7,9 +7,11 @@
 package assert
 
 import (
+	"errors"
 	"reflect"
 	"sync"
 	"testing"
+	"time"
 
 	"github.com/google/go-cmp/cmp"
 )
@@ -23,6 +25,7 @@ var errorCompareFn = func(e1, e2 error) bool {
 	return e1.Error() == e2.Error()
 }
 var errorCompareOpts = cmp.Options{cmp.Comparer(errorCompareFn)}
+var ErrCallbackTimedOut = errors.New("callback timed out")
 
 // RegisterOpts registers go-cmp options for a type. These options will be used when comparing two objects for equality.
 func RegisterOpts(t reflect.Type, opts ...cmp.Option) {
@@ -80,6 +83,29 @@ func NotNil(t testing.TB, obj interface{}, msg string, args ...interface{}) {
 	}
 }
 
+// RunWithTimeout runs the provided callback for a maximum of timeoutMS milliseconds. It returns the callback error
+// if the callback returned and ErrCallbackTimedOut if the timeout expired.
+func RunWithTimeout(callback func() error, timeout time.Duration) error {
+	done := make(chan struct{})
+	var err error
+	fullCallback := func() {
+		err = callback()
+		done <- struct{}{}
+	}
+
+	timer := time.NewTimer(timeout)
+	defer timer.Stop()
+
+	go fullCallback()
+
+	select {
+	case <-done:
+		return err
+	case <-timer.C:
+		return ErrCallbackTimedOut
+	}
+}
+
 func getCmpOpts(obj interface{}) cmp.Options {
 	opts, ok := cmpOpts.Load(reflect.TypeOf(obj))
 	if ok {

x/mongo/driver/topology/connection.go

@@ -297,19 +297,15 @@ func (c *connection) readWireMessage(ctx context.Context, dst []byte) ([]byte, e
 }
 
 func (c *connection) close() error {
-	if atomic.LoadInt32(&c.connected) != connected {
+	if !atomic.CompareAndSwapInt32(&c.connected, connected, disconnected) {
 		return nil
 	}
 
 	var err error
 	if c.nc != nil {
 		err = c.nc.Close()
 	}
-	atomic.StoreInt32(&c.connected, disconnected)
 
-	if c.pool != nil {
-		_ = c.pool.removeConnection(c)
-	}
 	return err
 }
 
@@ -458,9 +454,7 @@ func (c *Connection) Close() error {
 	if c.connection == nil {
 		return nil
 	}
-	if c.s != nil {
-		defer c.s.sem.Release(1)
-	}
+
 	err := c.pool.put(c.connection)
 	c.connection = nil
 	return err
@@ -473,10 +467,9 @@ func (c *Connection) Expire() error {
 	if c.connection == nil {
 		return nil
 	}
-	if c.s != nil {
-		c.s.sem.Release(1)
-	}
-	err := c.close()
+
+	_ = c.close()
+	err := c.pool.put(c.connection)
 	c.connection = nil
 	return err
 }

x/mongo/driver/topology/pool.go

@@ -8,12 +8,14 @@ package topology
 
 import (
 	"context"
+	"math"
 	"sync"
 	"sync/atomic"
 	"time"
 
 	"go.mongodb.org/mongo-driver/event"
 	"go.mongodb.org/mongo-driver/x/mongo/driver/address"
+	"golang.org/x/sync/semaphore"
 )
 
 // ErrPoolConnected is returned from an attempt to connect an already connected pool
@@ -67,6 +69,7 @@ type pool struct {
 	connected int32 // Must be accessed using the sync/atomic package.
 	nextid    uint64
 	opened    map[uint64]*connection // opened holds all of the currently open connections.
+	sem       *semaphore.Weighted
 	sync.Mutex
 }
 
@@ -116,6 +119,7 @@ func connectionCloseFunc(v interface{}) {
 		return
 	}
 
+	_ = c.pool.removeConnection(c)
 	go func() {
 		_ = c.pool.closeConnection(c)
 	}()
@@ -141,16 +145,23 @@ func newPool(config poolConfig, connOpts ...ConnectionOption) (*pool, error) {
 		opts = append(opts, WithIdleTimeout(func(_ time.Duration) time.Duration { return config.MaxIdleTime }))
 	}
 
+	var maxConns = config.MaxPoolSize
+	if maxConns == 0 {
+		maxConns = math.MaxInt64
+	}
+
 	pool := &pool{
 		address:   config.Address,
 		monitor:   config.PoolMonitor,
 		connected: disconnected,
 		opened:    make(map[uint64]*connection),
 		opts:      opts,
+		sem:       semaphore.NewWeighted(int64(maxConns)),
 	}
 
 	// we do not pass in config.MaxPoolSize because we manage the max size at this level rather than the resource pool level
 	rpc := resourcePoolConfig{
+		MaxSize:          maxConns,
 		MinSize:          config.MinPoolSize,
 		MaintainInterval: maintainInterval,
 		ExpiredFn:        connectionExpiredFunc,
@@ -162,7 +173,7 @@ func newPool(config poolConfig, connOpts ...ConnectionOption) (*pool, error) {
 		pool.monitor.Event(&event.PoolEvent{
 			Type: event.PoolCreated,
 			PoolOptions: &event.MonitorPoolOptions{
-				MaxPoolSize:        config.MaxPoolSize,
+				MaxPoolSize:        rpc.MaxSize,
 				MinPoolSize:        rpc.MinSize,
 				WaitQueueTimeoutMS: uint64(config.MaxIdleTime) / uint64(time.Millisecond),
 			},
@@ -249,9 +260,11 @@ func (p *pool) disconnect(ctx context.Context) error {
 				Reason:       event.ReasonPoolClosed,
 			})
 		}
+		_ = p.removeConnection(pc)
 		_ = p.closeConnection(pc) // We don't care about errors while closing the connection.
 	}
 	atomic.StoreInt32(&p.connected, disconnected)
+	p.conns.clearTotal()
 
 	if p.monitor != nil {
 		p.monitor.Event(&event.PoolEvent{
@@ -305,7 +318,6 @@ func (p *pool) makeNewConnection(ctx context.Context) (*connection, string, erro
 
 // Checkout returns a connection from the pool
 func (p *pool) get(ctx context.Context) (*connection, error) {
-
 	if ctx == nil {
 		ctx = context.Background()
 	}
@@ -321,81 +333,120 @@ func (p *pool) get(ctx context.Context) (*connection, error) {
 		return nil, ErrPoolDisconnected
 	}
 
-	connVal := p.conns.Get()
-	if c, ok := connVal.(*connection); ok && connVal != nil {
-		// call connect if not connected
-		if atomic.LoadInt32(&c.connected) == initialized {
-			c.connect(ctx)
+	err := p.sem.Acquire(ctx, 1)
+	if err != nil {
+		if p.monitor != nil {
+			p.monitor.Event(&event.PoolEvent{
+				Type:    event.GetFailed,
+				Address: p.address.String(),
+				Reason:  event.ReasonTimedOut,
+			})
 		}
+		return nil, ErrWaitQueueTimeout
+	}
 
-		err := c.wait()
-		if err != nil {
+	// This loop is so that we don't end up with more than maxPoolSize connections if p.conns.Maintain runs between
+	// calling p.conns.Get() and making the new connection
+	for {
+		if atomic.LoadInt32(&p.connected) != connected {
 			if p.monitor != nil {
 				p.monitor.Event(&event.PoolEvent{
 					Type:    event.GetFailed,
 					Address: p.address.String(),
-					Reason:  event.ReasonConnectionErrored,
+					Reason:  event.ReasonPoolClosed,
 				})
 			}
-			return nil, err
+			p.sem.Release(1)
+			return nil, ErrPoolDisconnected
 		}
 
-		if p.monitor != nil {
-			p.monitor.Event(&event.PoolEvent{
-				Type:         event.GetSucceeded,
-				Address:      p.address.String(),
-				ConnectionID: c.poolID,
-			})
-		}
-		return c, nil
-	}
+		connVal := p.conns.Get()
+		if c, ok := connVal.(*connection); ok && connVal != nil {
+			// call connect if not connected
+			if atomic.LoadInt32(&c.connected) == initialized {
+				c.connect(ctx)
+			}
 
-	select {
-	case <-ctx.Done():
-		if p.monitor != nil {
-			p.monitor.Event(&event.PoolEvent{
-				Type:    event.GetFailed,
-				Address: p.address.String(),
-				Reason:  event.ReasonTimedOut,
-			})
-		}
-		return nil, ctx.Err()
-	default:
-		c, reason, err := p.makeNewConnection(ctx)
+			err := c.wait()
+			if err != nil {
+				if p.monitor != nil {
+					p.monitor.Event(&event.PoolEvent{
+						Type:    event.GetFailed,
+						Address: p.address.String(),
+						Reason:  event.ReasonConnectionErrored,
+					})
+				}
+				p.conns.decrementTotal()
+				p.sem.Release(1)
+				return nil, err
+			}
 
-		if err != nil {
 			if p.monitor != nil {
 				p.monitor.Event(&event.PoolEvent{
-					Type:    event.GetFailed,
-					Address: p.address.String(),
-					Reason:  reason,
+					Type:         event.GetSucceeded,
+					Address:      p.address.String(),
+					ConnectionID: c.poolID,
 				})
 			}
-			return nil, err
+			return c, nil
 		}
 
-		c.connect(ctx)
-		// wait for conn to be connected
-		err = c.wait()
-		if err != nil {
+		select {
+		case <-ctx.Done():
 			if p.monitor != nil {
 				p.monitor.Event(&event.PoolEvent{
 					Type:    event.GetFailed,
 					Address: p.address.String(),
-					Reason:  reason,
+					Reason:  event.ReasonTimedOut,
 				})
 			}
-			return nil, err
-		}
+			p.sem.Release(1)
+			return nil, ctx.Err()
+		default:
+			made := p.conns.incrementTotal()
+			if !made {
+				continue
+			}
+			c, reason, err := p.makeNewConnection(ctx)
+
+			if err != nil {
+				if p.monitor != nil {
+					p.monitor.Event(&event.PoolEvent{
+						Type:    event.GetFailed,
+						Address: p.address.String(),
+						Reason:  reason,
+					})
+				}
+				p.sem.Release(1)
+				p.conns.decrementTotal()
+				return nil, err
+			}
 
-		if p.monitor != nil {
-			p.monitor.Event(&event.PoolEvent{
-				Type:         event.GetSucceeded,
-				Address:      p.address.String(),
-				ConnectionID: c.poolID,
-			})
+			c.connect(ctx)
+			// wait for conn to be connected
+			err = c.wait()
+			if err != nil {
+				if p.monitor != nil {
+					p.monitor.Event(&event.PoolEvent{
+						Type:    event.GetFailed,
+						Address: p.address.String(),
+						Reason:  reason,
+					})
+				}
+				p.sem.Release(1)
+				p.conns.decrementTotal()
+				return nil, err
+			}
+
+			if p.monitor != nil {
+				p.monitor.Event(&event.PoolEvent{
+					Type:         event.GetSucceeded,
+					Address:      p.address.String(),
+					ConnectionID: c.poolID,
+				})
+			}
+			return c, nil
 		}
-		return c, nil
 	}
 }
 
@@ -405,9 +456,6 @@ func (p *pool) closeConnection(c *connection) error {
 	if c.pool != p {
 		return ErrWrongPool
 	}
-	p.Lock()
-	delete(p.opened, c.poolID)
-	p.Unlock()
 
 	if atomic.LoadInt32(&c.connected) == connected {
 		c.closeConnectContext()
@@ -441,8 +489,10 @@ func (p *pool) removeConnection(c *connection) error {
 }
 
 // put returns a connection to this pool. If the pool is connected, the connection is not
-// stale, and there is space in the cache, the connection is returned to the cache.
+// stale, and there is space in the cache, the connection is returned to the cache. This
+// assumes that the connection has already been counted in p.conns.totalSize.
 func (p *pool) put(c *connection) error {
+	defer p.sem.Release(1)
 	if p.monitor != nil {
 		var cid uint64
 		var addr string