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/consumer.go

@@ -312,6 +312,10 @@ func (c *ConsumerOptions) SetClientProvidedName(clientProvidedName string) *Cons
 	return c
 }
 
+func (c *ConsumerOptions) GetClientProvidedName() string {
+	return c.ClientProvidedName
+}
+
 func (c *ConsumerOptions) SetFilter(filter *ConsumerFilter) *ConsumerOptions {
 	c.Filter = filter
 	return c

pkg/stream/environment.go

@@ -457,6 +457,10 @@ func (envOptions *EnvironmentOptions) SetRPCTimeout(timeout time.Duration) *Envi
 	return envOptions
 }
 
+type clientOptions interface {
+	GetClientProvidedName() string
+}
+
 type environmentCoordinator struct {
 	mutex             *sync.Mutex
 	clientsPerContext sync.Map
@@ -535,95 +539,87 @@ func (c *Client) maybeCleanConsumers(streamName string) {
 	})
 }
 
-func (cc *environmentCoordinator) newProducer(leader *Broker, tcpParameters *TCPParameters, saslConfiguration *SaslConfiguration, streamName string, options *ProducerOptions, rpcTimeout time.Duration, cleanUp func()) (*Producer, error) {
+func (cc *environmentCoordinator) newClientEntity(
+	isListFull func(int) bool,
+	defaultClientName string,
+	leader *Broker,
+	tcpParameters *TCPParameters,
+	saslConfiguration *SaslConfiguration,
+	options clientOptions,
+	rpcTimeout time.Duration,
+) (*Client, error) {
 	cc.mutex.Lock()
 	defer cc.mutex.Unlock()
 	var clientResult *Client
 
 	cc.clientsPerContext.Range(func(key, value any) bool {
-		if !cc.isProducerListFull(key.(int)) {
+		if !isListFull(key.(int)) {
 			clientResult = value.(*Client)
 			return false
 		}
 		return true
 	})
 
-	clientProvidedName := "go-stream-producer"
-	if options != nil && options.ClientProvidedName != "" {
-		clientProvidedName = options.ClientProvidedName
+	clientProvidedName := defaultClientName
+	if options != nil && options.GetClientProvidedName() != "" {
+		clientProvidedName = options.GetClientProvidedName()
 	}
 
 	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)
-	}
-
-	producer, err := clientResult.declarePublisher(streamName, options, cleanUp)
+	return cc.validateBrokerConnection(clientResult, leader,
+		func() *Client {
+			return cc.newClientForConnection(clientProvidedName, leader, tcpParameters, saslConfiguration, rpcTimeout)
+		})
+}
 
+func (cc *environmentCoordinator) newProducer(leader *Broker, tcpParameters *TCPParameters, saslConfiguration *SaslConfiguration, streamName string, options *ProducerOptions, rpcTimeout time.Duration, cleanUp func()) (*Producer, error) {
+	client, err := cc.newClientEntity(cc.isProducerListFull, "go-stream-producer", leader, tcpParameters, saslConfiguration, options, rpcTimeout)
 	if err != nil {
 		return nil, err
 	}
-
-	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++
-
-	cc.clientsPerContext.Store(cc.nextId, clientResult)
-	return clientResult
+	return client.declarePublisher(streamName, options, cleanUp)
 }
 
 func (cc *environmentCoordinator) newConsumer(connectionName string, leader *Broker, tcpParameters *TCPParameters, saslConfiguration *SaslConfiguration,
 	streamName string, messagesHandler MessagesHandler,
 	options *ConsumerOptions, rpcTimeout time.Duration, cleanUp func()) (*Consumer, error) {
-	cc.mutex.Lock()
-	defer cc.mutex.Unlock()
-	var clientResult *Client
-
-	cc.clientsPerContext.Range(func(key, value any) bool {
-		if !cc.isConsumerListFull(key.(int)) {
-			clientResult = value.(*Client)
-			return false
-		}
-		return true
-	})
-
-	if clientResult == nil {
-		clientResult = newClient(connectionName, leader, tcpParameters, saslConfiguration, rpcTimeout)
-		cc.nextId++
-		cc.clientsPerContext.Store(cc.nextId, clientResult)
-	}
-	// try to reconnect in case the socket is closed
-	err := clientResult.connect()
+	client, err := cc.newClientEntity(cc.isConsumerListFull, "go-stream-consumer", leader, tcpParameters, saslConfiguration, options, rpcTimeout)
 	if err != nil {
 		return nil, err
 	}
+	return client.declareSubscriber(streamName, messagesHandler, options, cleanUp)
+}
 
-	subscriber, err := clientResult.declareSubscriber(streamName, messagesHandler, options, cleanUp)
-	if err != nil {
-		return nil, err
+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 subscriber, nil
+	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
 }
 
 func (cc *environmentCoordinator) Close() error {
@@ -664,7 +660,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
 	}

pkg/stream/producer.go

@@ -189,6 +189,10 @@ func (po *ProducerOptions) SetClientProvidedName(name string) *ProducerOptions {
 	return po
 }
 
+func (po *ProducerOptions) GetClientProvidedName() string {
+	return po.ClientProvidedName
+}
+
 // SetFilter sets the filter for the producer. See ProducerOptions.Filter for more details
 func (po *ProducerOptions) SetFilter(filter *ProducerFilter) *ProducerOptions {
 	po.Filter = filter