session: expose API to get host pool information

Author: dkropachev

connectionpool.go

@@ -203,6 +203,16 @@ func (p *policyConnPool) getPoolByHostID(hostID string) (pool *hostConnPool, ok
 	return
 }
 
+func (p *policyConnPool) iteratePool(iter func(info HostPoolInfo) bool) {
+	p.mu.RLock()
+	for _, pool := range p.hostConnPools {
+		if !iter(pool) {
+			break
+		}
+	}
+	p.mu.RUnlock()
+}
+
 func (p *policyConnPool) Close() {
 	p.mu.Lock()
 	defer p.mu.Unlock()
@@ -264,12 +274,12 @@ type hostConnPool struct {
 	filling bool
 }
 
-func (h *hostConnPool) String() string {
-	h.mu.RLock()
-	defer h.mu.RUnlock()
-	size, _ := h.connPicker.Size()
+func (pool *hostConnPool) String() string {
+	pool.mu.RLock()
+	defer pool.mu.RUnlock()
+	size, _ := pool.connPicker.Size()
 	return fmt.Sprintf("[filling=%v closed=%v conns=%v size=%v host=%v]",
-		h.filling, h.closed, size, h.size, h.host)
+		pool.filling, pool.closed, size, pool.size, pool.host)
 }
 
 func newHostConnPool(session *Session, host *HostInfo, port, size int,
@@ -587,3 +597,31 @@ func (pool *hostConnPool) HandleError(conn *Conn, err error, closed bool) {
 	pool.connPicker.Remove(conn)
 	go pool.fill_debounce()
 }
+
+func (pool *hostConnPool) GetConnectionCount() int {
+	pool.mu.Lock()
+	defer pool.mu.Unlock()
+	return pool.connPicker.GetConnectionCount()
+}
+
+func (pool *hostConnPool) GetExcessConnectionCount() int {
+	pool.mu.Lock()
+	defer pool.mu.Unlock()
+	return pool.connPicker.GetExcessConnectionCount()
+}
+
+func (pool *hostConnPool) GetShardCount() int {
+	pool.mu.Lock()
+	defer pool.mu.Unlock()
+	return pool.connPicker.GetShardCount()
+}
+
+func (pool *hostConnPool) Host() HostInformation {
+	return pool.host
+}
+
+func (pool *hostConnPool) IsClosed() bool {
+	pool.mu.Lock()
+	defer pool.mu.Unlock()
+	return pool.closed
+}

connpicker.go

@@ -18,6 +18,10 @@ type ConnPicker interface {
 	// nrShard specifies how many shards the host has.
 	// If nrShards is zero, the caller shouldn't use shard-aware port.
 	NextShard() (shardID, nrShards int)
+
+	GetConnectionCount() int
+	GetExcessConnectionCount() int
+	GetShardCount() int
 }
 
 type defaultConnPicker struct {
@@ -27,6 +31,21 @@ type defaultConnPicker struct {
 	mu    sync.RWMutex
 }
 
+func (p *defaultConnPicker) GetConnectionCount() int {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	return len(p.conns)
+}
+
+func (p *defaultConnPicker) GetExcessConnectionCount() int {
+	return 0
+}
+
+func (p *defaultConnPicker) GetShardCount() int {
+	// It is not supposed to be used for scylla nodes and therefore does not know anything about shards count
+	return 0
+}
+
 func newDefaultConnPicker(size int) *defaultConnPicker {
 	if size <= 0 {
 		panic(fmt.Sprintf("invalid pool size %d", size))
@@ -111,6 +130,18 @@ func (*defaultConnPicker) NextShard() (shardID, nrShards int) {
 // to the point where we have first connection.
 type nopConnPicker struct{}
 
+func (p nopConnPicker) GetConnectionCount() int {
+	return 0
+}
+
+func (p nopConnPicker) GetExcessConnectionCount() int {
+	return 0
+}
+
+func (p nopConnPicker) GetShardCount() int {
+	return 0
+}
+
 func (nopConnPicker) Pick(Token, ExecutableQuery) *Conn {
 	return nil
 }

host_source.go

@@ -176,6 +176,7 @@ type HostInfo struct {
 	state                   nodeState
 	scyllaShardAwarePort    uint16
 	scyllaShardAwarePortTLS uint16
+	scyllaShardCount        int
 	graph                   bool
 }
 
@@ -488,8 +489,10 @@ func (h *HostInfo) String() string {
 func (h *HostInfo) setScyllaSupported(s scyllaSupported) {
 	h.mu.Lock()
 	defer h.mu.Unlock()
+	h.partitioner = s.partitioner
 	h.scyllaShardAwarePort = s.shardAwarePort
 	h.scyllaShardAwarePortTLS = s.shardAwarePortSSL
+	h.scyllaShardCount = s.nrShards
 }
 
 // ScyllaShardAwarePort returns the shard aware port of this host.
@@ -508,6 +511,13 @@ func (h *HostInfo) ScyllaShardAwarePortTLS() uint16 {
 	return h.scyllaShardAwarePortTLS
 }
 
+// ScyllaShardCount returns count of shards on the node.
+func (h *HostInfo) ScyllaShardCount() int {
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+	return h.scyllaShardCount
+}
+
 // Returns true if we are using system_schema.keyspaces instead of system.schema_keyspaces
 func checkSystemSchema(control controlConnection) (bool, error) {
 	iter := control.querySystem("SELECT * FROM system_schema.keyspaces")

integration_only.go

@@ -8,17 +8,6 @@ package gocql
 
 import "fmt"
 
-func (pool *hostConnPool) MissingConnections() (int, error) {
-	pool.mu.Lock()
-	defer pool.mu.Unlock()
-
-	if pool.closed {
-		return 0, fmt.Errorf("pool is closed")
-	}
-	_, missing := pool.connPicker.Size()
-	return missing, nil
-}
-
 func (p *policyConnPool) MissingConnections() (int, error) {
 	p.mu.Lock()
 	defer p.mu.Unlock()
@@ -27,9 +16,9 @@ func (p *policyConnPool) MissingConnections() (int, error) {
 
 	// close the pools
 	for _, pool := range p.hostConnPools {
-		missing, err := pool.MissingConnections()
-		if err != nil {
-			return 0, err
+		missing := pool.GetShardCount() - pool.GetConnectionCount()
+		if pool.IsClosed() {
+			return 0, fmt.Errorf("pool for %s is closed", pool.host.HostID())
 		}
 		total += missing
 	}

scylla.go

@@ -518,6 +518,18 @@ func (p *scyllaConnPicker) shouldCloseExcessConns() bool {
 	return len(p.excessConns) > int(p.excessConnsLimitRate*float32(p.nrShards))
 }
 
+func (p *scyllaConnPicker) GetConnectionCount() int {
+	return p.nrConns
+}
+
+func (p *scyllaConnPicker) GetExcessConnectionCount() int {
+	return len(p.excessConns)
+}
+
+func (p *scyllaConnPicker) GetShardCount() int {
+	return p.nrShards
+}
+
 func (p *scyllaConnPicker) Remove(conn *Conn) {
 	shard := conn.scyllaSupported.shard
 

session.go

@@ -2425,6 +2425,47 @@ func (s *Session) GetHosts() []*HostInfo {
 	return s.hostSource.getHostsList()
 }
 
+type HostInformation interface {
+	Peer() net.IP
+	ConnectAddress() net.IP
+	UntranslatedConnectAddress() net.IP
+	BroadcastAddress() net.IP
+	ListenAddress() net.IP
+	RPCAddress() net.IP
+	PreferredIP() net.IP
+	DataCenter() string
+	Rack() string
+	HostID() string
+	WorkLoad() string
+	Partitioner() string
+	ClusterName() string
+	Tokens() []string
+	Port() int
+	IsUp() bool
+	ScyllaShardAwarePort() uint16
+	ScyllaShardAwarePortTLS() uint16
+	ScyllaShardCount() int
+}
+
+type HostPoolInfo interface {
+	GetConnectionCount() int
+	GetExcessConnectionCount() int
+	GetShardCount() int
+	String() string
+	InFlight() int
+	Host() HostInformation
+	IsClosed() bool
+}
+
+func (s *Session) GetHostPoolByID(hostID string) HostPoolInfo {
+	hostPool, _ := s.pool.getPoolByHostID(hostID)
+	return hostPool
+}
+
+func (s *Session) IterateHostPools(iter func(info HostPoolInfo) bool) {
+	s.pool.iteratePool(iter)
+}
+
 type ObservedQuery struct {
 	// Start is a time when the query was attempted
 	Start time.Time