From a39ad5597305a592330dc15d7352471b4e2195f5 Mon Sep 17 00:00:00 2001
From: sylwiaszunejko <sylwia.szunejko@gmail.com>
Date: Fri, 2 Feb 2024 13:08:30 +0100
Subject: [PATCH] Fix checking if tablets should be used in scyllaConnPicker
 Pick

Previously the check was always based on p.conns[0], but there
is a posibility that it is nil and it should be checked.
The check should be based on not nil connection
This commit introduces searching for first not nil connection
and checking tablet-awareness based on this connection.
---
 scylla.go | 35 ++++++++++++++++++++++-------------
 1 file changed, 22 insertions(+), 13 deletions(-)

diff --git a/scylla.go b/scylla.go
index 420b18a01..7dece242a 100644
--- a/scylla.go
+++ b/scylla.go
@@ -370,24 +370,33 @@ func (p *scyllaConnPicker) Pick(t token, keyspace string, table string) *Conn {
 
 	idx := -1
 
-	p.conns[0].mu.Lock()
-	if p.conns[0].tabletsRoutingV1 {
-		tablets := p.conns[0].session.getTablets()
-
-		// Search for tablets with Keyspace and Table from the Query
-		l, r := findTablets(tablets, keyspace, table)
-
-		if l != -1 {
-			tablet := findTabletForToken(tablets, mmt, l, r)
+	for _, conn := range p.conns {
+		if conn == nil {
+			continue
+		}
 
-			for _, replica := range tablet.replicas {
-				if replica.hostId.String() == p.hostId {
-					idx = replica.shardId
+		conn.mu.Lock()
+		if conn.tabletsRoutingV1 {
+			tablets := conn.session.getTablets()
+	
+			// Search for tablets with Keyspace and Table from the Query
+			l, r := findTablets(tablets, keyspace, table)
+	
+			if l != -1 {
+				tablet := findTabletForToken(tablets, mmt, l, r)
+	
+				for _, replica := range tablet.replicas {
+					if replica.hostId.String() == p.hostId {
+						idx = replica.shardId
+					}
 				}
 			}
 		}
+		conn.mu.Unlock()
+
+		break
 	}
-	p.conns[0].mu.Unlock()
+	
 	if idx == -1 {
 		idx = p.shardOf(mmt)
 	}