From 598a4aede4a0fa465eb549c712ad7d02120a4153 Mon Sep 17 00:00:00 2001
From: Dmitry Kropachev <dmitry.kropachev@gmail.com>
Date: Fri, 21 Jun 2024 20:24:23 -0400
Subject: [PATCH] Simplify tokenAwareHostPolicy.Pick

---
 policies.go | 38 ++++++++++++--------------------------
 1 file changed, 12 insertions(+), 26 deletions(-)

diff --git a/policies.go b/policies.go
index 73cbbc4ad..82a0ea923 100644
--- a/policies.go
+++ b/policies.go
@@ -651,15 +651,12 @@ func (t *tokenAwareHostPolicy) Pick(qry ExecutableQuery) NextHost {
 	var replicas []*HostInfo
 
 	if qry.GetSession() != nil && qry.GetSession().tabletsRoutingV1 {
-		t.tablets.mu.Lock()
 		tablets := t.tablets.get()
 
 		// Search for tablets with Keyspace and Table from the Query
 		l, r := findTablets(tablets, qry.Keyspace(), qry.Table())
 		if l != -1 {
 			tablet := findTabletForToken(tablets, token, l, r)
-
-			replicas = []*HostInfo{}
 			hosts := t.hosts.get()
 			for _, replica := range tablet.Replicas() {
 				for _, host := range hosts {
@@ -669,36 +666,25 @@ func (t *tokenAwareHostPolicy) Pick(qry ExecutableQuery) NextHost {
 					}
 				}
 			}
-		} else {
-			ht := meta.replicas[qry.Keyspace()].replicasFor(token)
-
-			if ht == nil {
-				host, _ := meta.tokenRing.GetHostForToken(token)
-				replicas = []*HostInfo{host}
-			} else {
-				replicas = ht.hosts
-			}
-		}
-
-		if t.shuffleReplicas && !qry.IsLWT() {
-			replicas = shuffleHosts(replicas)
 		}
+	}
 
-		t.tablets.mu.Unlock()
-	} else {
+	if len(replicas) == 0 {
 		ht := meta.replicas[qry.Keyspace()].replicasFor(token)
-
-		if ht == nil {
-			host, _ := meta.tokenRing.GetHostForToken(token)
-			replicas = []*HostInfo{host}
-		} else {
+		if ht != nil {
 			replicas = ht.hosts
-			if t.shuffleReplicas && !qry.IsLWT() {
-				replicas = shuffleHosts(replicas)
-			}
 		}
 	}
 
+	if len(replicas) == 0 {
+		host, _ := meta.tokenRing.GetHostForToken(token)
+		replicas = []*HostInfo{host}
+	}
+
+	if t.shuffleReplicas && !qry.IsLWT() && len(replicas) > 1 {
+		replicas = shuffleHosts(replicas)
+	}
+
 	if s := qry.GetSession(); s != nil && t.avoidSlowReplicas {
 		healthyReplicas := make([]*HostInfo, 0, len(replicas))
 		unhealthyReplicas := make([]*HostInfo, 0, len(replicas))