Prefer replicas over leader for consumer connections

The Go client currently uses equal probability distribution for consumer
placement, including the leader node in the candidate pool alongside
replicas. This results in consumers connecting to the leader ~33% of the
time (with 2 replicas), which violates RabbitMQ Streams best practices
that recommend consumers connect to replicas rather than leaders.

This change modifies `BrokerForConsumerWithResolver` to only include the
leader in the candidate pool when no replicas are available. When replicas
exist, consumers randomly select from replicas only, guaranteeing zero
probability of leader placement. This matches the .NET client behavior and
follows RabbitMQ Streams best practices for consumer placement.

The leader is used as a fallback when `availableReplicas == 0`, ensuring
consumers can still connect when replicas are unavailable.

Author: lukebakken

pkg/stream/client.go

@@ -677,6 +677,10 @@ func (c *Client) queryPublisherSequence(publisherReference string, stream string
 }
 
 func (c *Client) BrokerLeader(stream string) (*Broker, error) {
+	return c.BrokerLeaderWithResolver(stream, nil)
+}
+
+func (c *Client) BrokerLeaderWithResolver(stream string, resolver *AddressResolver) (*Broker, error) {
 	streamsMetadata := c.metaData(stream)
 	if streamsMetadata == nil {
 		return nil, fmt.Errorf("leader error for stream for stream: %s", stream)
@@ -693,6 +697,13 @@ func (c *Client) BrokerLeader(stream string) (*Broker, error) {
 	streamMetadata.Leader.advPort = streamMetadata.Leader.Port
 	streamMetadata.Leader.advHost = streamMetadata.Leader.Host
 
+	// If AddressResolver is configured, use it directly and skip DNS lookup
+	if resolver != nil {
+		streamMetadata.Leader.Host = resolver.Host
+		streamMetadata.Leader.Port = strconv.Itoa(resolver.Port)
+		return streamMetadata.Leader, nil
+	}
+
 	res := net.Resolver{}
 	// see: https://github.com/rabbitmq/rabbitmq-stream-go-client/pull/317
 	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
@@ -738,12 +749,30 @@ func (c *Client) BrokerForConsumer(stream string) (*Broker, error) {
 	}
 
 	brokers := make([]*Broker, 0, 1+len(streamMetadata.Replicas))
-	brokers = append(brokers, streamMetadata.Leader)
+
+	// Count available replicas
+	availableReplicas := 0
+	for _, replica := range streamMetadata.Replicas {
+		if replica != nil {
+			availableReplicas++
+		}
+	}
+
+	// Only add leader if no replicas are available
+	if availableReplicas == 0 {
+		streamMetadata.Leader.advPort = streamMetadata.Leader.Port
+		streamMetadata.Leader.advHost = streamMetadata.Leader.Host
+		brokers = append(brokers, streamMetadata.Leader)
+	}
+
+	// Add all available replicas
 	for idx, replica := range streamMetadata.Replicas {
 		if replica == nil {
 			logs.LogWarn("Stream %s replica not ready: %d", stream, idx)
 			continue
 		}
+		replica.advPort = replica.Port
+		replica.advHost = replica.Host
 		brokers = append(brokers, replica)
 	}
 

pkg/stream/environment.go

@@ -619,13 +619,36 @@ func (cc *environmentCoordinator) newConsumer(connectionName string, leader *Bro
 		return nil, err
 	}
 
+	// Validate that connection reached the intended broker (replica vs leader)
+	// This is critical when using load balancers to ensure proper consumer placement
+	for clientResult.connectionProperties.host != leader.advHost ||
+		clientResult.connectionProperties.port != leader.advPort {
+		logs.LogDebug("connectionProperties host %s doesn't match with the advertised_host %s, advertised_port %s .. retry",
+			clientResult.connectionProperties.host,
+			leader.advHost, leader.advPort)
+		clientResult.Close()
+		clientResult = cc.newClientForConsumer(connectionName, leader, tcpParameters, saslConfiguration, rpcTimeout)
+		err = clientResult.connect()
+		if err != nil {
+			return nil, err
+		}
+		time.Sleep(1 * time.Second)
+	}
+
 	subscriber, err := clientResult.declareSubscriber(streamName, messagesHandler, options, cleanUp)
 	if err != nil {
 		return nil, err
 	}
 	return subscriber, nil
 }
 
+func (cc *environmentCoordinator) newClientForConsumer(connectionName string, leader *Broker, tcpParameters *TCPParameters, saslConfiguration *SaslConfiguration, rpcTimeout time.Duration) *Client {
+	clientResult := newClient(connectionName, leader, tcpParameters, saslConfiguration, rpcTimeout)
+	cc.nextId++
+	cc.clientsPerContext.Store(cc.nextId, clientResult)
+	return clientResult
+}
+
 func (cc *environmentCoordinator) Close() error {
 	cc.clientsPerContext.Range(func(_, value any) bool {
 		value.(*Client).coordinator.Close()
@@ -664,7 +687,7 @@ func (ps *producersEnvironment) newProducer(clientLocator *Client, streamName st
 	options *ProducerOptions, resolver *AddressResolver, rpcTimeOut time.Duration) (*Producer, error) {
 	ps.mutex.Lock()
 	defer ps.mutex.Unlock()
-	leader, err := clientLocator.BrokerLeader(streamName)
+	leader, err := clientLocator.BrokerLeaderWithResolver(streamName, resolver)
 	if err != nil {
 		return nil, err
 	}