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.

Use same retry delay as .NET stream client

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

@@ -554,26 +554,20 @@ func (cc *environmentCoordinator) newProducer(leader *Broker, tcpParameters *TCP
 	}
 
 	if clientResult == nil {
-		clientResult = cc.newClientForProducer(clientProvidedName, leader, tcpParameters, saslConfiguration, rpcTimeout)
+		clientResult = cc.newClientForConnection(clientProvidedName, leader, tcpParameters, saslConfiguration, rpcTimeout)
 	}
 
 	err := clientResult.connect()
 	if err != nil {
 		return nil, err
 	}
 
-	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.newClientForProducer(clientProvidedName, leader, tcpParameters, saslConfiguration, rpcTimeout)
-		err = clientResult.connect()
-		if err != nil {
-			return nil, err
-		}
-		time.Sleep(1 * time.Second)
+	clientResult, err = cc.validateBrokerConnection(clientResult, leader,
+		func() *Client {
+			return cc.newClientForConnection(clientProvidedName, leader, tcpParameters, saslConfiguration, rpcTimeout)
+		})
+	if err != nil {
+		return nil, err
 	}
 
 	producer, err := clientResult.declarePublisher(streamName, options, cleanUp)
@@ -585,10 +579,26 @@ func (cc *environmentCoordinator) newProducer(leader *Broker, tcpParameters *TCP
 	return producer, nil
 }
 
-func (cc *environmentCoordinator) newClientForProducer(connectionName string, leader *Broker, tcpParameters *TCPParameters, saslConfiguration *SaslConfiguration, rpcTimeOut time.Duration) *Client {
-	clientResult := newClient(connectionName, leader, tcpParameters, saslConfiguration, rpcTimeOut)
-	cc.nextId++
+func (cc *environmentCoordinator) validateBrokerConnection(client *Client, broker *Broker, newClientFunc func() *Client) (*Client, error) {
+	for client.connectionProperties.host != broker.advHost ||
+		client.connectionProperties.port != broker.advPort {
+		logs.LogDebug("connectionProperties host %s doesn't match with the advertised_host %s, advertised_port %s .. retry",
+			client.connectionProperties.host,
+			broker.advHost, broker.advPort)
+		client.Close()
+		client = newClientFunc()
+		err := client.connect()
+		if err != nil {
+			return nil, err
+		}
+		time.Sleep(time.Duration(500+rand.Intn(1000)) * time.Millisecond)
+	}
+	return client, nil
+}
 
+func (cc *environmentCoordinator) newClientForConnection(connectionName string, broker *Broker, tcpParameters *TCPParameters, saslConfiguration *SaslConfiguration, rpcTimeout time.Duration) *Client {
+	clientResult := newClient(connectionName, broker, tcpParameters, saslConfiguration, rpcTimeout)
+	cc.nextId++
 	cc.clientsPerContext.Store(cc.nextId, clientResult)
 	return clientResult
 }
@@ -619,6 +629,14 @@ func (cc *environmentCoordinator) newConsumer(connectionName string, leader *Bro
 		return nil, err
 	}
 
+	clientResult, err = cc.validateBrokerConnection(clientResult, leader,
+		func() *Client {
+			return cc.newClientForConnection(connectionName, leader, tcpParameters, saslConfiguration, rpcTimeout)
+		})
+	if err != nil {
+		return nil, err
+	}
+
 	subscriber, err := clientResult.declareSubscriber(streamName, messagesHandler, options, cleanUp)
 	if err != nil {
 		return nil, err
@@ -664,7 +682,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
 	}