revise_id

Signed-off-by: Song Gao <[email protected]>
lf-edge · Jul 5, 2024 · cf82a65 · cf82a65
1 parent e0ae706
commit cf82a65
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/internal/io/websocket/websocket_sink.go b/internal/io/websocket/websocket_sink.go
@@ -46,11 +46,11 @@ func (w *WebsocketSink) Provision(ctx api.StreamContext, configs map[string]any)
 
 func (w *WebsocketSink) Close(ctx api.StreamContext) error {
 	pubsub.RemovePub(w.topic)
-	return connection.DetachConnection(ctx, w.cfg.Endpoint, w.props)
+	return connection.DetachConnection(ctx, buildWebsocketEpID(w.cfg.Endpoint), w.props)
 }
 
 func (w *WebsocketSink) 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
 	}
@@ -77,3 +77,7 @@ func GetSink() api.Sink {
 }
 
 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
@@ -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
 	}