Notify forced disconnection (#62)

* Notify socket connection closed
* Check when the socket is disconnected in an unexpected way
* Add 2e2 test for the disconnection
* add github.com/michaelklishin/rabbit-hole/v2 lib for testing 

---------
Signed-off-by: Gabriele Santomaggio <G.santomaggio@gmail.com>
Signed-off-by: Aitor Pérez Cedres <acedres@vmware.com>
Co-authored-by: Aitor Pérez Cedres <acedres@vmware.com>

Author: Gsantomaggio

.github/workflows/build_test_linux.yml

@@ -32,10 +32,11 @@ jobs:
     runs-on: ubuntu-latest
     services:
       rabbitmq:
-        image: "rabbitmq:3"
+        image: "rabbitmq:3-management"
         ports:
           - 5552
           - 5672
+          - 15672:15672
         options: >-
           --health-cmd "rabbitmq-diagnostics check_running"
           --health-interval 10s
@@ -54,4 +55,4 @@ jobs:
       - name: E2E Tests
         env:
           RABBITMQ_CONTAINER_NAME: ${{ job.services.rabbitmq.id }}
-        run: make e2e_tests RMQ_E2E_SKIP_CONTAINER_START="please" RABBITMQ_URI="rabbitmq-stream://guest:guest@localhost:${{ job.services.rabbitmq.ports[5552] }}/%2F"
+        run: make e2e_tests RMQ_E2E_SKIP_CONTAINER_START="please" RABBITMQ_URI="rabbitmq-stream://guest:guest@localhost:${{ job.services.rabbitmq.ports[5552] }}/%2F" CI="actions"

go.mod

@@ -14,6 +14,7 @@ require (
 	github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572 // indirect
 	github.com/google/go-cmp v0.5.9 // indirect
 	github.com/google/pprof v0.0.0-20230502171905-255e3b9b56de // indirect
+	github.com/michaelklishin/rabbit-hole/v2 v2.13.0 // indirect
 	golang.org/x/net v0.9.0 // indirect
 	golang.org/x/sys v0.8.0 // indirect
 	golang.org/x/text v0.9.0 // indirect

go.sum

@@ -20,7 +20,7 @@ func NewPeerPropertiesRequest(connectionName string) *PeerPropertiesRequest {
 	p.clientProperties["copyright"] = "Copyright (c) 2023 VMware, Inc. or its affiliates."
 	p.clientProperties["information"] = "Licensed under the MPL 2.0. See https://www.rabbitmq.com/"
 	//c.clientProperties.items["version"] = ClientVersion
-	p.clientProperties["platform"] = "Golang V2"
+	p.clientProperties["platform"] = "Golang"
 	return p
 }
 

internal/peer_properties.go

@@ -34,6 +34,13 @@ func main() {
 	}
 	log.Info("connection status", "open", streamClient.IsOpen())
 
+	var closeChan = streamClient.NotifyConnectionClosed()
+	go func() {
+		for c := range closeChan {
+			log.Info("connection closed", "reason", c, "isOpen", streamClient.IsOpen())
+		}
+	}()
+
 	err = streamClient.DeclareStream(ctx, stream, map[string]string{"name": "test-stream"})
 	if err != nil && err.Error() != "stream already exists" {
 		log.Error("error in declaring stream", "error", err)
@@ -50,8 +57,6 @@ func main() {
 	const batchSize = 100
 	const iterations = 2
 	const totalMessages = iterations * batchSize
-	// PublishingId 555 --> [] messages 551, 552, 553, 554, 555
-	// ONLY 555
 	publishChan := streamClient.NotifyPublish(make(chan *raw.PublishConfirm, 100))
 	go func() {
 		var confirmed int

main.go

@@ -67,3 +67,10 @@ const (
 	CompressionNone uint8 = 0x00
 	CompressionGzip uint8 = 0x01
 )
+
+// Connections states
+const (
+	ConnectionClosed  = 0x01
+	ConnectionOpen    = 0x02
+	ConnectionClosing = 0x03
+)

pkg/constants/types.go

@@ -72,23 +72,26 @@ var _ = SynchronizedBeforeSuite(func(ctx SpecContext) {
 		containerName = defaultContainerName
 	}
 
-	pluginsCmd := exec.Command(
-		"docker",
-		"exec",
-		"--user=rabbitmq",
-		containerName,
-		"rabbitmq-plugins",
-		"enable",
-		"rabbitmq_stream",
-		"rabbitmq_stream_management",
-		"rabbitmq_management",
-	)
-	session, err := gexec.Start(pluginsCmd, GinkgoWriter, GinkgoWriter)
-	Expect(err).ToNot(HaveOccurred())
-	Eventually(session).
-		WithTimeout(time.Second*15).
-		WithPolling(time.Second).
-		Should(gexec.Exit(0), "expected to enable stream plugin")
+	if isCi := os.Getenv("CI"); isCi != "" {
+		pluginsCmd := exec.Command(
+			"docker",
+			"exec",
+			"--user=rabbitmq",
+			containerName,
+			"rabbitmq-plugins",
+			"enable",
+			"rabbitmq_stream",
+			"rabbitmq_stream_management",
+			"rabbitmq_management",
+		)
+
+		session, err := gexec.Start(pluginsCmd, GinkgoWriter, GinkgoWriter)
+		Expect(err).ToNot(HaveOccurred())
+		Eventually(session).
+			WithTimeout(time.Second*15).
+			WithPolling(time.Second).
+			Should(gexec.Exit(0), "expected to enable stream plugin")
+	}
 
 	if rabbitDebugLog {
 		debugCmd := exec.Command(

pkg/e2e/e2e_suite_test.go

@@ -217,6 +217,44 @@ var _ = Describe("E2E", Serial, Label("e2e"), func() {
 		Expect(streamClient.Close(itCtx)).To(Succeed())
 	}, SpecTimeout(15*time.Second))
 
+	// Test the notification in case of force disconnection.
+	// The disconnection is based on the connection name.
+	// With the HTTP API, we can check the connection name and kill it.
+	// The client has to notify the disconnection.
+	It("connection name and notify disconnection", Label("behaviour"), func(ctx SpecContext) {
+		h := slog.HandlerOptions{Level: slog.LevelDebug}.NewTextHandler(GinkgoWriter)
+		debugLogger := slog.New(h)
+		itCtx := raw.NewContextWithLogger(ctx, *debugLogger)
+		streamClientConfiguration, err := raw.NewClientConfiguration(rabbitmqUri)
+		Expect(err).ToNot(HaveOccurred())
+		connectionName := "notify-disconnection-test-1"
+		streamClientConfiguration.SetConnectionName(connectionName)
+
+		By("preparing the environment")
+		streamClient, err := raw.DialConfig(itCtx, streamClientConfiguration)
+		Expect(err).ToNot(HaveOccurred())
+		// Force close the connection is done with the HTTP API and based on the connection name.
+		httpUtils := NewHTTPUtils()
+		Eventually(func() string {
+			name, _ := httpUtils.GetConnectionByConnectionName(connectionName)
+			return name
+		}, 10*time.Second).WithPolling(1*time.Second).Should(Equal(connectionName),
+			"expected connection to be present")
+
+		c := streamClient.NotifyConnectionClosed()
+		By("Forcing closing the connection")
+		errClose := httpUtils.ForceCloseConnectionByConnectionName(connectionName)
+		Eventually(errClose).WithTimeout(time.Second * 10).WithPolling(time.Second).ShouldNot(HaveOccurred())
+
+		// the channel should be open and the notification should be received
+		select {
+		case notify, ok := <-c:
+			Expect(ok).To(BeTrue(), "expected the channel to be open")
+			Expect(notify).To(Equal(raw.ErrConnectionClosed))
+		case <-itCtx.Done():
+			Fail("expected to receive a closed notification")
+		}
+	}, SpecTimeout(20*time.Second))
 })
 
 func wrap[T any](v T) []T {

pkg/e2e/end_to_end_test.go

@@ -0,0 +1,52 @@
+package e2e_test
+
+import (
+	"fmt"
+	hole "github.com/michaelklishin/rabbit-hole/v2"
+)
+
+type HTTPUtils struct {
+	HttpClient *hole.Client
+}
+
+func NewHTTPUtils() *HTTPUtils {
+	client, err := hole.NewClient("http://localhost:15672", "guest", "guest")
+	if err != nil {
+		return nil
+	}
+	return &HTTPUtils{
+		HttpClient: client,
+	}
+}
+
+func (h *HTTPUtils) getConnectionInfoByName(connectionName string) (*hole.ConnectionInfo, error) {
+	connections, err := h.HttpClient.ListConnections()
+	if err != nil {
+		return nil, err
+	}
+	for i := range connections {
+		if connections[i].ClientProperties["connection_name"] != nil &&
+			connections[i].ClientProperties["connection_name"] == connectionName {
+			return &connections[i], nil
+		}
+	}
+	return nil, fmt.Errorf("connection not found %s", connectionName)
+}
+
+func (h *HTTPUtils) GetConnectionByConnectionName(connectionName string) (string, error) {
+	connection, err := h.getConnectionInfoByName(connectionName)
+	if err != nil {
+		return "", err
+	}
+	return connection.ClientProperties["connection_name"].(string), nil
+}
+
+func (h *HTTPUtils) ForceCloseConnectionByConnectionName(connectionName string) error {
+
+	connection, err := h.getConnectionInfoByName(connectionName)
+	if err != nil {
+		return err
+	}
+	_, err = h.HttpClient.CloseConnection(connection.Name)
+	return err
+}

pkg/e2e/http_utils_test.go

@@ -45,7 +45,6 @@ type Client struct {
 	mu sync.Mutex
 	// this channel is used for correlation-less incoming frames from the server
 	frameBodyListener chan internal.SyncCommandRead
-	isOpen            bool
 	// this function is used during shutdown to stop the background loop that reads from the network
 	ioLoopCancelFn       context.CancelFunc
 	connection           *internal.Connection
@@ -57,16 +56,23 @@ type Client struct {
 	publishErrorCh       chan *PublishError
 	chunkCh              chan *Chunk
 	notifyCh             chan *CreditError
-	metadataUpdateCh     chan *MetadataUpdate
-	consumerUpdateCh     chan *ConsumerUpdate
+	// see constants states of the connection
+	// we need different states to handle the case where the connection is closed by the server
+	// Open/ConnectionClosing/Closed
+	// ConnectionClosing is used to set the connection status no longer open, but still waiting for the server to close the connection
+	connectionStatus uint8
+	// socketClosedCh is used to notify the client that the socket has been closed
+	socketClosedCh   chan error
+	metadataUpdateCh chan *MetadataUpdate
+	consumerUpdateCh chan *ConsumerUpdate
 }
 
 // IsOpen returns true if the connection is open, false otherwise
 // IsOpen is thread-safe
 func (tc *Client) IsOpen() bool {
 	tc.mu.Lock()
 	defer tc.mu.Unlock()
-	return tc.isOpen
+	return tc.connectionStatus == constants.ConnectionOpen
 }
 
 // NewClient returns a common.Clienter implementation to interact with RabbitMQ stream using low level primitives.
@@ -76,7 +82,7 @@ func NewClient(connection net.Conn, configuration *ClientConfiguration) Clienter
 	rawClient := &Client{
 		frameBodyListener: make(chan internal.SyncCommandRead),
 		connection:        internal.NewConnection(connection),
-		isOpen:            false,
+		connectionStatus:  constants.ConnectionClosed,
 		correlationsMap:   sync.Map{},
 		configuration:     configuration,
 	}
@@ -222,6 +228,24 @@ func (tc *Client) handleIncoming(ctx context.Context) error {
 			if errors.Is(err, io.ErrClosedPipe) {
 				return nil
 			}
+
+			if errors.Is(err, io.EOF) {
+				// EOF is returned when the connection is closed
+				if tc.socketClosedCh != nil {
+					if tc.IsOpen() {
+						tc.socketClosedCh <- ErrConnectionClosed
+					}
+					// the TCP connection here is closed
+					// we close the channel since we don't need to send more than one message
+					close(tc.socketClosedCh)
+					tc.socketClosedCh = nil
+				}
+				// set the shutdown flag to false since we don't want to close the connection
+				// since it's already closed
+				_ = tc.shutdown(false)
+				return nil
+			}
+
 			if err != nil {
 				// TODO: some errors may be recoverable. We only need to return if reconnection
 				// 	is needed
@@ -278,7 +302,7 @@ func (tc *Client) handleIncoming(ctx context.Context) error {
 					// we do not return here because we must execute the shutdown process and close the socket
 				}
 
-				err = tc.shutdown()
+				err = tc.shutdown(true)
 				if err != nil && !errors.Is(err, io.ErrClosedPipe) {
 					return err
 				}
@@ -594,21 +618,34 @@ func (tc *Client) open(ctx context.Context, brokerIndex int) error {
 	return streamErrorOrNil(openResp.ResponseCode())
 }
 
-func (tc *Client) shutdown() error {
+func (tc *Client) shutdown(closeConnection bool) error {
 	tc.mu.Lock()
 	defer tc.mu.Unlock()
-	tc.isOpen = false
+	// The method can be called by the Close() method or by the
+	// connection error handler, see: handleIncoming EOF error
+	// the shutdown method has to be idempotent since it can be called multiple times
+	// In case of unexpected connection error, the shutdown is called just once
+	tc.connectionStatus = constants.ConnectionClosed
 	tc.ioLoopCancelFn()
 
 	if tc.confirmsCh != nil {
 		close(tc.confirmsCh)
+		tc.confirmsCh = nil
 	}
 
 	if tc.chunkCh != nil {
 		close(tc.chunkCh)
+		tc.chunkCh = nil
 	}
 
-	return tc.connection.Close()
+	// if the caller is handleIncoming EOF error closeConnection is false,
+	// The connection is already closed
+	// So we don't need to close it again
+
+	if closeConnection {
+		return tc.connection.Close()
+	}
+	return nil
 }
 
 func (tc *Client) handleClose(ctx context.Context, req *internal.CloseRequest) error {
@@ -863,7 +900,7 @@ func (tc *Client) Connect(ctx context.Context) error {
 	}
 
 	tc.mu.Lock()
-	tc.isOpen = true
+	tc.connectionStatus = constants.ConnectionOpen
 	defer tc.mu.Unlock()
 	log.Info("connection is open")
 
@@ -1113,6 +1150,8 @@ func (tc *Client) Close(ctx context.Context) error {
 		return errNilContext
 	}
 
+	tc.connectionStatus = constants.ConnectionClosing
+
 	log := loggerFromCtxOrDiscard(ctx).WithGroup("close")
 	log.Debug("starting connection close")
 
@@ -1127,7 +1166,7 @@ func (tc *Client) Close(ctx context.Context) error {
 		log.Error("close response code is not OK", "error", err)
 	}
 
-	err = tc.shutdown()
+	err = tc.shutdown(true)
 	if err != nil && !errors.Is(err, io.ErrClosedPipe) {
 		return err
 	}
@@ -1238,6 +1277,20 @@ func (tc *Client) NotifyChunk(c chan *Chunk) <-chan *Chunk {
 	return c
 }
 
+// NotifyConnectionClosed receives notifications about connection closed events.
+// It is raised only once when the connection is closed in unexpected way.
+// Connection gracefully closed by the client will not raise this event.
+func (tc *Client) NotifyConnectionClosed() <-chan error {
+	tc.mu.Lock()
+	defer tc.mu.Unlock()
+	// The user can't decide the size of the channel, so we use a buffer of 1.
+	// NotifyConnectionClosed is one shot notification, so we don't need a buffer.
+	// buffer greater than 1 cloud cause a deadlock since the channel is closed after the first notification.
+	c := make(chan error, 1)
+	tc.socketClosedCh = c
+	return c
+}
+
 // NotifyCreditError TODO: go docs
 func (tc *Client) NotifyCreditError(notification chan *CreditError) <-chan *CreditError {
 	tc.mu.Lock()