Endless query execution fix

Author: tengu-alt

policies.go

@@ -323,6 +323,7 @@ func (host *selectedHost) Info() *HostInfo {
 func (host *selectedHost) Mark(err error) {}
 
 // NextHost is an iteration function over picked hosts
+// Should return nil if SelectedHost is not Up to prevent endless query execution.
 type NextHost func() SelectedHost
 
 // RoundRobinHostPolicy is a round-robin load balancing policy, where each host

query_executor.go

@@ -27,6 +27,7 @@ package gocql
 import (
 	"context"
 	"sync"
+	"sync/atomic"
 	"time"
 )
 
@@ -89,14 +90,17 @@ func (q *queryExecutor) executeQuery(qry ExecutableQuery) (*Iter, error) {
 	// check if the host id is specified for the query,
 	// if it is, the query should be executed at the corresponding host.
 	if hostID := qry.GetHostID(); hostID != "" {
+		pool, ok := q.pool.getPoolByHostID(hostID)
+		if !ok {
+			return nil, ErrNoConnections
+		}
+		host := pool.host
+		var returnedHostOnce int32 = 0
 		hostIter = func() SelectedHost {
-			pool, ok := q.pool.getPoolByHostID(hostID)
-			// if the specified host is down
-			// we return nil to avoid endless query execution in queryExecutor.do()
-			if !ok || !pool.host.IsUp() {
-				return nil
+			if atomic.CompareAndSwapInt32(&returnedHostOnce, 0, 1) {
+				return (*selectedHost)(host)
 			}
-			return (*selectedHost)(pool.host)
+			return nil
 		}
 	}