Endless query execution fix proposal

Author: tengu-alt

cassandra_test.go

@@ -44,7 +44,7 @@ import (
 	"time"
 	"unicode"
 
-	inf "gopkg.in/inf.v0"
+	"gopkg.in/inf.v0"
 
 	"github.com/stretchr/testify/require"
 )
@@ -3955,3 +3955,71 @@ func TestRoutingKeyCacheUsesOverriddenKeyspace(t *testing.T) {
 
 	session.Query("DROP KEYSPACE IF EXISTS gocql_test_routing_key_cache").Exec()
 }
+
+func DownedHostSelectionPolicy() HostSelectionPolicy {
+	return &downedHostSelectionPolicy{}
+	//return &roundRobinHostPolicy{}
+}
+
+type downedHostSelectionPolicy struct {
+	hosts           cowHostList
+	lastUsedHostIdx uint64
+}
+
+func (r *downedHostSelectionPolicy) IsLocal(*HostInfo) bool              { return true }
+func (r *downedHostSelectionPolicy) KeyspaceChanged(KeyspaceUpdateEvent) {}
+func (r *downedHostSelectionPolicy) SetPartitioner(partitioner string)   {}
+func (r *downedHostSelectionPolicy) Init(*Session)                       {}
+
+func (r *downedHostSelectionPolicy) Pick(qry ExecutableQuery) NextHost {
+	return dummyPolicy()
+}
+
+func dummyPolicy() NextHost {
+	return func() SelectedHost {
+		h := &HostInfo{
+			state:          NodeDown,
+			connectAddress: net.ParseIP("172.30.0.2"),
+		}
+		return (*selectedHost)(h)
+
+	}
+}
+
+func (r *downedHostSelectionPolicy) AddHost(host *HostInfo) {
+	r.hosts.add(host)
+}
+
+func (r *downedHostSelectionPolicy) RemoveHost(host *HostInfo) {
+	r.hosts.remove(host.ConnectAddress())
+}
+
+func (r *downedHostSelectionPolicy) HostUp(host *HostInfo) {
+	r.AddHost(host)
+}
+
+func (r *downedHostSelectionPolicy) HostDown(host *HostInfo) {
+	r.RemoveHost(host)
+}
+
+func TestQuery_Do(t *testing.T) {
+	cluster := createCluster()
+	session := createSessionFromCluster(cluster, t)
+	defer session.Close()
+
+	if err := createTable(session, `CREATE TABLE gocql_test.policy (id int primary key)`); err != nil {
+		t.Fatal(err)
+	}
+
+	session.executor.policy = DownedHostSelectionPolicy()
+	err := session.Query("select * from gocql_test.policy").Exec()
+	if err != ErrNoConnections {
+		t.Fatalf("expected ErrNoConnections, got %v", err)
+	}
+
+	sp := &SimpleSpeculativeExecution{NumAttempts: 3, TimeoutDelay: 200 * time.Millisecond}
+	err = session.Query("select * from gocql_test.policy").SetSpeculativeExecutionPolicy(sp).Idempotent(true).Exec()
+	if err == nil {
+		t.Fatalf("expected ErrNoConnections, got %v", err)
+	}
+}

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

@@ -150,12 +150,17 @@ func (q *queryExecutor) executeQuery(qry ExecutableQuery) (*Iter, error) {
 func (q *queryExecutor) do(ctx context.Context, qry ExecutableQuery, hostIter NextHost) *Iter {
 	selectedHost := hostIter()
 	rt := qry.retryPolicy()
+	count := 0
+	q.pool.mu.Lock()
+	threshold := len(q.pool.hostConnPools)
+	q.pool.mu.Unlock()
 
 	var lastErr error
 	var iter *Iter
-	for selectedHost != nil {
+	for selectedHost != nil && count < threshold {
 		host := selectedHost.Info()
 		if host == nil || !host.IsUp() {
+			count++
 			selectedHost = hostIter()
 			continue
 		}