Repository: lf-edge/ekuiper
Pull request #2985 (2024-07-04 08:20:06)

feat: support websocket sink

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/lf-edge/ekuiper/pull/2985?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lf-edge) Report
Attention: Patch coverage is `84.88372%` with `13 lines` in your changes missing coverage. Please review.
> Project coverage is 68.18%. Comparing base [(`e82fe62`)](https://app.codecov.io/gh/lf-edge/ekuiper/commit/e82fe622b1efb7b0956bb94fe6ab2b4bc994fdff?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lf-edge) to head [(`151f597`)](https://app.codecov.io/gh/lf-edge/ekuiper/commit/151f597acb3fe574ce41f98b13a70a543f178ea4?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lf-edge).

| [Files](https://app.codecov.io/gh/lf-edge/ekuiper/pull/2985?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lf-edge) | Patch % | Lines |
|---|---|---|
| [internal/io/websocket/websocket\_sink.go](https://app.codecov.io/gh/lf-edge/ekuiper/pull/2985?src=pr&el=tree&filepath=internal%2Fio%2Fwebsocket%2Fwebsocket_sink.go&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lf-edge#diff-aW50ZXJuYWwvaW8vd2Vic29ja2V0L3dlYnNvY2tldF9zaW5rLmdv) | 72.41% | [5 Missing and 3 partials :warning: ](https://app.codecov.io/gh/lf-edge/ekuiper/pull/2985?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lf-edge) |
| [internal/io/http/httpserver/websocket\_server.go](https://app.codecov.io/gh/lf-edge/ekuiper/pull/2985?src=pr&el=tree&filepath=internal%2Fio%2Fhttp%2Fhttpserver%2Fwebsocket_server.go&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lf-edge#diff-aW50ZXJuYWwvaW8vaHR0cC9odHRwc2VydmVyL3dlYnNvY2tldF9zZXJ2ZXIuZ28=) | 90.57% | [4 Missing and 1 partial :warning: ](https://app.codecov.io/gh/lf-edge/ekuiper/pull/2985?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lf-edge) |

<details><summary>Additional details and impacted files</summary>


```diff
@@              Coverage Diff              @@
##           master-v2    #2985      +/-   ##
=============================================
+ Coverage      68.15%   68.18%   +0.03%     
=============================================
  Files            338      339       +1     
  Lines          32857    32931      +74     
=============================================
+ Hits           22393    22452      +59     
- Misses          8455     8465      +10     
- Partials        2009     2014       +5     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/lf-edge/ekuiper/pull/2985?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lf-edge).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lf-edge).


Patch:
diff --git a/internal/binder/io/builtin.go b/internal/binder/io/builtin.go
index e38ed8c7..eb7cc113 100644
--- a/internal/binder/io/builtin.go
+++ b/internal/binder/io/builtin.go
@@ -47,7 +47,7 @@ func init() {
 	modules.RegisterSink("memory", func() api.Sink { return memory.GetSink() })
 	modules.RegisterSink("neuron", neuron.GetSink)
 	modules.RegisterSink("file", file.GetSink)
-	// modules.RegisterSink("websocket", func() api.Sink { return &websocket.WebSocketSink{} })
+	modules.RegisterSink("websocket", func() api.Sink { return websocket.GetSink() })
 
 	modules.RegisterLookupSource("memory", memory.GetLookupSource)
 	// modules.RegisterLookupSource("httppull", func() api.LookupSource { return http.GetLookUpSource() })
diff --git a/internal/io/http/httpserver/data_server.go b/internal/io/http/httpserver/data_server.go
index 8db68fc8..d2b22b6a 100644
--- a/internal/io/http/httpserver/data_server.go
+++ b/internal/io/http/httpserver/data_server.go
@@ -33,6 +33,7 @@ import (
 
 type GlobalServerManager struct {
 	sync.RWMutex
+	instanceID        int
 	endpoint          map[string]string
 	server            *http.Server
 	router            *mux.Router
diff --git a/internal/io/http/httpserver/websocketConn.go b/internal/io/http/httpserver/websocketConn.go
index 22c98c51..c106b974 100644
--- a/internal/io/http/httpserver/websocketConn.go
+++ b/internal/io/http/httpserver/websocketConn.go
@@ -22,6 +22,7 @@ import (
 
 type WebsocketConnection struct {
 	RecvTopic string
+	SendTopic string
 	cfg       *connectionCfg
 }
 
@@ -46,12 +47,13 @@ func createWebsocketServerConnection(ctx api.StreamContext, props map[string]any
 	if err := cast.MapToStruct(props, cfg); err != nil {
 		return nil, err
 	}
-	recvTopic, err := RegisterWebSocketEndpoint(ctx, cfg.Datasource)
+	rTopic, sTopic, err := RegisterWebSocketEndpoint(ctx, cfg.Datasource)
 	if err != nil {
 		return nil, err
 	}
 	return &WebsocketConnection{
-		RecvTopic: recvTopic,
+		RecvTopic: rTopic,
+		SendTopic: sTopic,
 		cfg:       cfg,
 	}, nil
 }
diff --git a/internal/io/http/httpserver/websocket_server.go b/internal/io/http/httpserver/websocket_server.go
index d2137bdf..fbef1d61 100644
--- a/internal/io/http/httpserver/websocket_server.go
+++ b/internal/io/http/httpserver/websocket_server.go
@@ -32,12 +32,20 @@ const (
 	WebsocketTopicPrefix = "$$websocket/"
 )
 
+func recvTopic(endpoint string) string {
+	return fmt.Sprintf("recv/%s/%s", WebsocketTopicPrefix, endpoint)
+}
+
+func sendTopic(endpoint string) string {
+	return fmt.Sprintf("send/%s/%s", WebsocketTopicPrefix, endpoint)
+}
+
 type websocketEndpointContext struct {
 	wg    *sync.WaitGroup
 	conns map[*websocket.Conn]context.CancelFunc
 }
 
-func RegisterWebSocketEndpoint(ctx api.StreamContext, endpoint string) (string, error) {
+func RegisterWebSocketEndpoint(ctx api.StreamContext, endpoint string) (string, string, error) {
 	return manager.RegisterWebSocketEndpoint(ctx, endpoint)
 }
 
@@ -49,14 +57,56 @@ func UnRegisterWebSocketEndpoint(endpoint string) {
 	}
 }
 
-func (m *GlobalServerManager) recvProcess(ctx api.StreamContext, endpoint string, c *websocket.Conn, wg *sync.WaitGroup) {
+func (m *GlobalServerManager) handleProcess(ctx api.StreamContext, endpoint string, instanceID int, c *websocket.Conn, cancel context.CancelFunc, parWg *sync.WaitGroup) {
 	defer func() {
 		m.CloseEndpointConnection(endpoint, c)
-		conf.Log.Infof("websocket endpoint %v stop recvProcess", endpoint)
+		parWg.Done()
+	}()
+	subWg := &sync.WaitGroup{}
+	subWg.Add(2)
+	go m.recvProcess(ctx, endpoint, c, cancel, subWg)
+	go m.sendProcess(ctx, endpoint, instanceID, c, cancel, subWg)
+	subWg.Wait()
+}
+
+func (m *GlobalServerManager) sendProcess(ctx api.StreamContext, endpoint string, instanceID int, c *websocket.Conn, cancel context.CancelFunc, wg *sync.WaitGroup) {
+	conf.Log.Infof("websocket endpoint %v start sendProcess", endpoint)
+	topic := sendTopic(endpoint)
+	sourceID := fmt.Sprintf("ws/send/%v", instanceID)
+	defer func() {
+		pubsub.CloseSourceConsumerChannel(topic, sourceID)
+		cancel()
+		c.Close()
+		wg.Done()
+		conf.Log.Infof("websocket send endpoint %v stop sendProcess", endpoint)
+	}()
+	ch := pubsub.CreateSub(topic, nil, sourceID, 1024)
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case d := <-ch:
+			data := d.([]byte)
+			if err := c.WriteMessage(websocket.TextMessage, data); err != nil {
+				if websocket.IsCloseError(err) || strings.Contains(err.Error(), "close") {
+					conf.Log.Infof("websocket endpoint %s connection get closed: %v", endpoint, err)
+					return
+				}
+				conf.Log.Warnf("websocket endpoint %v send data meet error: %v", endpoint, err)
+			}
+		}
+	}
+}
+
+func (m *GlobalServerManager) recvProcess(ctx api.StreamContext, endpoint string, c *websocket.Conn, cancel context.CancelFunc, wg *sync.WaitGroup) {
+	defer func() {
+		cancel()
+		c.Close()
 		wg.Done()
+		conf.Log.Infof("websocket recv endpoint %v stop recvProcess", endpoint)
 	}()
 	conf.Log.Infof("websocket endpoint %v start recvProcess", endpoint)
-	topic := fmt.Sprintf("recv/%s/%s", WebsocketTopicPrefix, endpoint)
+	topic := recvTopic(endpoint)
 	for {
 		select {
 		case <-ctx.Done():
@@ -81,12 +131,13 @@ func (m *GlobalServerManager) recvProcess(ctx api.StreamContext, endpoint string
 	}
 }
 
-func (m *GlobalServerManager) RegisterWebSocketEndpoint(ctx api.StreamContext, endpoint string) (string, error) {
+func (m *GlobalServerManager) RegisterWebSocketEndpoint(ctx api.StreamContext, endpoint string) (string, string, error) {
 	conf.Log.Infof("websocket endpoint %v register", endpoint)
 	m.Lock()
 	defer m.Unlock()
-	recvTopic := fmt.Sprintf("recv/%s/%s", WebsocketTopicPrefix, endpoint)
-	pubsub.CreatePub(recvTopic)
+	rTopic := recvTopic(endpoint)
+	sTopic := sendTopic(endpoint)
+	pubsub.CreatePub(rTopic)
 	m.router.HandleFunc(endpoint, func(w http.ResponseWriter, r *http.Request) {
 		c, err := m.upgrader.Upgrade(w, r, nil)
 		if err != nil {
@@ -95,16 +146,16 @@ func (m *GlobalServerManager) RegisterWebSocketEndpoint(ctx api.StreamContext, e
 		}
 		subCtx, cancel := ctx.WithCancel()
 		wg := m.AddEndpointConnection(endpoint, c, cancel)
+		go m.handleProcess(subCtx, endpoint, m.FetchInstanceID(), c, cancel, wg)
 		conf.Log.Infof("websocket endpint %v create connection", endpoint)
-		wg.Add(1)
-		go m.recvProcess(subCtx, endpoint, c, wg)
 	})
 	conf.Log.Infof("websocker endpoint %v registered success", endpoint)
-	return recvTopic, nil
+	return rTopic, sTopic, nil
 }
 
 func (m *GlobalServerManager) UnRegisterWebSocketEndpoint(endpoint string) *websocketEndpointContext {
 	conf.Log.Infof("websocket endpoint %v unregister", endpoint)
+	pubsub.RemovePub(recvTopic(endpoint))
 	m.Lock()
 	defer m.Unlock()
 	wctx, ok := m.websocketEndpoint[endpoint]
@@ -140,6 +191,7 @@ func (m *GlobalServerManager) AddEndpointConnection(endpoint string, c *websocke
 		return wctx.wg
 	}
 	wg := &sync.WaitGroup{}
+	wg.Add(1)
 	m.websocketEndpoint[endpoint] = &websocketEndpointContext{
 		wg: wg,
 		conns: map[*websocket.Conn]context.CancelFunc{
@@ -149,6 +201,11 @@ func (m *GlobalServerManager) AddEndpointConnection(endpoint string, c *websocke
 	return wg
 }
 
+func (m *GlobalServerManager) FetchInstanceID() int {
+	m.instanceID++
+	return m.instanceID
+}
+
 // getEndpointConnections only for unit test
 func (m *GlobalServerManager) getEndpointConnections(endpoint string) *websocketEndpointContext {
 	m.RLock()
diff --git a/internal/io/websocket/websocket_sink.go b/internal/io/websocket/websocket_sink.go
new file mode 100644
index 00000000..613a21d5
--- /dev/null
+++ b/internal/io/websocket/websocket_sink.go
@@ -0,0 +1,83 @@
+// Copyright 2024 EMQ Technologies Co., Ltd.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package websocket
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/lf-edge/ekuiper/contract/v2/api"
+	"github.com/lf-edge/ekuiper/v2/internal/io/http/httpserver"
+	"github.com/lf-edge/ekuiper/v2/internal/io/memory/pubsub"
+	"github.com/lf-edge/ekuiper/v2/pkg/cast"
+	"github.com/lf-edge/ekuiper/v2/pkg/connection"
+)
+
+type WebsocketSink struct {
+	cfg   *WebsocketConfig
+	props map[string]any
+	topic string
+}
+
+func (w *WebsocketSink) Provision(ctx api.StreamContext, configs map[string]any) error {
+	cfg := &WebsocketConfig{}
+	if err := cast.MapToStruct(configs, cfg); err != nil {
+		return err
+	}
+	if !strings.HasPrefix(cfg.Endpoint, "/") {
+		return fmt.Errorf("websocket endpoint should start with /")
+	}
+	w.cfg = cfg
+	w.props = configs
+	return nil
+}
+
+func (w *WebsocketSink) Close(ctx api.StreamContext) error {
+	pubsub.RemovePub(w.topic)
+	return connection.DetachConnection(ctx, buildWebsocketEpID(w.cfg.Endpoint), w.props)
+}
+
+func (w *WebsocketSink) Connect(ctx api.StreamContext) error {
+	conn, err := connection.FetchConnection(ctx, buildWebsocketEpID(w.cfg.Endpoint), "websocket", w.props)
+	if err != nil {
+		return err
+	}
+	c, ok := conn.(*httpserver.WebsocketConnection)
+	if !ok {
+		return fmt.Errorf("should use websocket connection")
+	}
+	w.topic = c.SendTopic
+	pubsub.CreatePub(w.topic)
+	return nil
+}
+
+func (w *WebsocketSink) Collect(ctx api.StreamContext, item api.RawTuple) error {
+	return w.collect(ctx, item.Raw())
+}
+
+func (w *WebsocketSink) collect(ctx api.StreamContext, data []byte) error {
+	pubsub.ProduceAny(ctx, w.topic, data)
+	return nil
+}
+
+func GetSink() api.Sink {
+	return &WebsocketSink{}
+}
+
+var _ api.BytesCollector = &WebsocketSink{}
+
+func buildWebsocketEpID(endpoint string) string {
+	return fmt.Sprintf("$$ws/%s", endpoint)
+}
diff --git a/internal/io/websocket/websocket_source.go b/internal/io/websocket/websocket_source.go
index e2f689cf..5dc6d2a1 100644
--- a/internal/io/websocket/websocket_source.go
+++ b/internal/io/websocket/websocket_source.go
@@ -53,11 +53,11 @@ func (w *WebsocketSource) Provision(ctx api.StreamContext, configs map[string]an
 
 func (w *WebsocketSource) Close(ctx api.StreamContext) error {
 	pubsub.CloseSourceConsumerChannel(w.topic, w.sourceID)
-	return connection.DetachConnection(ctx, w.cfg.Endpoint, w.props)
+	return connection.DetachConnection(ctx, buildWebsocketEpID(w.cfg.Endpoint), w.props)
 }
 
 func (w *WebsocketSource) Connect(ctx api.StreamContext) error {
-	conn, err := connection.FetchConnection(ctx, w.cfg.Endpoint, "websocket", w.props)
+	conn, err := connection.FetchConnection(ctx, buildWebsocketEpID(w.cfg.Endpoint), "websocket", w.props)
 	if err != nil {
 		return err
 	}


Test patch:
diff --git a/internal/io/http/httpserver/websocket_server_test.go b/internal/io/http/httpserver/websocket_server_test.go
index 5117350f..02f47afd 100644
--- a/internal/io/http/httpserver/websocket_server_test.go
+++ b/internal/io/http/httpserver/websocket_server_test.go
@@ -16,6 +16,7 @@ package httpserver
 
 import (
 	"testing"
+	"time"
 
 	"github.com/gorilla/websocket"
 	"github.com/stretchr/testify/require"
@@ -32,14 +33,16 @@ func TestWebsocketServerRecvData(t *testing.T) {
 	defer ShutDown()
 	ctx := mockContext.NewMockContext("1", "2")
 	endpint := "/e1"
-	recvTopic, err := RegisterWebSocketEndpoint(ctx, endpint)
+	rTopic, _, err := RegisterWebSocketEndpoint(ctx, endpint)
 	require.NoError(t, err)
-	subCh := pubsub.CreateSub(recvTopic, nil, "test", 1024)
-	defer pubsub.CloseSourceConsumerChannel(recvTopic, "test")
+	subCh := pubsub.CreateSub(rTopic, nil, "test", 1024)
+	defer pubsub.CloseSourceConsumerChannel(rTopic, "test")
 	conn, err := testx.CreateWebsocketClient(ip, port, endpint)
 	require.NoError(t, err)
 	defer conn.Close()
 	data := []byte("123")
+	// wait goroutine process started
+	time.Sleep(10 * time.Millisecond)
 	require.NoError(t, conn.WriteMessage(websocket.TextMessage, data))
 	recvData := <-subCh
 	require.Equal(t, data, recvData.([]byte))
@@ -53,7 +56,7 @@ func TestWebsocketServerRecvDataCancel(t *testing.T) {
 	defer ShutDown()
 	ctx := mockContext.NewMockContext("1", "2")
 	endpint := "/e1"
-	_, err := RegisterWebSocketEndpoint(ctx, endpint)
+	_, _, err := RegisterWebSocketEndpoint(ctx, endpint)
 	require.NoError(t, err)
 	UnRegisterWebSocketEndpoint(endpint)
 }
@@ -65,13 +68,47 @@ func TestWebsocketServerRecvDataOther(t *testing.T) {
 	defer ShutDown()
 	ctx := mockContext.NewMockContext("1", "2")
 	endpint := "/e1"
-	_, err := RegisterWebSocketEndpoint(ctx, endpint)
+	_, _, err := RegisterWebSocketEndpoint(ctx, endpint)
 	require.NoError(t, err)
 	conn, err := testx.CreateWebsocketClient(ip, port, endpint)
 	require.NoError(t, err)
+	defer conn.Close()
+	// wait goroutine process started
+	time.Sleep(10 * time.Millisecond)
 	require.NoError(t, conn.WriteMessage(websocket.PingMessage, []byte("123")))
 	require.NoError(t, conn.WriteMessage(websocket.CloseMessage, []byte("123")))
 	wctx := manager.getEndpointConnections(endpint)
 	wctx.wg.Wait()
 	require.Equal(t, 0, len(wctx.conns))
 }
+
+func TestWebsocketServerSendData(t *testing.T) {
+	endpoint := "/e1"
+	topic := sendTopic(endpoint)
+	pubsub.CreatePub(topic)
+	ip := "127.0.0.1"
+	port := 10085
+	InitGlobalServerManager(ip, port, nil)
+	defer ShutDown()
+	ctx := mockContext.NewMockContext("1", "2")
+	_, sTopic, err := RegisterWebSocketEndpoint(ctx, endpoint)
+	require.NoError(t, err)
+	require.Equal(t, topic, sTopic)
+	conn, err := testx.CreateWebsocketClient(ip, port, endpoint)
+	require.NoError(t, err)
+	defer conn.Close()
+	// wait goroutine process started
+	time.Sleep(10 * time.Millisecond)
+	assertCh := make(chan struct{})
+	go func() {
+		msgTyp, data, err := conn.ReadMessage()
+		require.NoError(t, err)
+		require.Equal(t, websocket.TextMessage, msgTyp)
+		require.Equal(t, []byte("123"), data)
+		assertCh <- struct{}{}
+	}()
+	time.Sleep(10 * time.Millisecond)
+	pubsub.ProduceAny(ctx, topic, []byte("123"))
+	<-assertCh
+	UnRegisterWebSocketEndpoint(endpoint)
+}
diff --git a/internal/io/websocket/websocket_sink_test.go b/internal/io/websocket/websocket_sink_test.go
new file mode 100644
index 00000000..55b9cd5e
--- /dev/null
+++ b/internal/io/websocket/websocket_sink_test.go
@@ -0,0 +1,68 @@
+// Copyright 2024 EMQ Technologies Co., Ltd.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package websocket
+
+import (
+	"testing"
+	"time"
+
+	"github.com/gorilla/websocket"
+	"github.com/stretchr/testify/require"
+
+	"github.com/lf-edge/ekuiper/v2/internal/io/http/httpserver"
+	"github.com/lf-edge/ekuiper/v2/internal/testx"
+	"github.com/lf-edge/ekuiper/v2/pkg/connection"
+	mockContext "github.com/lf-edge/ekuiper/v2/pkg/mock/context"
+	"github.com/lf-edge/ekuiper/v2/pkg/modules"
+)
+
+func init() {
+	modules.RegisterConnection("websocket", httpserver.CreateWebsocketConnection)
+}
+
+func TestWebsocketSink(t *testing.T) {
+	connection.InitConnectionManager4Test()
+	ip := "127.0.0.1"
+	port := 10081
+	endpoint := "/e1"
+	httpserver.InitGlobalServerManager(ip, port, nil)
+	defer httpserver.ShutDown()
+	ctx := mockContext.NewMockContext("1", "2")
+	ws := &WebsocketSink{}
+	props := map[string]any{
+		"datasource": endpoint,
+	}
+	require.Error(t, ws.Provision(ctx, map[string]any{
+		"datasource": "",
+	}))
+	require.NoError(t, ws.Provision(ctx, props))
+	require.NoError(t, ws.Connect(ctx))
+	expData := []byte("123")
+	assertCh := make(chan struct{})
+	conn, err := testx.CreateWebsocketClient(ip, port, endpoint)
+	require.NoError(t, err)
+	go func() {
+		msgTyp, data, err := conn.ReadMessage()
+		require.NoError(t, err)
+		require.Equal(t, websocket.TextMessage, msgTyp)
+		require.Equal(t, expData, data)
+		assertCh <- struct{}{}
+	}()
+	// wait goroutine start
+	time.Sleep(10 * time.Millisecond)
+	require.NoError(t, ws.collect(ctx, expData))
+	<-assertCh
+	ws.Close(ctx)
+}