Merge pull request #480 from Zariel/remove-simple-pool

Remove SimpleConnPool, remove custom pools

cass1batch_test.go

@@ -10,7 +10,7 @@ import (
 func TestProto1BatchInsert(t *testing.T) {
 	session := createSession(t)
 	if err := createTable(session, "CREATE TABLE large (id int primary key)"); err != nil {
-		t.Fatal("create table:", err)
+		t.Fatal(err)
 	}
 	defer session.Close()
 

cassandra_test.go

@@ -96,7 +96,7 @@ func createKeyspace(tb testing.TB, cluster *ClusterConfig, keyspace string) {
 	}
 
 	// should reuse the same conn apparently
-	conn := session.Pool.Pick(nil)
+	conn := session.pool.Pick(nil)
 	if conn == nil {
 		tb.Fatal("no connections available in the pool")
 	}
@@ -189,7 +189,9 @@ func TestRingDiscovery(t *testing.T) {
 		time.Sleep(*flagAutoWait)
 	}
 
-	size := len(session.Pool.(*SimplePool).connPool)
+	session.pool.mu.RLock()
+	size := len(session.pool.hostConnPools)
+	session.pool.mu.RUnlock()
 
 	if *clusterSize != size {
 		t.Logf("WARN: Expected a cluster size of %d, but actual size was %d", *clusterSize, size)
@@ -1139,7 +1141,7 @@ func injectInvalidPreparedStatement(t *testing.T, session *Session, table string
 		t.Fatal("create:", err)
 	}
 	stmt := "INSERT INTO " + table + " (foo, bar) VALUES (?, 7)"
-	conn := session.Pool.Pick(nil)
+	conn := session.pool.Pick(nil)
 	flight := new(inflightPrepare)
 	stmtsLRU.Lock()
 	stmtsLRU.lru.Add(conn.addr+stmt, flight)
@@ -1165,7 +1167,7 @@ func injectInvalidPreparedStatement(t *testing.T, session *Session, table string
 
 func TestMissingSchemaPrepare(t *testing.T) {
 	s := createSession(t)
-	conn := s.Pool.Pick(nil)
+	conn := s.pool.Pick(nil)
 	defer s.Close()
 
 	insertQry := &Query{stmt: "INSERT INTO invalidschemaprep (val) VALUES (?)", values: []interface{}{5}, cons: s.cons,
@@ -1214,7 +1216,7 @@ func TestQueryInfo(t *testing.T) {
 	session := createSession(t)
 	defer session.Close()
 
-	conn := session.Pool.Pick(nil)
+	conn := session.pool.Pick(nil)
 	info, err := conn.prepareStatement("SELECT release_version, host_id FROM system.local WHERE key = ?", nil)
 
 	if err != nil {
@@ -2016,7 +2018,7 @@ func TestRoutingKey(t *testing.T) {
 // Integration test of the token-aware policy-based connection pool
 func TestTokenAwareConnPool(t *testing.T) {
 	cluster := createCluster()
-	cluster.ConnPoolType = NewTokenAwareConnPool
+	cluster.PoolConfig.HostSelectionPolicy = TokenAwareHostPolicy(RoundRobinHostPolicy())
 	cluster.DiscoverHosts = true
 
 	// Drop and re-create the keyspace once. Different tests should use their own
@@ -2037,8 +2039,8 @@ func TestTokenAwareConnPool(t *testing.T) {
 		time.Sleep(*flagAutoWait)
 	}
 
-	if session.Pool.Size() != cluster.NumConns*len(cluster.Hosts) {
-		t.Errorf("Expected pool size %d but was %d", cluster.NumConns*len(cluster.Hosts), session.Pool.Size())
+	if session.pool.Size() != cluster.NumConns*len(cluster.Hosts) {
+		t.Errorf("Expected pool size %d but was %d", cluster.NumConns*len(cluster.Hosts), session.pool.Size())
 	}
 
 	if err := createTable(session, "CREATE TABLE test_token_aware (id int, data text, PRIMARY KEY (id))"); err != nil {
@@ -2077,7 +2079,7 @@ func TestStream0(t *testing.T) {
 			break
 		}
 
-		conn = session.Pool.Pick(nil)
+		conn = session.pool.Pick(nil)
 	}
 
 	if conn == nil {
@@ -2116,7 +2118,7 @@ func TestNegativeStream(t *testing.T) {
 			break
 		}
 
-		conn = session.Pool.Pick(nil)
+		conn = session.pool.Pick(nil)
 	}
 
 	if conn == nil {
@@ -2222,7 +2224,7 @@ func TestLexicalUUIDType(t *testing.T) {
 // Issue 475
 func TestSessionBindRoutingKey(t *testing.T) {
 	cluster := createCluster()
-	cluster.ConnPoolType = NewTokenAwareConnPool
+	cluster.PoolConfig.HostSelectionPolicy = TokenAwareHostPolicy(RoundRobinHostPolicy())
 
 	session := createSessionFromCluster(cluster, t)
 	defer session.Close()

cluster.go

@@ -50,6 +50,33 @@ type DiscoveryConfig struct {
 	Sleep time.Duration
 }
 
+// PoolConfig configures the connection pool used by the driver, it defaults to
+// using a round robbin host selection policy and a round robbin connection selection
+// policy for each host.
+type PoolConfig struct {
+	// HostSelectionPolicy sets the policy for selecting which host to use for a
+	// given query (default: RoundRobinHostPolicy())
+	HostSelectionPolicy HostSelectionPolicy
+
+	// ConnSelectionPolicy sets the policy factory for selecting a connection to use for
+	// each host for a query (default: RoundRobinConnPolicy())
+	ConnSelectionPolicy func() ConnSelectionPolicy
+}
+
+func (p PoolConfig) buildPool(cfg *ClusterConfig) (*policyConnPool, error) {
+	hostSelection := p.HostSelectionPolicy
+	if hostSelection == nil {
+		hostSelection = RoundRobinHostPolicy()
+	}
+
+	connSelection := p.ConnSelectionPolicy
+	if connSelection == nil {
+		connSelection = RoundRobinConnPolicy()
+	}
+
+	return newPolicyConnPool(cfg, hostSelection, connSelection)
+}
+
 // ClusterConfig is a struct to configure the default cluster implementation
 // of gocoql. It has a varity of attributes that can be used to modify the
 // behavior to fit the most common use cases. Applications that requre a
@@ -68,7 +95,6 @@ type ClusterConfig struct {
 	Authenticator     Authenticator     // authenticator (default: nil)
 	RetryPolicy       RetryPolicy       // Default retry policy to use for queries (default: 0)
 	SocketKeepalive   time.Duration     // The keepalive period to use, enabled if > 0 (default: 0)
-	ConnPoolType      NewPoolFunc       // The function used to create the connection pool for the session (default: NewSimplePool)
 	DiscoverHosts     bool              // If set, gocql will attempt to automatically discover other members of the Cassandra cluster (default: false)
 	MaxPreparedStmts  int               // Sets the maximum cache size for prepared statements globally for gocql (default: 1000)
 	MaxRoutingKeyInfo int               // Sets the maximum cache size for query info about statements for each session (default: 1000)
@@ -77,6 +103,9 @@ type ClusterConfig struct {
 	Discovery         DiscoveryConfig
 	SslOpts           *SslOptions
 	DefaultTimestamp  bool // Sends a client side timestamp for all requests which overrides the timestamp at which it arrives at the server. (default: true, only enabled for protocol 3 and above)
+	// PoolConfig configures the underlying connection pool, allowing the
+	// configuration of host selection and connection selection policies.
+	PoolConfig PoolConfig
 }
 
 // NewCluster generates a new config for the default cluster implementation.
@@ -89,7 +118,6 @@ func NewCluster(hosts ...string) *ClusterConfig {
 		Port:              9042,
 		NumConns:          2,
 		Consistency:       Quorum,
-		ConnPoolType:      NewSimplePool,
 		DiscoverHosts:     false,
 		MaxPreparedStmts:  defaultMaxPreparedStmts,
 		MaxRoutingKeyInfo: 1000,

conn_test.go

@@ -259,8 +259,7 @@ func TestConnClosing(t *testing.T) {
 	wg.Wait()
 
 	time.Sleep(1 * time.Second) //Sleep so the fillPool can complete.
-	pool := db.Pool.(ConnectionPool)
-	conns := pool.Size()
+	conns := db.pool.Size()
 
 	if conns != numConns {
 		t.Errorf("Expected to have %d connections but have %d", numConns, conns)
@@ -390,7 +389,8 @@ func TestRoundRobinConnPoolRoundRobin(t *testing.T) {
 
 	// create a new cluster using the policy-based round robin conn pool
 	cluster := NewCluster(addrs...)
-	cluster.ConnPoolType = NewRoundRobinConnPool
+	cluster.PoolConfig.HostSelectionPolicy = RoundRobinHostPolicy()
+	cluster.PoolConfig.ConnSelectionPolicy = RoundRobinConnPolicy()
 
 	db, err := cluster.CreateSession()
 	if err != nil {
@@ -420,7 +420,7 @@ func TestRoundRobinConnPoolRoundRobin(t *testing.T) {
 
 	// wait for the pool to drain
 	time.Sleep(100 * time.Millisecond)
-	size := db.Pool.Size()
+	size := db.pool.Size()
 	if size != 0 {
 		t.Errorf("connection pool did not drain, still contains %d connections", size)
 	}
@@ -450,7 +450,8 @@ func TestPolicyConnPoolSSL(t *testing.T) {
 	defer srv.Stop()
 
 	cluster := createTestSslCluster(srv.Address, defaultProto, true)
-	cluster.ConnPoolType = NewRoundRobinConnPool
+	cluster.PoolConfig.HostSelectionPolicy = RoundRobinHostPolicy()
+	cluster.PoolConfig.ConnSelectionPolicy = RoundRobinConnPolicy()
 
 	db, err := cluster.CreateSession()
 	if err != nil {
@@ -465,7 +466,7 @@ func TestPolicyConnPoolSSL(t *testing.T) {
 
 	// wait for the pool to drain
 	time.Sleep(100 * time.Millisecond)
-	size := db.Pool.Size()
+	size := db.pool.Size()
 	if size != 0 {
 		t.Errorf("connection pool did not drain, still contains %d connections", size)
 	}

connectionpool.go

@@ -17,89 +17,6 @@ import (
 	"time"
 )
 
-/*ConnectionPool represents the interface gocql will use to work with a collection of connections.
-
-Purpose
-
-The connection pool in gocql opens and closes connections as well as selects an available connection
-for gocql to execute a query against. The pool is also responsible for handling connection errors that
-are caught by the connection experiencing the error.
-
-A connection pool should make a copy of the variables used from the ClusterConfig provided to the pool
-upon creation. ClusterConfig is a pointer and can be modified after the creation of the pool. This can
-lead to issues with variables being modified outside the expectations of the ConnectionPool type.
-
-Example of Single Connection Pool:
-
-	type SingleConnection struct {
-		conn *Conn
-		cfg *ClusterConfig
-	}
-
-	func NewSingleConnection(cfg *ClusterConfig) ConnectionPool {
-		addr := JoinHostPort(cfg.Hosts[0], cfg.Port)
-
-		connCfg := ConnConfig{
-			ProtoVersion:  cfg.ProtoVersion,
-			CQLVersion:    cfg.CQLVersion,
-			Timeout:       cfg.Timeout,
-			NumStreams:    cfg.NumStreams,
-			Compressor:    cfg.Compressor,
-			Authenticator: cfg.Authenticator,
-			Keepalive:     cfg.SocketKeepalive,
-		}
-		pool := SingleConnection{cfg:cfg}
-		pool.conn = Connect(addr,connCfg,pool)
-		return &pool
-	}
-
-	func (s *SingleConnection) HandleError(conn *Conn, err error, closed bool) {
-		if closed {
-			connCfg := ConnConfig{
-				ProtoVersion:  cfg.ProtoVersion,
-				CQLVersion:    cfg.CQLVersion,
-				Timeout:       cfg.Timeout,
-				NumStreams:    cfg.NumStreams,
-				Compressor:    cfg.Compressor,
-				Authenticator: cfg.Authenticator,
-				Keepalive:     cfg.SocketKeepalive,
-			}
-			s.conn = Connect(conn.Address(),connCfg,s)
-		}
-	}
-
-	func (s *SingleConnection) Pick(qry *Query) *Conn {
-		if s.conn.isClosed {
-			return nil
-		}
-		return s.conn
-	}
-
-	func (s *SingleConnection) Size() int {
-		return 1
-	}
-
-	func (s *SingleConnection) Close() {
-		s.conn.Close()
-	}
-
-This is a very simple example of a type that exposes the connection pool interface. To assign
-this type as the connection pool to use you would assign it to the ClusterConfig like so:
-
-		cluster := NewCluster("127.0.0.1")
-		cluster.ConnPoolType = NewSingleConnection
-		...
-		session, err := cluster.CreateSession()
-
-To see a more complete example of a ConnectionPool implementation please see the SimplePool type.
-*/
-type ConnectionPool interface {
-	SetHosts
-	Pick(*Query) *Conn
-	Size() int
-	Close()
-}
-
 // interface to implement to receive the host information
 type SetHosts interface {
 	SetHosts(hosts []HostInfo)
@@ -110,35 +27,6 @@ type SetPartitioner interface {
 	SetPartitioner(partitioner string)
 }
 
-//NewPoolFunc is the type used by ClusterConfig to create a pool of a specific type.
-type NewPoolFunc func(*ClusterConfig) (ConnectionPool, error)
-
-//SimplePool is the current implementation of the connection pool inside gocql. This
-//pool is meant to be a simple default used by gocql so users can get up and running
-//quickly.
-type SimplePool struct {
-	cfg      *ClusterConfig
-	hostPool *RoundRobin
-	connPool map[string]*RoundRobin
-	conns    map[*Conn]struct{}
-	keyspace string
-
-	hostMu sync.RWMutex
-	// this is the set of current hosts which the pool will attempt to connect to
-	hosts map[string]*HostInfo
-
-	// protects hostpool, connPoll, conns, quit
-	mu sync.Mutex
-
-	cFillingPool chan int
-
-	quit     bool
-	quitWait chan bool
-	quitOnce sync.Once
-
-	tlsConfig *tls.Config
-}
-
 func setupTLSConfig(sslOpts *SslOptions) (*tls.Config, error) {
 	// ca cert is optional
 	if sslOpts.CaPath != "" {
@@ -169,309 +57,6 @@ func setupTLSConfig(sslOpts *SslOptions) (*tls.Config, error) {
 	return &sslOpts.Config, nil
 }
 
-//NewSimplePool is the function used by gocql to create the simple connection pool.
-//This is the default if no other pool type is specified.
-func NewSimplePool(cfg *ClusterConfig) (ConnectionPool, error) {
-	pool := &SimplePool{
-		cfg:          cfg,
-		hostPool:     NewRoundRobin(),
-		connPool:     make(map[string]*RoundRobin),
-		conns:        make(map[*Conn]struct{}),
-		quitWait:     make(chan bool),
-		cFillingPool: make(chan int, 1),
-		keyspace:     cfg.Keyspace,
-		hosts:        make(map[string]*HostInfo),
-	}
-
-	for _, host := range cfg.Hosts {
-		// seed hosts have unknown topology
-		// TODO: Handle populating this during SetHosts
-		pool.hosts[host] = &HostInfo{Peer: host}
-	}
-
-	if cfg.SslOpts != nil {
-		config, err := setupTLSConfig(cfg.SslOpts)
-		if err != nil {
-			return nil, err
-		}
-		pool.tlsConfig = config
-	}
-
-	//Walk through connecting to hosts. As soon as one host connects
-	//defer the remaining connections to cluster.fillPool()
-	for i := 0; i < len(cfg.Hosts); i++ {
-		addr := JoinHostPort(cfg.Hosts[i], cfg.Port)
-
-		if pool.connect(addr) == nil {
-			pool.cFillingPool <- 1
-			go pool.fillPool()
-			break
-		}
-	}
-
-	return pool, nil
-}
-
-func (c *SimplePool) connect(addr string) error {
-
-	cfg := ConnConfig{
-		ProtoVersion:  c.cfg.ProtoVersion,
-		CQLVersion:    c.cfg.CQLVersion,
-		Timeout:       c.cfg.Timeout,
-		NumStreams:    c.cfg.NumStreams,
-		Compressor:    c.cfg.Compressor,
-		Authenticator: c.cfg.Authenticator,
-		Keepalive:     c.cfg.SocketKeepalive,
-		tlsConfig:     c.tlsConfig,
-	}
-
-	conn, err := Connect(addr, cfg, c)
-	if err != nil {
-		log.Printf("connect: failed to connect to %q: %v", addr, err)
-		return err
-	}
-
-	return c.addConn(conn)
-}
-
-func (c *SimplePool) addConn(conn *Conn) error {
-	c.mu.Lock()
-	defer c.mu.Unlock()
-	if c.quit {
-		conn.Close()
-		return nil
-	}
-
-	//Set the connection's keyspace if any before adding it to the pool
-	if c.keyspace != "" {
-		if err := conn.UseKeyspace(c.keyspace); err != nil {
-			log.Printf("error setting connection keyspace. %v", err)
-			conn.Close()
-			return err
-		}
-	}
-
-	connPool := c.connPool[conn.Address()]
-	if connPool == nil {
-		connPool = NewRoundRobin()
-		c.connPool[conn.Address()] = connPool
-		c.hostPool.AddNode(connPool)
-	}
-
-	connPool.AddNode(conn)
-	c.conns[conn] = struct{}{}
-
-	return nil
-}
-
-//fillPool manages the pool of connections making sure that each host has the correct
-//amount of connections defined. Also the method will test a host with one connection
-//instead of flooding the host with number of connections defined in the cluster config
-func (c *SimplePool) fillPool() {
-	//Debounce large amounts of requests to fill pool
-	select {
-	case <-time.After(1 * time.Millisecond):
-		return
-	case <-c.cFillingPool:
-		defer func() { c.cFillingPool <- 1 }()
-	}
-
-	c.mu.Lock()
-	isClosed := c.quit
-	c.mu.Unlock()
-	//Exit if cluster(session) is closed
-	if isClosed {
-		return
-	}
-
-	c.hostMu.RLock()
-
-	//Walk through list of defined hosts
-	var wg sync.WaitGroup
-	for host := range c.hosts {
-		addr := JoinHostPort(host, c.cfg.Port)
-
-		numConns := 1
-		//See if the host already has connections in the pool
-		c.mu.Lock()
-		conns, ok := c.connPool[addr]
-		c.mu.Unlock()
-
-		if ok {
-			//if the host has enough connections just exit
-			numConns = conns.Size()
-			if numConns >= c.cfg.NumConns {
-				continue
-			}
-		} else {
-			//See if the host is reachable
-			if err := c.connect(addr); err != nil {
-				continue
-			}
-		}
-
-		//This is reached if the host is responsive and needs more connections
-		//Create connections for host synchronously to mitigate flooding the host.
-		wg.Add(1)
-		go func(a string, conns int) {
-			defer wg.Done()
-			for ; conns < c.cfg.NumConns; conns++ {
-				c.connect(a)
-			}
-		}(addr, numConns)
-	}
-
-	c.hostMu.RUnlock()
-
-	//Wait until we're finished connecting to each host before returning
-	wg.Wait()
-}
-
-// Should only be called if c.mu is locked
-func (c *SimplePool) removeConnLocked(conn *Conn) {
-	conn.Close()
-	connPool := c.connPool[conn.addr]
-	if connPool == nil {
-		return
-	}
-	connPool.RemoveNode(conn)
-	if connPool.Size() == 0 {
-		c.hostPool.RemoveNode(connPool)
-		delete(c.connPool, conn.addr)
-	}
-	delete(c.conns, conn)
-}
-
-func (c *SimplePool) removeConn(conn *Conn) {
-	c.mu.Lock()
-	defer c.mu.Unlock()
-	c.removeConnLocked(conn)
-}
-
-//HandleError is called by a Connection object to report to the pool an error has occured.
-//Logic is then executed within the pool to clean up the erroroneous connection and try to
-//top off the pool.
-func (c *SimplePool) HandleError(conn *Conn, err error, closed bool) {
-	if !closed {
-		// ignore all non-fatal errors
-		return
-	}
-	c.removeConn(conn)
-	c.mu.Lock()
-	poolClosed := c.quit
-	c.mu.Unlock()
-	if !poolClosed {
-		go c.fillPool() // top off pool.
-	}
-}
-
-//Pick selects a connection to be used by the query.
-func (c *SimplePool) Pick(qry *Query) *Conn {
-	//Check if connections are available
-	c.mu.Lock()
-	conns := len(c.conns)
-	c.mu.Unlock()
-
-	if conns == 0 {
-		//try to populate the pool before returning.
-		c.fillPool()
-	}
-
-	return c.hostPool.Pick(qry)
-}
-
-//Size returns the number of connections currently active in the pool
-func (p *SimplePool) Size() int {
-	p.mu.Lock()
-	conns := len(p.conns)
-	p.mu.Unlock()
-	return conns
-}
-
-//Close kills the pool and all associated connections.
-func (c *SimplePool) Close() {
-	c.quitOnce.Do(func() {
-		c.mu.Lock()
-		defer c.mu.Unlock()
-		c.quit = true
-		close(c.quitWait)
-		for conn := range c.conns {
-			c.removeConnLocked(conn)
-		}
-	})
-}
-
-func (c *SimplePool) SetHosts(hosts []HostInfo) {
-
-	c.hostMu.Lock()
-	toRemove := make(map[string]struct{})
-	for k := range c.hosts {
-		toRemove[k] = struct{}{}
-	}
-
-	for _, host := range hosts {
-		host := host
-		delete(toRemove, host.Peer)
-		// we already have it
-		if _, ok := c.hosts[host.Peer]; ok {
-			// TODO: Check rack, dc, token range is consistent, trigger topology change
-			// update stored host
-			continue
-		}
-
-		c.hosts[host.Peer] = &host
-	}
-
-	// can we hold c.mu whilst iterating this loop?
-	for addr := range toRemove {
-		c.removeHostLocked(addr)
-	}
-	c.hostMu.Unlock()
-
-	c.fillPool()
-}
-
-func (c *SimplePool) removeHostLocked(addr string) {
-	if _, ok := c.hosts[addr]; !ok {
-		return
-	}
-	delete(c.hosts, addr)
-
-	c.mu.Lock()
-	defer c.mu.Unlock()
-
-	if _, ok := c.connPool[addr]; !ok {
-		return
-	}
-
-	for conn := range c.conns {
-		if conn.Address() == addr {
-			c.removeConnLocked(conn)
-		}
-	}
-}
-
-//NewRoundRobinConnPool creates a connection pool which selects hosts by
-//round-robin, and then selects a connection for that host by round-robin.
-func NewRoundRobinConnPool(cfg *ClusterConfig) (ConnectionPool, error) {
-	return NewPolicyConnPool(
-		cfg,
-		NewRoundRobinHostPolicy(),
-		NewRoundRobinConnPolicy,
-	)
-}
-
-//NewTokenAwareConnPool creates a connection pool which selects hosts by
-//a token aware policy, and then selects a connection for that host by
-//round-robin.
-func NewTokenAwareConnPool(cfg *ClusterConfig) (ConnectionPool, error) {
-	return NewPolicyConnPool(
-		cfg,
-		NewTokenAwareHostPolicy(NewRoundRobinHostPolicy()),
-		NewRoundRobinConnPolicy,
-	)
-}
-
 type policyConnPool struct {
 	port     int
 	numConns int
@@ -484,18 +69,13 @@ type policyConnPool struct {
 	hostConnPools map[string]*hostConnPool
 }
 
-//Creates a policy based connection pool. This func isn't meant to be directly
-//used as a NewPoolFunc in ClusterConfig, instead a func should be created
-//which satisfies the NewPoolFunc type, which calls this func with the desired
-//hostPolicy and connPolicy; see NewRoundRobinConnPool or NewTokenAwareConnPool
-//for examples.
-func NewPolicyConnPool(
-	cfg *ClusterConfig,
-	hostPolicy HostSelectionPolicy,
-	connPolicy func() ConnSelectionPolicy,
-) (ConnectionPool, error) {
-	var err error
-	var tlsConfig *tls.Config
+func newPolicyConnPool(cfg *ClusterConfig, hostPolicy HostSelectionPolicy,
+	connPolicy func() ConnSelectionPolicy) (*policyConnPool, error) {
+
+	var (
+		err       error
+		tlsConfig *tls.Config
+	)
 
 	if cfg.SslOpts != nil {
 		tlsConfig, err = setupTLSConfig(cfg.SslOpts)
@@ -594,9 +174,12 @@ func (p *policyConnPool) Size() int {
 func (p *policyConnPool) Pick(qry *Query) *Conn {
 	nextHost := p.hostPolicy.Pick(qry)
 
+	var (
+		host *HostInfo
+		conn *Conn
+	)
+
 	p.mu.RLock()
-	var host *HostInfo
-	var conn *Conn
 	for conn == nil {
 		host = nextHost()
 		if host == nil {
@@ -639,14 +222,8 @@ type hostConnPool struct {
 	filling bool
 }
 
-func newHostConnPool(
-	host string,
-	port int,
-	size int,
-	connCfg ConnConfig,
-	keyspace string,
-	policy ConnSelectionPolicy,
-) *hostConnPool {
+func newHostConnPool(host string, port int, size int, connCfg ConnConfig,
+	keyspace string, policy ConnSelectionPolicy) *hostConnPool {
 
 	pool := &hostConnPool{
 		host:     host,

host_source.go

@@ -24,14 +24,10 @@ type ringDescriber struct {
 	closeChan       chan bool
 }
 
-func (r *ringDescriber) GetHosts() (
-	hosts []HostInfo,
-	partitioner string,
-	err error,
-) {
+func (r *ringDescriber) GetHosts() (hosts []HostInfo, partitioner string, err error) {
 	// we need conn to be the same because we need to query system.peers and system.local
 	// on the same node to get the whole cluster
-	conn := r.session.Pool.Pick(nil)
+	conn := r.session.pool.Pick(nil)
 	if conn == nil {
 		return r.prevHosts, r.prevPartitioner, nil
 	}
@@ -106,12 +102,11 @@ func (h *ringDescriber) run(sleep time.Duration) {
 			hosts, partitioner, err := h.GetHosts()
 			if err != nil {
 				log.Println("RingDescriber: unable to get ring topology:", err)
-			} else {
-				h.session.Pool.SetHosts(hosts)
-				if v, ok := h.session.Pool.(SetPartitioner); ok {
-					v.SetPartitioner(partitioner)
-				}
+				continue
 			}
+
+			h.session.pool.SetHosts(hosts)
+			h.session.pool.SetPartitioner(partitioner)
 		case <-h.closeChan:
 			return
 		}

policies.go

@@ -10,8 +10,8 @@ import (
 	"sync/atomic"
 )
 
-//RetryableQuery is an interface that represents a query or batch statement that
-//exposes the correct functions for the retry policy logic to evaluate correctly.
+// RetryableQuery is an interface that represents a query or batch statement that
+// exposes the correct functions for the retry policy logic to evaluate correctly.
 type RetryableQuery interface {
 	Attempts() int
 	GetConsistency() Consistency
@@ -48,8 +48,8 @@ func (s *SimpleRetryPolicy) Attempt(q RetryableQuery) bool {
 	return q.Attempts() <= s.NumRetries
 }
 
-//HostSelectionPolicy is an interface for selecting
-//the most appropriate host to execute a given query.
+// HostSelectionPolicy is an interface for selecting
+// the most appropriate host to execute a given query.
 type HostSelectionPolicy interface {
 	SetHosts
 	SetPartitioner
@@ -57,11 +57,12 @@ type HostSelectionPolicy interface {
 	Pick(*Query) NextHost
 }
 
-//NextHost is an iteration function over picked hosts
+// NextHost is an iteration function over picked hosts
 type NextHost func() *HostInfo
 
-//NewRoundRobinHostPolicy is a round-robin load balancing policy
-func NewRoundRobinHostPolicy() HostSelectionPolicy {
+// RoundRobinHostPolicy is a round-robin load balancing policy, where each host
+// is tried sequentially for each query.
+func RoundRobinHostPolicy() HostSelectionPolicy {
 	return &roundRobinHostPolicy{hosts: []HostInfo{}}
 }
 
@@ -105,8 +106,10 @@ func (r *roundRobinHostPolicy) Pick(qry *Query) NextHost {
 	}
 }
 
-//NewTokenAwareHostPolicy is a token aware host selection policy
-func NewTokenAwareHostPolicy(fallback HostSelectionPolicy) HostSelectionPolicy {
+// TokenAwareHostPolicy is a token aware host selection policy, where hosts are
+// selected based on the partition key, so queries are sent to the host which
+// owns the partition. Fallback is used when routing information is not available.
+func TokenAwareHostPolicy(fallback HostSelectionPolicy) HostSelectionPolicy {
 	return &tokenAwareHostPolicy{fallback: fallback, hosts: []HostInfo{}}
 }
 
@@ -227,8 +230,10 @@ type roundRobinConnPolicy struct {
 	mu    sync.RWMutex
 }
 
-func NewRoundRobinConnPolicy() ConnSelectionPolicy {
-	return &roundRobinConnPolicy{}
+func RoundRobinConnPolicy() func() ConnSelectionPolicy {
+	return func() ConnSelectionPolicy {
+		return &roundRobinConnPolicy{}
+	}
 }
 
 func (r *roundRobinConnPolicy) SetConns(conns []*Conn) {

policies_test.go

@@ -8,7 +8,7 @@ import "testing"
 
 // Tests of the round-robin host selection policy implementation
 func TestRoundRobinHostPolicy(t *testing.T) {
-	policy := NewRoundRobinHostPolicy()
+	policy := RoundRobinHostPolicy()
 
 	hosts := []HostInfo{
 		HostInfo{HostId: "0"},
@@ -46,7 +46,7 @@ func TestRoundRobinHostPolicy(t *testing.T) {
 // Tests of the token-aware host selection policy implementation with a
 // round-robin host selection policy fallback.
 func TestTokenAwareHostPolicy(t *testing.T) {
-	policy := NewTokenAwareHostPolicy(NewRoundRobinHostPolicy())
+	policy := TokenAwareHostPolicy(RoundRobinHostPolicy())
 
 	query := &Query{}
 
@@ -101,7 +101,7 @@ func TestTokenAwareHostPolicy(t *testing.T) {
 
 // Tests of the round-robin connection selection policy implementation
 func TestRoundRobinConnPolicy(t *testing.T) {
-	policy := NewRoundRobinConnPolicy()
+	policy := RoundRobinConnPolicy()()
 
 	conn0 := &Conn{}
 	conn1 := &Conn{}

session.go

@@ -28,7 +28,7 @@ import (
 // and automatically sets a default consinstency level on all operations
 // that do not have a consistency level set.
 type Session struct {
-	Pool                ConnectionPool
+	pool                *policyConnPool
 	cons                Consistency
 	pageSize            int
 	prefetch            float64
@@ -60,47 +60,44 @@ func NewSession(cfg ClusterConfig) (*Session, error) {
 		cfg.NumStreams = maxStreams
 	}
 
-	pool, err := cfg.ConnPoolType(&cfg)
-	if err != nil {
-		return nil, err
-	}
-
 	//Adjust the size of the prepared statements cache to match the latest configuration
 	stmtsLRU.Lock()
 	initStmtsLRU(cfg.MaxPreparedStmts)
 	stmtsLRU.Unlock()
 
 	s := &Session{
-		Pool:     pool,
 		cons:     cfg.Consistency,
 		prefetch: 0.25,
 		cfg:      cfg,
+		pageSize: cfg.PageSize,
+	}
+
+	pool, err := cfg.PoolConfig.buildPool(&s.cfg)
+	if err != nil {
+		return nil, err
 	}
+	s.pool = pool
 
 	//See if there are any connections in the pool
-	if pool.Size() > 0 {
-		s.routingKeyInfoCache.lru = lru.New(cfg.MaxRoutingKeyInfo)
-
-		s.SetConsistency(cfg.Consistency)
-		s.SetPageSize(cfg.PageSize)
-
-		if cfg.DiscoverHosts {
-			s.hostSource = &ringDescriber{
-				session:    s,
-				dcFilter:   cfg.Discovery.DcFilter,
-				rackFilter: cfg.Discovery.RackFilter,
-				closeChan:  make(chan bool),
-			}
+	if pool.Size() == 0 {
+		s.Close()
+		return nil, ErrNoConnectionsStarted
+	}
+
+	s.routingKeyInfoCache.lru = lru.New(cfg.MaxRoutingKeyInfo)
 
-			go s.hostSource.run(cfg.Discovery.Sleep)
+	if cfg.DiscoverHosts {
+		s.hostSource = &ringDescriber{
+			session:    s,
+			dcFilter:   cfg.Discovery.DcFilter,
+			rackFilter: cfg.Discovery.RackFilter,
+			closeChan:  make(chan bool),
 		}
 
-		return s, nil
+		go s.hostSource.run(cfg.Discovery.Sleep)
 	}
 
-	s.Close()
-
-	return nil, ErrNoConnectionsStarted
+	return s, nil
 }
 
 // SetConsistency sets the default consistency level for this session. This
@@ -185,7 +182,7 @@ func (s *Session) Close() {
 	}
 	s.isClosed = true
 
-	s.Pool.Close()
+	s.pool.Close()
 
 	if s.hostSource != nil {
 		close(s.hostSource.closeChan)
@@ -210,7 +207,7 @@ func (s *Session) executeQuery(qry *Query) *Iter {
 	qry.attempts = 0
 	qry.totalLatency = 0
 	for {
-		conn := s.Pool.Pick(qry)
+		conn := s.pool.Pick(qry)
 
 		//Assign the error unavailable to the iterator
 		if conn == nil {
@@ -294,7 +291,7 @@ func (s *Session) routingKeyInfo(stmt string) (*routingKeyInfo, error) {
 	)
 
 	// get the query info for the statement
-	conn := s.Pool.Pick(nil)
+	conn := s.pool.Pick(nil)
 	if conn == nil {
 		// no connections
 		inflight.err = ErrNoConnections
@@ -389,7 +386,7 @@ func (s *Session) executeBatch(batch *Batch) (*Iter, error) {
 	batch.attempts = 0
 	batch.totalLatency = 0
 	for {
-		conn := s.Pool.Pick(nil)
+		conn := s.pool.Pick(nil)
 
 		//Assign the error unavailable and break loop
 		if conn == nil {

session_test.go

@@ -10,13 +10,13 @@ import (
 func TestSessionAPI(t *testing.T) {
 
 	cfg := &ClusterConfig{}
-	pool, err := NewSimplePool(cfg)
+	pool, err := cfg.PoolConfig.buildPool(cfg)
 	if err != nil {
 		t.Fatal(err)
 	}
 
 	s := &Session{
-		Pool: pool,
+		pool: pool,
 		cfg:  *cfg,
 		cons: Quorum,
 	}
@@ -160,13 +160,13 @@ func TestQueryShouldPrepare(t *testing.T) {
 func TestBatchBasicAPI(t *testing.T) {
 
 	cfg := &ClusterConfig{RetryPolicy: &SimpleRetryPolicy{NumRetries: 2}}
-	pool, err := NewSimplePool(cfg)
+	pool, err := cfg.PoolConfig.buildPool(cfg)
 	if err != nil {
 		t.Fatal(err)
 	}
 
 	s := &Session{
-		Pool: pool,
+		pool: pool,
 		cfg:  *cfg,
 		cons: Quorum,
 	}

stress_test.go

@@ -36,7 +36,6 @@ func BenchmarkConnStress(b *testing.B) {
 
 	b.SetParallelism(workers)
 	b.RunParallel(writer)
-
 }
 
 func BenchmarkConnRoutingKey(b *testing.B) {
@@ -45,7 +44,7 @@ func BenchmarkConnRoutingKey(b *testing.B) {
 	cluster := createCluster()
 	cluster.NumConns = 1
 	cluster.NumStreams = workers
-	cluster.ConnPoolType = NewTokenAwareConnPool
+	cluster.PoolConfig.HostSelectionPolicy = TokenAwareHostPolicy(RoundRobinHostPolicy())
 	session := createSessionFromCluster(cluster, b)
 	defer session.Close()