Fix super stream reconnection (#456)

* Fixes: #453
* Fixes: #458
---------

Signed-off-by: Gabriele Santomaggio <G.santomaggio@gmail.com>

Author: Gsantomaggio

examples/reliable_super_stream_getting_started/reliable_super_stream_getting_started.go

@@ -3,6 +3,7 @@ package main
 import (
 	"errors"
 	"fmt"
+	"time"
 
 	"github.com/rabbitmq/rabbitmq-stream-go-client/pkg/amqp"
 	"github.com/rabbitmq/rabbitmq-stream-go-client/pkg/ha"
@@ -27,7 +28,7 @@ func main() {
 	// Create a super stream
 	streamName := "my-super-stream"
 	// It is highly recommended to define the stream retention policy
-	err = env.DeclareSuperStream(streamName, stream.NewPartitionsOptions(3).
+	err = env.DeclareSuperStream(streamName, stream.NewPartitionsOptions(1).
 		SetMaxLengthBytes(stream.ByteCapacity{}.GB(2)))
 
 	// ignore the error if the stream already exists
@@ -75,17 +76,22 @@ func main() {
 		return
 	}
 
-	// Send a message
-	for i := range 10 {
+	// Send a messages
+	for i := range 10000 {
 		msg := amqp.NewMessage([]byte(fmt.Sprintf("Hello stream:%d", i)))
 		msg.Properties = &amqp.MessageProperties{
 			MessageID: fmt.Sprintf("msg-%d", i),
 		}
 		err = producer.Send(msg)
 		if err != nil {
 			fmt.Printf("Error sending message: %v\n", err)
-			return
+			time.Sleep(1 * time.Second)
 		}
+		if i%1000 == 0 {
+			fmt.Printf("Sent %d messages\n", i)
+		}
+		// add a small delay in case you want to kill connection to see the reliable reconnection in action
+		time.Sleep(200 * time.Millisecond)
 	}
 
 	// press any key to exit

pkg/ha/ha_super_stream_consumer.go

@@ -64,25 +64,29 @@ func NewReliableSuperStreamConsumer(env *stream.Environment, superStream string,
 
 func (r *ReliableSuperStreamConsumer) handleNotifyClose(channelClose chan stream.CPartitionClose) {
 	go func() {
-		cPartitionClose := <-channelClose
-		if strings.EqualFold(cPartitionClose.Event.Reason, stream.SocketClosed) || strings.EqualFold(cPartitionClose.Event.Reason, stream.MetaDataUpdate) || strings.EqualFold(cPartitionClose.Event.Reason, stream.ZombieConsumer) {
-			r.setStatus(StatusReconnecting)
-			logs.LogWarn("[Reliable] - %s closed unexpectedly %s.. Reconnecting..", r.getInfo(), cPartitionClose.Event.Reason)
-			r.bootstrap = false
-			err, reconnected := retry(1, r, cPartitionClose.Partition)
-			if err != nil {
-				logs.LogInfo(""+
-					"[Reliable] - %s won't be reconnected. Error: %s", r.getInfo(), err)
-			}
-			if reconnected {
-				r.setStatus(StatusOpen)
+		// for channelClose until closed
+		for cPartitionClose := range channelClose {
+			if strings.EqualFold(cPartitionClose.Event.Reason, stream.SocketClosed) || strings.EqualFold(cPartitionClose.Event.Reason, stream.MetaDataUpdate) || strings.EqualFold(cPartitionClose.Event.Reason, stream.ZombieConsumer) {
+				r.setStatus(StatusReconnecting)
+				logs.LogWarn("[Reliable] - %s closed unexpectedly %s.. Reconnecting..", r.getInfo(), cPartitionClose.Event.Reason)
+				r.bootstrap = false
+				err, reconnected := retry(1, r, cPartitionClose.Partition)
+				if err != nil {
+					logs.LogInfo(""+
+						"[Reliable] - %s won't be reconnected. Error: %s", r.getInfo(), err)
+				}
+				if reconnected {
+					r.setStatus(StatusOpen)
+				} else {
+					r.setStatus(StatusClosed)
+				}
 			} else {
+				logs.LogInfo("[Reliable] - %s closed normally. Reason: %s", r.getInfo(), cPartitionClose.Event.Reason)
 				r.setStatus(StatusClosed)
+				break
 			}
-		} else {
-			logs.LogInfo("[Reliable] - %s closed normally. Reason: %s", r.getInfo(), cPartitionClose.Event.Reason)
-			r.setStatus(StatusClosed)
 		}
+		logs.LogDebug("[ReliableSuperStreamConsumer] - cPartitionClose closed %s", r.getInfo())
 	}()
 }
 
@@ -103,6 +107,7 @@ func (r *ReliableSuperStreamConsumer) getEnv() *stream.Environment {
 
 func (r *ReliableSuperStreamConsumer) getNewInstance(partition string) newEntityInstance {
 	return func() error {
+		c := r.consumer.Load()
 		// by default the consumer will start from the consumerOptions.Offset
 		off := r.consumerOptions.Offset
 		var restartOffset int64
@@ -113,7 +118,8 @@ func (r *ReliableSuperStreamConsumer) getNewInstance(partition string) newEntity
 			restartOffset = v.(int64)
 			off = stream.OffsetSpecification{}.Offset(restartOffset + 1)
 		}
-		return r.consumer.Load().ConnectPartition(partition, off)
+
+		return c.ConnectPartition(partition, off)
 	}
 }
 

pkg/ha/ha_super_stream_publisher.go

@@ -71,27 +71,30 @@ func (r *ReliableSuperStreamProducer) handlePublishConfirm(confirm chan stream.P
 
 func (r *ReliableSuperStreamProducer) handleNotifyClose(channelClose chan stream.PPartitionClose) {
 	go func() {
-		cPartitionClose := <-channelClose
-		if strings.EqualFold(cPartitionClose.Event.Reason, stream.SocketClosed) || strings.EqualFold(cPartitionClose.Event.Reason, stream.MetaDataUpdate) || strings.EqualFold(cPartitionClose.Event.Reason, stream.ZombieConsumer) {
-			r.setStatus(StatusReconnecting)
-			logs.LogWarn("[Reliable] - %s closed unexpectedly %s.. Reconnecting..", r.getInfo(), cPartitionClose.Event.Reason)
-			err, reconnected := retry(1, r, cPartitionClose.Partition)
-			if err != nil {
-				logs.LogInfo(""+
-					"[Reliable] - %s won't be reconnected. Error: %s", r.getInfo(), err)
-			}
-			if reconnected {
-				r.setStatus(StatusOpen)
+		for cPartitionClose := range channelClose {
+			if strings.EqualFold(cPartitionClose.Event.Reason, stream.SocketClosed) || strings.EqualFold(cPartitionClose.Event.Reason, stream.MetaDataUpdate) || strings.EqualFold(cPartitionClose.Event.Reason, stream.ZombieConsumer) {
+				r.setStatus(StatusReconnecting)
+				logs.LogWarn("[Reliable] - %s closed unexpectedly %s.. Reconnecting..", r.getInfo(), cPartitionClose.Event.Reason)
+				err, reconnected := retry(1, r, cPartitionClose.Partition)
+				if err != nil {
+					logs.LogInfo(""+
+						"[Reliable] - %s won't be reconnected. Error: %s", r.getInfo(), err)
+				}
+				if reconnected {
+					r.setStatus(StatusOpen)
+				} else {
+					r.setStatus(StatusClosed)
+				}
 			} else {
+				logs.LogInfo("[Reliable] - %s closed normally. Reason: %s", r.getInfo(), cPartitionClose.Event.Reason)
 				r.setStatus(StatusClosed)
+				break
 			}
-		} else {
-			logs.LogInfo("[Reliable] - %s closed normally. Reason: %s", r.getInfo(), cPartitionClose.Event.Reason)
-			r.setStatus(StatusClosed)
+			r.reconnectionSignal.L.Lock()
+			r.reconnectionSignal.Broadcast()
+			r.reconnectionSignal.L.Unlock()
 		}
-		r.reconnectionSignal.L.Lock()
-		r.reconnectionSignal.Broadcast()
-		r.reconnectionSignal.L.Unlock()
+		logs.LogDebug("[ReliableSuperStreamProducer] - closed %s", r.getInfo())
 	}()
 }
 
@@ -112,7 +115,8 @@ func (r *ReliableSuperStreamProducer) getEnv() *stream.Environment {
 
 func (r *ReliableSuperStreamProducer) getNewInstance(streamName string) newEntityInstance {
 	return func() error {
-		return r.producer.Load().ConnectPartition(streamName)
+		p := r.producer.Load()
+		return p.ConnectPartition(streamName)
 	}
 }
 

pkg/stream/environment.go

@@ -502,8 +502,13 @@ func (cc *environmentCoordinator) isConsumerListFull(clientsPerContextId int) bo
 }
 
 func (cc *environmentCoordinator) maybeCleanClients() {
-	cc.mutex.Lock()
-	defer cc.mutex.Unlock()
+	// Note: Mutex is not needed here because:
+	// 1. sync.Map operations (Range, Delete) are thread-safe and can be called concurrently
+	// 2. Deleting the current entry during Range iteration is safe per Go's sync.Map documentation
+	// 3. This function is called from cleanup callbacks which may run concurrently, but
+	//    sync.Map handles concurrent access safely without requiring external synchronization
+	// 4. We only delete entries for clients that are already closed (socket.isOpen() == false),
+	//    so there's no risk of deleting active clients that are being used elsewhere
 
 	cc.clientsPerContext.Range(func(key, value any) bool {
 		client := value.(*Client)
@@ -620,7 +625,13 @@ func (cc *environmentCoordinator) validateBrokerConnection(client *Client, broke
 		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()
+		// Safety check: Only close the client if there are no active consumers or producers.
+		// This prevents premature disconnection during active operations, which could cause
+		// message loss or connection errors. If there are active producers/consumers, we
+		// create a new client without closing the old one, allowing graceful migration.
+		if client.coordinator.ConsumersCount() == 0 && client.coordinator.ProducersCount() == 0 {
+			client.Close()
+		}
 		client = newClientFunc()
 		err := client.connect()
 		if err != nil {

pkg/stream/super_stream_consumer.go

@@ -87,9 +87,11 @@ type SuperStreamConsumer struct {
 	// in a normal situation len(partitions) == len(consumers)
 	// but in case of disconnection the len(partitions) can be > len(consumers)
 	// since the consumer is in reconnection
-	partitions                  []string
-	env                         *Environment
-	mutex                       sync.Mutex
+	partitions []string
+	env        *Environment
+	mutex      sync.Mutex
+
+	chSuperStreamPartitionMutex sync.Mutex
 	chSuperStreamPartitionClose chan CPartitionClose
 
 	SuperStream                string
@@ -225,15 +227,15 @@ func (s *SuperStreamConsumer) ConnectPartition(partition string, offset OffsetSp
 			}
 		}
 		s.mutex.Unlock()
+		s.chSuperStreamPartitionMutex.Lock()
 		if s.chSuperStreamPartitionClose != nil {
-			s.mutex.Lock()
 			s.chSuperStreamPartitionClose <- CPartitionClose{
 				Partition: gpartion,
 				Event:     event,
 				Context:   s,
 			}
-			s.mutex.Unlock()
 		}
+		s.chSuperStreamPartitionMutex.Unlock()
 		logs.LogDebug("[SuperStreamConsumer] chSuperStreamPartitionClose for partition: %s", gpartion)
 	}(partition, closedEvent)
 
@@ -255,11 +257,11 @@ func (s *SuperStreamConsumer) Close() error {
 	// give the time to raise the close event
 	go func() {
 		time.Sleep(2 * time.Second)
-		s.mutex.Lock()
+		s.chSuperStreamPartitionMutex.Lock()
 		if s.chSuperStreamPartitionClose != nil {
 			close(s.chSuperStreamPartitionClose)
 		}
-		s.mutex.Unlock()
+		s.chSuperStreamPartitionMutex.Unlock()
 	}()
 
 	logs.LogDebug("[SuperStreamConsumer] Closed SuperStreamConsumer for: %s", s.SuperStream)

pkg/stream/super_stream_producer.go

@@ -172,9 +172,12 @@ type SuperStreamProducer struct {
 	// since the producer is in reconnection
 	partitions []string
 
-	env                         *Environment
-	mutex                       sync.Mutex
+	env   *Environment
+	mutex sync.Mutex
+
 	chNotifyPublishConfirmation chan PartitionPublishConfirm
+
+	chSuperStreamPartitionMutex sync.Mutex
 	chSuperStreamPartitionClose chan PPartitionClose
 
 	// public
@@ -220,8 +223,10 @@ func (s *SuperStreamProducer) init() error {
 		return err
 	}
 	for _, p := range partitions {
+		logs.LogDebug("Producing partition: %s", p)
 		err = s.ConnectPartition(p)
 		if err != nil {
+			logs.LogError("Failed to connect partition: %s", p)
 			return err
 		}
 	}
@@ -233,6 +238,10 @@ func (s *SuperStreamProducer) init() error {
 // that are hidden to the user.
 // with the ConnectPartition the user can re-connect a partition to the SuperStreamProducer
 // that should be used only in case of disconnection
+//
+// Note on mutex usage: The mutex is held while checking partition validity and activeProducers,
+// but is released before calling NewProducer() to avoid potential deadlocks. The mutex is
+// re-acquired after NewProducer() returns to safely update the activeProducers slice.
 func (s *SuperStreamProducer) ConnectPartition(partition string) error {
 	logs.LogDebug("[SuperStreamProducer] ConnectPartition for partition: %s", partition)
 
@@ -255,17 +264,25 @@ func (s *SuperStreamProducer) ConnectPartition(partition string) error {
 		}
 	}
 
-	s.mutex.Unlock()
 	var options = NewProducerOptions()
 	if s.SuperStreamProducerOptions.ClientProvidedName != "" {
 		options.ClientProvidedName = s.SuperStreamProducerOptions.ClientProvidedName
 	}
 	options = options.SetFilter(s.SuperStreamProducerOptions.Filter)
 
+	// Unlock mutex before calling NewProducer() to avoid potential deadlocks.
+	// NewProducer() may perform network I/O and could block, and it may also
+	// need to acquire other locks internally. Holding this mutex during that
+	// call could cause deadlocks if other goroutines are waiting on this mutex
+	// or if NewProducer() needs to acquire locks that conflict with this one.
+	s.mutex.Unlock()
+
 	producer, err := s.env.NewProducer(partition, options)
 	if err != nil {
 		return err
 	}
+
+	// Re-acquire mutex to safely update activeProducers slice
 	s.mutex.Lock()
 	s.activeProducers = append(s.activeProducers, producer)
 	chSingleStreamPublishConfirmation := producer.NotifyPublishConfirmation()
@@ -277,21 +294,23 @@ func (s *SuperStreamProducer) ConnectPartition(partition string) error {
 		event := <-_closedEvent
 
 		s.mutex.Lock()
-		defer s.mutex.Unlock()
 		for i := range s.activeProducers {
 			if s.activeProducers[i].GetStreamName() == gpartion {
 				s.activeProducers = append(s.activeProducers[:i], s.activeProducers[i+1:]...)
 				break
 			}
 		}
+		s.mutex.Unlock()
 
+		s.chSuperStreamPartitionMutex.Lock()
 		if s.chSuperStreamPartitionClose != nil {
 			s.chSuperStreamPartitionClose <- PPartitionClose{
 				Partition: gpartion,
 				Event:     event,
 				Context:   s,
 			}
 		}
+		s.chSuperStreamPartitionMutex.Unlock()
 		logs.LogDebug("[SuperStreamProducer] chSuperStreamPartitionClose for partition: %s", gpartion)
 	}(partition, closedEvent)
 
@@ -368,11 +387,12 @@ func (s *SuperStreamProducer) Send(message message.StreamMessage) error {
 	for _, p := range ps {
 		producer := s.getProducer(p)
 		if producer == nil {
-			// the producer is not. It can happen if the tcp connection for the partition is dropped
+			// the producer is not in the list. It can happen if the tcp connection for the partition is dropped
 			// the user can reconnect the partition using the ConnectPartition
 			// The client returns an error. Even there could be other partitions where the message can be sent.
 			// but won't to that to break the expectation of the user. The routing should be always the same
 			// for the same message. The user has to handle the error and decide to send the message again
+
 			return ErrProducerNotFound
 		}
 
@@ -404,11 +424,14 @@ func (s *SuperStreamProducer) Close() error {
 			close(s.chNotifyPublishConfirmation)
 			s.chNotifyPublishConfirmation = nil
 		}
+		s.mutex.Unlock()
+
+		s.chSuperStreamPartitionMutex.Lock()
 		if s.chSuperStreamPartitionClose != nil {
 			close(s.chSuperStreamPartitionClose)
 			s.chSuperStreamPartitionClose = nil
 		}
-		s.mutex.Unlock()
+		s.chSuperStreamPartitionMutex.Unlock()
 	}()
 	logs.LogDebug("[SuperStreamProducer] Closed SuperStreamProducer for: %s", s.SuperStream)
 	return nil