Repository: lf-edge/ekuiper
Pull request #3061 (2024-07-26 06:59:38)

fix: fix mqtt client ping


Patch:
diff --git a/internal/io/mqtt/client/client.go b/internal/io/mqtt/client/client.go
index fa32eef5..860beda0 100644
--- a/internal/io/mqtt/client/client.go
+++ b/internal/io/mqtt/client/client.go
@@ -114,11 +114,10 @@ func (conn *Connection) Close(ctx api.StreamContext) error {
 }
 
 func (conn *Connection) Ping(ctx api.StreamContext) error {
-	if conn.Client.IsConnected() {
+	if conn.connected.Load() {
 		return nil
-	} else {
-		return errors.New("mqtt ping failed")
 	}
+	return errors.New("failed to connect to broker")
 }
 
 func CreateConnection(ctx api.StreamContext, props map[string]any) (modules.Connection, error) {
diff --git a/pkg/connection/connection_pool.go b/pkg/connection/connection_pool.go
index ecddbae0..bdf5cf58 100644
--- a/pkg/connection/connection_pool.go
+++ b/pkg/connection/connection_pool.go
@@ -69,25 +69,39 @@ func GetConnectionRef(id string) int {
 }
 
 func GetAllConnectionStatus(ctx api.StreamContext) map[string]ConnectionStatus {
+	cws := make(map[string]*ConnWrapper)
 	globalConnectionManager.RLock()
-	defer globalConnectionManager.RUnlock()
-	s := make(map[string]ConnectionStatus)
 	for id, meta := range globalConnectionManager.connectionPool {
-		status := ConnectionStatus{
-			Status: ConnectionRunning,
-		}
-		conn, err := meta.cw.Wait()
-		if err != nil {
-			status.Status = ConnectionFail
-			status.ErrMsg = err.Error()
-		} else {
+		cws[id] = meta.cw
+	}
+	globalConnectionManager.RUnlock()
+	s := make(map[string]ConnectionStatus)
+	for id, cw := range cws {
+		if cw.IsInitialized() {
+			conn, err := cw.Wait()
+			if err != nil {
+				s[id] = ConnectionStatus{
+					Status: ConnectionFail,
+					ErrMsg: err.Error(),
+				}
+				continue
+			}
 			err = conn.Ping(ctx)
 			if err != nil {
-				status.Status = ConnectionFail
-				status.ErrMsg = err.Error()
+				s[id] = ConnectionStatus{
+					Status: ConnectionFail,
+					ErrMsg: err.Error(),
+				}
+				continue
+			}
+			s[id] = ConnectionStatus{
+				Status: ConnectionRunning,
+			}
+		} else {
+			s[id] = ConnectionStatus{
+				Status: ConnectionIntializing,
 			}
 		}
-		s[id] = status
 	}
 	return s
 }
@@ -371,13 +385,14 @@ const (
 var DefaultBackoffMaxElapsedDuration = 3 * time.Minute
 
 const (
-	ConnectionRunning = "running"
-	ConnectionFail    = "fail"
+	ConnectionRunning     = "running"
+	ConnectionIntializing = "initializing"
+	ConnectionFail        = "fail"
 )
 
 type ConnectionStatus struct {
-	Status string
-	ErrMsg string
+	Status string `json:"status"`
+	ErrMsg string `json:"errMsg,omitempty"`
 }
 
 func extractSelID(props map[string]interface{}) string {
@@ -419,6 +434,12 @@ func (cw *ConnWrapper) Wait() (modules.Connection, error) {
 	return cw.conn, cw.err
 }
 
+func (cw *ConnWrapper) IsInitialized() bool {
+	cw.cond.L.Lock()
+	defer cw.cond.L.Unlock()
+	return cw.initialized
+}
+
 func newConnWrapper(ctx api.StreamContext, meta *ConnectionMeta) *ConnWrapper {
 	cw := &ConnWrapper{
 		ID:   meta.ID,


Test patch:
diff --git a/internal/io/mqtt/client/client_test.go b/internal/io/mqtt/client/client_test.go
index 21741854..78de4a6c 100644
--- a/internal/io/mqtt/client/client_test.go
+++ b/internal/io/mqtt/client/client_test.go
@@ -16,11 +16,13 @@ package client
 
 import (
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"github.com/lf-edge/ekuiper/v2/internal/testx"
+	mockContext "github.com/lf-edge/ekuiper/v2/pkg/mock/context"
 )
 
 func TestValidate(t *testing.T) {
@@ -59,3 +61,21 @@ func TestValidate(t *testing.T) {
 		})
 	}
 }
+
+func TestMqttClientPing(t *testing.T) {
+	url, cancel, err := testx.InitBroker("TestMqttClientPing")
+	require.NoError(t, err)
+	ctx := mockContext.NewMockContext("1", "2")
+	c, err := CreateClient(ctx, map[string]any{
+		"server":     url,
+		"datasource": "demo",
+	})
+	require.NoError(t, err)
+	// wait connection done
+	time.Sleep(100 * time.Millisecond)
+	require.NoError(t, c.Ping(ctx))
+	cancel()
+	// wait status done
+	time.Sleep(100 * time.Millisecond)
+	require.Error(t, c.Ping(ctx))
+}
diff --git a/pkg/connection/connection_pool_test.go b/pkg/connection/connection_pool_test.go
index 9c63cc2b..aa7b129c 100644
--- a/pkg/connection/connection_pool_test.go
+++ b/pkg/connection/connection_pool_test.go
@@ -17,6 +17,7 @@ package connection
 import (
 	"sync"
 	"testing"
+	"time"
 
 	"github.com/pingcap/failpoint"
 	"github.com/stretchr/testify/require"
@@ -120,6 +121,7 @@ func TestConnectionStatus(t *testing.T) {
 	conf.WriteCfgIntoKVStorage("connections", "mockErr", "a1", map[string]interface{}{})
 	conf.WriteCfgIntoKVStorage("connections", "mock", "a2", map[string]interface{}{})
 	require.NoError(t, ReloadConnection())
+	time.Sleep(100 * time.Millisecond)
 	ctx := context.Background()
 	allStatus := GetAllConnectionStatus(ctx)
 	s, ok := allStatus["a1"]
@@ -135,6 +137,44 @@ func TestConnectionStatus(t *testing.T) {
 	}, s)
 }
 
+func TestGetAllConnectionStatus(t *testing.T) {
+	require.NoError(t, InitConnectionManager4Test())
+	ctx := mockContext.NewMockContext("id", "2")
+	cw1, err := FetchConnection(ctx, "id1", "mock", nil)
+	require.NoError(t, err)
+	cw2, err := FetchConnection(ctx, "id2", "mockErr", nil)
+	require.NoError(t, err)
+	cw3, err := FetchConnection(ctx, "id3", "blockconn", nil)
+	require.NoError(t, err)
+	cw1.Wait()
+	cw2.Wait()
+	allStatus := GetAllConnectionStatus(ctx)
+	s, ok := allStatus["id2"]
+	require.True(t, ok)
+	require.Equal(t, ConnectionStatus{
+		Status: ConnectionFail,
+		ErrMsg: "mockErr",
+	}, s)
+	s, ok = allStatus["id1"]
+	require.True(t, ok)
+	require.Equal(t, ConnectionStatus{
+		Status: ConnectionRunning,
+	}, s)
+	s, ok = allStatus["id3"]
+	require.True(t, ok)
+	require.Equal(t, ConnectionStatus{
+		Status: ConnectionIntializing,
+	}, s)
+	blockCh <- struct{}{}
+	cw3.Wait()
+	allStatus = GetAllConnectionStatus(ctx)
+	s, ok = allStatus["id3"]
+	require.True(t, ok)
+	require.Equal(t, ConnectionStatus{
+		Status: ConnectionRunning,
+	}, s)
+}
+
 func TestNonStoredConnection(t *testing.T) {
 	require.NoError(t, InitConnectionManager4Test())
 	ctx := mockContext.NewMockContext("id", "2")
@@ -151,13 +191,9 @@ func TestNonStoredConnection(t *testing.T) {
 	require.False(t, IsConnectionExists("id1"))
 }
 
-var blockCh chan any
-
 func TestConnectionLock(t *testing.T) {
 	require.NoError(t, InitConnectionManager4Test())
 	ctx := mockContext.NewMockContext("id", "2")
-	modules.RegisterConnection("blockconn", CreateBlockConnection)
-	blockCh = make(chan any, 10)
 	wg := sync.WaitGroup{}
 	wg.Add(1)
 	go func() {
@@ -170,7 +206,16 @@ func TestConnectionLock(t *testing.T) {
 	require.True(t, CheckConn("ccc1"))
 }
 
-type blockConnection struct{}
+var blockCh chan any
+
+func init() {
+	blockCh = make(chan any, 10)
+	modules.RegisterConnection("blockconn", CreateBlockConnection)
+}
+
+type blockConnection struct {
+	blochCh chan any
+}
 
 func (b blockConnection) Ping(ctx api.StreamContext) error {
 	return nil