Fix JedisBroadcastException in functionLoadReplace for Redis Cluster

pom.xml

@@ -348,6 +348,7 @@
 						<include>src/test/java/redis/clients/jedis/commands/jedis/ClusterStreamsCommandsTest.java</include>
 						<include>src/test/java/redis/clients/jedis/commands/jedis/PooledStreamsCommandsTest.java</include>
 						<include>src/test/java/redis/clients/jedis/resps/StreamEntryDeletionResultTest.java</include>
+						<include>**/*FunctionCommandsTest*</include>
 					</includes>
 				</configuration>
 				<executions>

src/main/java/redis/clients/jedis/JedisClusterInfoCache.java

@@ -39,6 +39,7 @@ public class JedisClusterInfoCache {
   private static final Logger logger = LoggerFactory.getLogger(JedisClusterInfoCache.class);
 
   private final Map<String, ConnectionPool> nodes = new HashMap<>();
+  private final Map<String, ConnectionPool> primaryNodesCache = new HashMap<>();
   private final ConnectionPool[] slots = new ConnectionPool[Protocol.CLUSTER_HASHSLOTS];
   private final HostAndPort[] slotNodes = new HostAndPort[Protocol.CLUSTER_HASHSLOTS];
   private final List<ConnectionPool>[] replicaSlots;
@@ -176,6 +177,7 @@ public void discoverClusterNodesAndSlots(Connection jedis) {
           HostAndPort targetNode = generateHostAndPort(hostInfos);
           setupNodeIfNotExist(targetNode);
           if (i == MASTER_NODE_INDEX) {
+            primaryNodesCache.put(getNodeKey(targetNode), getNode(targetNode));
             assignSlotsToNode(slotNums, targetNode);
           } else if (clientConfig.isReadOnlyForRedisClusterReplicas()) {
             assignSlotsToReplicaNode(slotNums, targetNode);
@@ -425,6 +427,26 @@ public Map<String, ConnectionPool> getNodes() {
     }
   }
 
+  public Map<String, ConnectionPool> getPrimaryNodes() {
+    r.lock();
+    try {
+      return new HashMap<>(primaryNodesCache);
+    } finally {
+      r.unlock();
+    }
+  }
+
+  public List<ConnectionPool> getShuffledPrimaryNodesPool() {
+    r.lock();
+    try {
+      List<ConnectionPool> pools = new ArrayList<>(primaryNodesCache.values());
+      Collections.shuffle(pools);
+      return pools;
+    } finally {
+      r.unlock();
+    }
+  }
+
   public List<ConnectionPool> getShuffledNodesPool() {
     r.lock();
     try {
@@ -475,6 +497,7 @@ private void resetNodes() {
       }
     }
     nodes.clear();
+    primaryNodesCache.clear();
   }
 
   public void close() {

src/main/java/redis/clients/jedis/executors/ClusterCommandExecutor.java

@@ -45,7 +45,7 @@ public void close() {
 
   @Override
   public final <T> T broadcastCommand(CommandObject<T> commandObject) {
-    Map<String, ConnectionPool> connectionMap = provider.getConnectionMap();
+    Map<String, ConnectionPool> connectionMap = provider.getPrimaryNodesConnectionMap();
 
     boolean isErrored = false;
     T reply = null;

src/main/java/redis/clients/jedis/providers/ClusterConnectionProvider.java

@@ -112,6 +112,10 @@ public Map<String, ConnectionPool> getNodes() {
     return cache.getNodes();
   }
 
+  public Map<String, ConnectionPool> getPrimaryNodes() {
+    return cache.getPrimaryNodes();
+  }
+
   public HostAndPort getNode(int slot) {
     return slot >= 0 ? cache.getSlotNode(slot) : null;
   }
@@ -136,7 +140,7 @@ public Connection getConnection() {
     // In antirez's redis-rb-cluster implementation, getRandomConnection always return
     // valid connection (able to ping-pong) or exception if all connections are invalid
 
-    List<ConnectionPool> pools = cache.getShuffledNodesPool();
+    List<ConnectionPool> pools = cache.getShuffledPrimaryNodesPool();
 
     JedisException suppressed = null;
     for (ConnectionPool pool : pools) {
@@ -205,8 +209,15 @@ public Connection getReplicaConnectionFromSlot(int slot) {
     return getConnectionFromSlot(slot);
   }
 
+
   @Override
   public Map<String, ConnectionPool> getConnectionMap() {
     return Collections.unmodifiableMap(getNodes());
   }
+
+  @Override
+  public Map<String, ConnectionPool> getPrimaryNodesConnectionMap() {
+    return Collections.unmodifiableMap(getPrimaryNodes());
+  }
+
 }

src/main/java/redis/clients/jedis/providers/ConnectionProvider.java

@@ -15,4 +15,9 @@ public interface ConnectionProvider extends AutoCloseable {
     final Connection c = getConnection();
     return Collections.singletonMap(c.toString(), c);
   }
+
+  default Map<?, ?> getPrimaryNodesConnectionMap() {
+    final Connection c = getConnection();
+    return Collections.singletonMap(c.toString(), c);
+  }
 }

src/test/java/redis/clients/jedis/ClusterCommandExecutorTest.java

@@ -19,7 +19,6 @@
 import org.mockito.InOrder;
 import org.mockito.Mockito;
 import org.mockito.invocation.InvocationOnMock;
-import org.mockito.stubbing.Answer;
 
 import redis.clients.jedis.exceptions.JedisAskDataException;
 import redis.clients.jedis.exceptions.JedisClusterOperationException;

src/test/java/redis/clients/jedis/JedisClusterInfoCacheTest.java

@@ -15,12 +15,16 @@
 import java.util.stream.Collectors;
 
 import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.aMapWithSize;
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.hasEntry;
 import static org.hamcrest.Matchers.hasItem;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.Mockito.when;
+import static redis.clients.jedis.JedisClusterInfoCache.getNodeKey;
 import static redis.clients.jedis.Protocol.Command.CLUSTER;
 import static redis.clients.jedis.util.CommandArgumentMatchers.commandWithArgs;
 
@@ -49,7 +53,7 @@ public void testReplicaNodeRemovalAndRediscovery() {
 
     // Mock the cluster slots responses
     when(mockConnection.executeCommand(argThat(commandWithArgs(CLUSTER, "SLOTS")))).thenReturn(
-            masterReplicaSlotsResponse()).thenReturn(masterOnlySlotsResponse())
+            masterReplicaSlotsResponse(MASTER_HOST, REPLICA_1_HOST)).thenReturn(masterOnlySlotsResponse())
         .thenReturn(masterReplica2SlotsResponse());
 
     // Initial discovery with one master and one replica (replica-1)
@@ -78,7 +82,7 @@ public void testResetWithReplicaSlots() {
 
     // Mock the cluster slots responses
     when(mockConnection.executeCommand(argThat(commandWithArgs(CLUSTER, "SLOTS")))).thenReturn(
-        masterReplicaSlotsResponse());
+        masterReplicaSlotsResponse(MASTER_HOST, REPLICA_1_HOST));
 
     // Initial discovery
     cache.discoverClusterNodesAndSlots(mockConnection);
@@ -94,10 +98,68 @@ public void testResetWithReplicaSlots() {
     assertReplicasAvailable(cache, REPLICA_1_HOST);
   }
 
-  private List<Object> masterReplicaSlotsResponse() {
+  @Test
+  public void getPrimaryNodesAfterReplicaNodeRemovalAndRediscovery() {
+    // Create client config with read-only replicas enabled
+    JedisClientConfig clientConfig = DefaultJedisClientConfig.builder()
+            .readOnlyForRedisClusterReplicas().build();
+
+    Set<HostAndPort> startNodes = new HashSet<>();
+    startNodes.add(MASTER_HOST);
+
+    JedisClusterInfoCache cache = new JedisClusterInfoCache(clientConfig, startNodes);
+
+    // Mock the cluster slots responses
+    when(mockConnection.executeCommand(argThat(commandWithArgs(CLUSTER, "SLOTS")))).thenReturn(
+                    masterReplicaSlotsResponse(MASTER_HOST, REPLICA_1_HOST)).thenReturn(masterOnlySlotsResponse())
+            .thenReturn(masterReplica2SlotsResponse());
+
+    // Initial discovery with one master and one replica (replica-1)
+    cache.discoverClusterNodesAndSlots(mockConnection);
+    assertThat(cache.getPrimaryNodes(),aMapWithSize(1));
+    assertThat(cache.getPrimaryNodes(),
+                    hasEntry(equalTo(getNodeKey(MASTER_HOST)), equalTo(cache.getNode(MASTER_HOST))));
+
+    // Simulate rediscovery - master only
+    cache.discoverClusterNodesAndSlots(mockConnection);
+    assertThat(  cache.getPrimaryNodes(),aMapWithSize(1));
+    assertThat(cache.getPrimaryNodes(),
+            hasEntry(equalTo(getNodeKey(MASTER_HOST)), equalTo(cache.getNode(MASTER_HOST))));
+  }
+
+  @Test
+  public void getPrimaryNodesAfterMasterReplicaFailover() {
+    // Create client config with read-only replicas enabled
+    JedisClientConfig clientConfig = DefaultJedisClientConfig.builder()
+            .readOnlyForRedisClusterReplicas().build();
+
+    Set<HostAndPort> startNodes = new HashSet<>();
+    startNodes.add(MASTER_HOST);
+
+    JedisClusterInfoCache cache = new JedisClusterInfoCache(clientConfig, startNodes);
+
+    // Mock the cluster slots responses
+    when(mockConnection.executeCommand(argThat(commandWithArgs(CLUSTER, "SLOTS"))))
+            .thenReturn(masterReplicaSlotsResponse(MASTER_HOST, REPLICA_1_HOST))
+            .thenReturn(masterReplicaSlotsResponse(REPLICA_1_HOST, MASTER_HOST));
+
+    // Initial discovery with one master and one replica (replica-1)
+    cache.discoverClusterNodesAndSlots(mockConnection);
+    assertThat(cache.getPrimaryNodes(),aMapWithSize(1));
+    assertThat(cache.getPrimaryNodes(),
+            hasEntry(equalTo(getNodeKey(MASTER_HOST)), equalTo(cache.getNode(MASTER_HOST))));
+
+    // Simulate rediscovery - master only
+    cache.discoverClusterNodesAndSlots(mockConnection);
+    assertThat(  cache.getPrimaryNodes(),aMapWithSize(1));
+    assertThat(cache.getPrimaryNodes(),
+            hasEntry(equalTo(getNodeKey(REPLICA_1_HOST)), equalTo(cache.getNode(REPLICA_1_HOST))));
+  }
+
+  private List<Object> masterReplicaSlotsResponse(HostAndPort masterHost, HostAndPort replicaHost) {
     return createClusterSlotsResponse(
-        new SlotRange.Builder(0, 16383).master(MASTER_HOST, "master-id-1")
-            .replica(REPLICA_1_HOST, "replica-id-1").build());
+            new SlotRange.Builder(0, 16383).master(masterHost, masterHost.toString() + "-id")
+                    .replica(replicaHost, replicaHost.toString() + "-id").build());
   }
 
   private List<Object> masterOnlySlotsResponse() {