Endless query execution fix

Author: tengu-alt

events_ccm_test.go

@@ -104,7 +104,7 @@ func TestEventNodeDownControl(t *testing.T) {
 	}
 	session.pool.mu.RUnlock()
 
-	host := session.ring.getHost(node.Addr)
+	host, _ := session.ring.getHost(node.Addr)
 	if host == nil {
 		t.Fatal("node not in metadata ring")
 	} else if host.IsUp() {
@@ -146,7 +146,7 @@ func TestEventNodeDown(t *testing.T) {
 		t.Fatal("node not removed after remove event")
 	}
 
-	host := session.ring.getHost(node.Addr)
+	host, _ := session.ring.getHost(node.Addr)
 	if host == nil {
 		t.Fatal("node not in metadata ring")
 	} else if host.IsUp() {
@@ -203,7 +203,7 @@ func TestEventNodeUp(t *testing.T) {
 		t.Fatal("node not added after node added event")
 	}
 
-	host := session.ring.getHost(node.Addr)
+	host, _ := session.ring.getHost(node.Addr)
 	if host == nil {
 		t.Fatal("node not in metadata ring")
 	} else if !host.IsUp() {

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 eventually 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,16 @@ 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 != "" {
+		host, ok := q.pool.session.ring.getHost(hostID)
+		if !ok {
+			return nil, ErrNoConnections
+		}
+		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
 		}
 	}
 

ring.go

@@ -67,11 +67,11 @@ func (r *ring) getHostByIP(ip string) (*HostInfo, bool) {
 	return r.hosts[hi], ok
 }
 
-func (r *ring) getHost(hostID string) *HostInfo {
+func (r *ring) getHost(hostID string) (host *HostInfo, ok bool) {
 	r.mu.RLock()
-	host := r.hosts[hostID]
+	host, ok = r.hosts[hostID]
 	r.mu.RUnlock()
-	return host
+	return
 }
 
 func (r *ring) allHosts() []*HostInfo {