Endless query execution fix

Fix for the endless query execution when HostSelectionPolicy returns
the same downed host.
Internal HostSelectionPolicy will return the host only once if HostID is set.
Documentation for external policies was added.

patch by Oleksandr Luzhniy; reviewed by João Reis, for CASSGO-50

Author: tengu-alt

CHANGELOG.md

@@ -63,6 +63,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Don't panic in MapExecuteBatchCAS if no `[applied]` column is returned (CASSGO-42)
 
+- Endless query execution fix (CASSGO-50)
+
 ## [1.7.0] - 2024-09-23
 
 This release is the first after the donation of gocql to the Apache Software Foundation (ASF)

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 {