refactor: update ConnPicker interface to return error from Put method

Co-authored-by: Dmitry Kropachev <[email protected]>

Author: nikagra

connectionpool.go

@@ -544,7 +544,9 @@ func (pool *hostConnPool) connect() (err error) {
 
 	// lazily initialize the connPicker when we know the required type
 	pool.initConnPicker(conn)
-	pool.connPicker.Put(conn)
+	if err := pool.connPicker.Put(conn); err != nil {
+		return err
+	}
 	conn.finalizeConnection()
 
 	return nil

connpicker.go

@@ -8,7 +8,7 @@ import (
 
 type ConnPicker interface {
 	Pick(Token, ExecutableQuery) *Conn
-	Put(*Conn)
+	Put(*Conn) error
 	Remove(conn *Conn)
 	InFlight() int
 	Size() (int, int)
@@ -96,10 +96,11 @@ func (p *defaultConnPicker) Pick(Token, ExecutableQuery) *Conn {
 	return leastBusyConn
 }
 
-func (p *defaultConnPicker) Put(conn *Conn) {
+func (p *defaultConnPicker) Put(conn *Conn) error {
 	p.mu.Lock()
 	p.conns = append(p.conns, conn)
 	p.mu.Unlock()
+	return nil
 }
 
 func (*defaultConnPicker) NextShard() (shardID, nrShards int) {
@@ -115,7 +116,8 @@ func (nopConnPicker) Pick(Token, ExecutableQuery) *Conn {
 	return nil
 }
 
-func (nopConnPicker) Put(*Conn) {
+func (nopConnPicker) Put(*Conn) error {
+	return nil
 }
 
 func (nopConnPicker) Remove(conn *Conn) {

integration_only.go

@@ -45,7 +45,7 @@ func (s *Session) MissingConnections() (int, error) {
 
 type ConnPickerIntegration interface {
 	Pick(Token, ExecutableQuery) *Conn
-	Put(*Conn)
+	Put(*Conn) error
 	Remove(conn *Conn)
 	InFlight() int
 	Size() (int, int)

scylla.go

@@ -452,19 +452,19 @@ func (p *scyllaConnPicker) shardOf(token int64Token) int {
 	return int(sum >> 32)
 }
 
-func (p *scyllaConnPicker) Put(conn *Conn) {
+func (p *scyllaConnPicker) Put(conn *Conn) error {
 	var (
 		nrShards = conn.scyllaSupported.nrShards
 		shard    = conn.scyllaSupported.shard
 	)
 
 	if nrShards == 0 {
-		panic(fmt.Sprintf("scylla: %s not a sharded connection", p.address))
+		return errors.New("server reported that it has no shards")
 	}
 
 	if nrShards != len(p.conns) {
 		if nrShards != p.nrShards {
-			panic(fmt.Sprintf("scylla: %s invalid number of shards", p.address))
+			return fmt.Errorf("server %s reported that number of shard changed from %d to %d", p.address, p.nrShards, nrShards)
 		}
 		conns := p.conns
 		p.conns = make([]*Conn, nrShards, nrShards)
@@ -509,6 +509,8 @@ func (p *scyllaConnPicker) Put(conn *Conn) {
 	if p.shouldCloseExcessConns() {
 		p.closeExcessConns()
 	}
+
+	return nil
 }
 
 func (p *scyllaConnPicker) shouldCloseExcessConns() bool {