transport: return errors on missing connections

Author: Kulezi

query.go

@@ -31,7 +31,12 @@ func (q *Query) Exec(ctx context.Context) (Result, error) {
 	for n != nil {
 	sameNodeRetries:
 		for {
-			conn := n.Conn(info)
+			conn, err := n.Conn(info)
+			if err != nil {
+				lastErr = err
+				break sameNodeRetries
+			}
+
 			res, err := q.exec(ctx, conn, q.stmt, nil)
 			if err != nil {
 				ri := transport.RetryInfo{
@@ -60,15 +65,19 @@ func (q *Query) Exec(ctx context.Context) (Result, error) {
 		i++
 		n = q.session.cfg.HostSelectionPolicy.Node(info, i)
 	}
+
+	if lastErr == nil {
+		return Result{}, ErrNoConnection
+	}
 	return Result{}, lastErr
 }
 
 func (q *Query) pickConn(qi transport.QueryInfo) (*transport.Conn, error) {
 	n := q.session.cfg.HostSelectionPolicy.Node(qi, 0)
 
-	conn := n.Conn(qi)
-	if conn == nil {
-		return nil, errNoConnection
+	conn, err := n.Conn(qi)
+	if err != nil {
+		return nil, ErrNoConnection
 	}
 
 	return conn, nil
@@ -281,6 +290,7 @@ type iterWorker struct {
 	pickNode  func(transport.QueryInfo, int) *transport.Node
 	nodeIdx   int
 	conn      *transport.Conn
+	connErr   error
 
 	rd transport.RetryDecider
 
@@ -295,7 +305,7 @@ func (w *iterWorker) loop(ctx context.Context) {
 		w.errCh <- fmt.Errorf("can't pick a node to execute request")
 		return
 	}
-	w.conn = n.Conn(w.queryInfo)
+	w.conn, w.connErr = n.Conn(w.queryInfo)
 
 	for {
 		_, ok := <-w.requestCh
@@ -324,6 +334,10 @@ func (w *iterWorker) exec(ctx context.Context) (transport.QueryResult, error) {
 	for {
 	sameNodeRetries:
 		for {
+			if w.connErr != nil {
+				lastErr = w.connErr
+				break
+			}
 			res, err := w.queryExec(ctx, w.conn, w.stmt, w.pagingState)
 			if err != nil {
 				ri := transport.RetryInfo{
@@ -349,9 +363,12 @@ func (w *iterWorker) exec(ctx context.Context) (transport.QueryResult, error) {
 		w.nodeIdx++
 		n := w.pickNode(w.queryInfo, w.nodeIdx)
 		if n == nil {
+			if lastErr == nil {
+				return transport.QueryResult{}, ErrNoConnection
+			}
 			return transport.QueryResult{}, lastErr
 		}
 
-		w.conn = n.Conn(w.queryInfo)
+		w.conn, w.connErr = n.Conn(w.queryInfo)
 	}
 }

session.go

@@ -56,7 +56,7 @@ var (
 		"SERIAL      Consistency = 0x0008\n" +
 		"LOCALSERIAL Consistency = 0x0009\n" +
 		"LOCALONE    Consistency = 0x000A")
-	errNoConnection = fmt.Errorf("no working connection")
+	ErrNoConnection = fmt.Errorf("no connection to execute the query on")
 )
 
 type Compression = frame.Compression

transport/node.go

@@ -53,10 +53,10 @@ func (n *Node) Close() {
 	n.setStatus(statusDown)
 }
 
-func (n *Node) LeastBusyConn() *Conn {
+func (n *Node) LeastBusyConn() (*Conn, error) {
 	return n.pool.LeastBusyConn()
 }
-func (n *Node) Conn(qi QueryInfo) *Conn {
+func (n *Node) Conn(qi QueryInfo) (*Conn, error) {
 	if qi.tokenAware {
 		return n.pool.Conn(qi.token)
 	}
@@ -65,7 +65,11 @@ func (n *Node) Conn(qi QueryInfo) *Conn {
 }
 
 func (n *Node) Prepare(ctx context.Context, s Statement) (Statement, error) {
-	return n.LeastBusyConn().Prepare(ctx, s)
+	conn, err := n.LeastBusyConn()
+	if err != nil {
+		return Statement{}, err
+	}
+	return conn.Prepare(ctx, s)
 }
 
 var versionQuery = Statement{
@@ -74,7 +78,11 @@ var versionQuery = Statement{
 }
 
 func (n *Node) FetchSchemaVersion(ctx context.Context) (frame.UUID, error) {
-	conn := n.LeastBusyConn()
+	conn, err := n.LeastBusyConn()
+	if err != nil {
+		return frame.UUID{}, err
+	}
+
 	res, err := conn.Query(ctx, versionQuery, nil)
 	if err != nil {
 		return frame.UUID{}, err

transport/pool.go

@@ -41,13 +41,13 @@ func (p *ConnPool) String() string {
 	return fmt.Sprintf("pool %s [shards=%d]", p.host, p.nrShards)
 }
 
-func (p *ConnPool) Conn(token Token) *Conn {
+func (p *ConnPool) Conn(token Token) (*Conn, error) {
 	idx := p.shardOf(token)
 	if conn := p.loadConn(idx); conn != nil {
 		if isHeavyLoaded(conn) {
-			return p.maybeReplaceWithLessBusyConn(conn)
+			return p.maybeReplaceWithLessBusyConn(conn), nil
 		}
-		return conn
+		return conn, nil
 	}
 	return p.LeastBusyConn()
 }
@@ -57,7 +57,7 @@ func isHeavyLoaded(conn *Conn) bool {
 }
 
 func (p *ConnPool) maybeReplaceWithLessBusyConn(conn *Conn) *Conn {
-	if lb := p.LeastBusyConn(); conn.Waiting()-lb.Waiting() > maxStreamID<<1/10 {
+	if lb, err := p.LeastBusyConn(); err == nil && conn.Waiting()-lb.Waiting() > maxStreamID<<1/10 {
 		if p.connObs != nil {
 			p.connObs.OnPickReplacedWithLessBusyConn(conn.Event())
 		}
@@ -66,7 +66,7 @@ func (p *ConnPool) maybeReplaceWithLessBusyConn(conn *Conn) *Conn {
 	return conn
 }
 
-func (p *ConnPool) LeastBusyConn() *Conn {
+func (p *ConnPool) LeastBusyConn() (*Conn, error) {
 	var (
 		leastBusyConn *Conn
 		minBusy       = maxStreamID + 2 // adding 2 more is required due to atomics
@@ -80,7 +80,11 @@ func (p *ConnPool) LeastBusyConn() *Conn {
 			}
 		}
 	}
-	return leastBusyConn
+
+	if leastBusyConn == nil {
+		return nil, fmt.Errorf("no connections available for host %s", p.host)
+	}
+	return leastBusyConn, nil
 }
 
 func (p *ConnPool) shardOf(token Token) int {

transport/pool_integration_test.go

@@ -72,21 +72,27 @@ func TestConnPoolConnIntegration(t *testing.T) {
 	defer p.Close()
 
 	t0 := MurmurToken([]byte(""))
-	if conn := p.Conn(t0); conn == nil || conn.Shard() != 0 {
+	if conn, err := p.Conn(t0); err != nil || conn.Shard() != 0 {
 		t.Fatal("invalid return of Conn")
 	}
 
 	load := uint32(math.Floor(maxStreamID/2 + 1))
-	p.Conn(t0).stats.inQueue.Store(load)
 
-	if conn := p.Conn(t0); conn == nil {
+	conn, err := p.Conn(t0)
+	if err != nil {
+		t.Fatal("invalid return of Conn")
+	}
+
+	conn.stats.inQueue.Store(load)
+
+	if conn, err := p.Conn(t0); err != nil {
 		t.Fatal("invalid return of Conn")
 	} else if conn.Shard() == 0 {
 		t.Fatalf("invalid load distribution")
 	}
 
 	t1 := MurmurToken([]byte("0")) // Very big number approx. 3 * 10^18.
-	if conn := p.Conn(t1); conn == nil {
+	if _, err := p.Conn(t1); err != nil {
 		t.Fatal("invalid return of Conn")
 	}
 }