Repository: lf-edge/ekuiper
Pull request #2802 (2024-04-24 13:28:06)

refactor(sink): split compression op


Patch:
diff --git a/internal/io/mqtt/sink.go b/internal/io/mqtt/sink.go
index 762fcce4..cf8d92e7 100644
--- a/internal/io/mqtt/sink.go
+++ b/internal/io/mqtt/sink.go
@@ -19,21 +19,19 @@ import (
 	"strings"
 
 	"github.com/lf-edge/ekuiper/contract/v2/api"
-	"github.com/lf-edge/ekuiper/v2/internal/compressor"
 	"github.com/lf-edge/ekuiper/v2/internal/io"
 	"github.com/lf-edge/ekuiper/v2/internal/topo/connection/clients"
 	mqttClient "github.com/lf-edge/ekuiper/v2/internal/topo/connection/clients/mqtt"
 	"github.com/lf-edge/ekuiper/v2/pkg/cast"
 	"github.com/lf-edge/ekuiper/v2/pkg/errorx"
-	"github.com/lf-edge/ekuiper/v2/pkg/message"
 )
 
 // AdConf is the advanced configuration for the mqtt sink
 type AdConf struct {
-	Tpc         string `json:"topic"`
-	Qos         byte   `json:"qos"`
-	Retained    bool   `json:"retained"`
-	Compression string `json:"compression"`
+	Tpc      string `json:"topic"`
+	Qos      byte   `json:"qos"`
+	Retained bool   `json:"retained"`
+
 	ResendTopic string `json:"resendDestination"`
 }
 
@@ -41,7 +39,6 @@ type MQTTSink struct {
 	adconf     *AdConf
 	config     map[string]interface{}
 	cli        io.MessageClient
-	compressor message.Compressor
 	sendParams map[string]any
 }
 
@@ -61,12 +58,6 @@ func (ms *MQTTSink) Provision(_ api.StreamContext, ps map[string]any) error {
 	if adconf.Qos != 0 && adconf.Qos != 1 && adconf.Qos != 2 {
 		return fmt.Errorf("invalid qos value %v, the value could be only int 0 or 1 or 2", adconf.Qos)
 	}
-	if adconf.Compression != "" {
-		ms.compressor, err = compressor.GetCompressor(adconf.Compression)
-		if err != nil {
-			return fmt.Errorf("invalid compression method %s", adconf.Compression)
-		}
-	}
 	ms.config = ps
 	if adconf.ResendTopic == "" {
 		adconf.ResendTopic = adconf.Tpc
diff --git a/internal/topo/node/compress_op.go b/internal/topo/node/compress_op.go
new file mode 100644
index 00000000..5ee1850b
--- /dev/null
+++ b/internal/topo/node/compress_op.go
@@ -0,0 +1,69 @@
+// 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 node
+
+import (
+	"fmt"
+
+	"github.com/lf-edge/ekuiper/contract/v2/api"
+	"github.com/lf-edge/ekuiper/v2/internal/compressor"
+	"github.com/lf-edge/ekuiper/v2/internal/pkg/def"
+	"github.com/lf-edge/ekuiper/v2/pkg/infra"
+	"github.com/lf-edge/ekuiper/v2/pkg/message"
+)
+
+type CompressOp struct {
+	*defaultSinkNode
+	tool message.Compressor
+}
+
+func NewCompressOp(name string, rOpt *def.RuleOption, compressMethod string) (*CompressOp, error) {
+	dc, err := compressor.GetCompressor(compressMethod)
+	if err != nil {
+		return nil, fmt.Errorf("get compressor %s fail with error: %v", compressMethod, err)
+	}
+	return &CompressOp{
+		defaultSinkNode: newDefaultSinkNode(name, rOpt),
+		tool:            dc,
+	}, nil
+}
+
+func (o *CompressOp) Exec(ctx api.StreamContext, errCh chan<- error) {
+	o.prepareExec(ctx, errCh, "op")
+	go func() {
+		err := infra.SafeRun(func() error {
+			runWithOrder(ctx, o.defaultSinkNode, o.concurrency, o.Worker)
+			return nil
+		})
+		if err != nil {
+			infra.DrainError(ctx, err, errCh)
+		}
+	}()
+}
+
+func (o *CompressOp) Worker(_ api.Logger, item any) []any {
+	o.statManager.ProcessTimeStart()
+	defer o.statManager.ProcessTimeEnd()
+	switch d := item.(type) {
+	case []byte:
+		if r, err := o.tool.Compress(d); err != nil {
+			return []any{err}
+		} else {
+			return []any{r}
+		}
+	default:
+		return []any{fmt.Errorf("unsupported data received: %v", d)}
+	}
+}
diff --git a/internal/topo/node/props.go b/internal/topo/node/props.go
index d7d6111a..d3f96093 100644
--- a/internal/topo/node/props.go
+++ b/internal/topo/node/props.go
@@ -36,6 +36,7 @@ type SinkConf struct {
 	DataField      string   `json:"dataField"`
 	BatchSize      int      `json:"batchSize"`
 	LingerInterval int      `json:"lingerInterval"`
+	Compression    string   `json:"compression"`
 	conf.SinkConf
 }
 
diff --git a/internal/topo/planner/planner_sink.go b/internal/topo/planner/planner_sink.go
index 1fced69d..38898bdb 100644
--- a/internal/topo/planner/planner_sink.go
+++ b/internal/topo/planner/planner_sink.go
@@ -103,6 +103,16 @@ func splitSink(tp *topo.Topo, inputs []node.Emitter, s api.Sink, sinkName string
 		index++
 		tp.AddOperator(newInputs, encodeOp)
 		newInputs = []node.Emitter{encodeOp}
+
+		if sc.Compression != "" {
+			compressOp, err := node.NewCompressOp(fmt.Sprintf("%s_%d_compress", sinkName, index), options, sc.Compression)
+			if err != nil {
+				return nil, err
+			}
+			index++
+			tp.AddOperator(newInputs, compressOp)
+			newInputs = []node.Emitter{compressOp}
+		}
 	}
 	return newInputs, nil
 }


Test patch:
diff --git a/internal/io/mqtt/sink_test.go b/internal/io/mqtt/sink_test.go
index 23799afc..bcaf23c7 100644
--- a/internal/io/mqtt/sink_test.go
+++ b/internal/io/mqtt/sink_test.go
@@ -76,7 +76,6 @@ func TestSinkConfigure(t *testing.T) {
 				Tpc:         "testTopic3",
 				Qos:         0,
 				Retained:    false,
-				Compression: "",
 				ResendTopic: "testTopic3",
 			},
 		},
@@ -94,7 +93,6 @@ func TestSinkConfigure(t *testing.T) {
 				Tpc:         "testTopic4",
 				Qos:         1,
 				Retained:    false,
-				Compression: "zlib",
 				ResendTopic: "testTopic4",
 			},
 		},
diff --git a/internal/topo/node/compress_op_test.go b/internal/topo/node/compress_op_test.go
new file mode 100644
index 00000000..5c637875
--- /dev/null
+++ b/internal/topo/node/compress_op_test.go
@@ -0,0 +1,73 @@
+// 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 node
+
+import (
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/lf-edge/ekuiper/v2/internal/pkg/def"
+	mockContext "github.com/lf-edge/ekuiper/v2/pkg/mock/context"
+)
+
+func TestNewCompressOp(t *testing.T) {
+	_, err := NewCompressOp("test", &def.RuleOption{}, "non")
+	assert.Error(t, err)
+	assert.Equal(t, "get compressor non fail with error: unsupported compressor: non", err.Error())
+	_, err = NewCompressOp("test", &def.RuleOption{}, "gzip")
+	assert.NoError(t, err)
+}
+
+func TestCompressOp_Exec(t *testing.T) {
+	op, err := NewCompressOp("test", &def.RuleOption{BufferLength: 10, SendError: true}, "gzip")
+	assert.NoError(t, err)
+	op.tool = &MockCompresser{}
+	out := make(chan any, 100)
+	err = op.AddOutput(out, "test")
+	assert.NoError(t, err)
+	ctx := mockContext.NewMockContext("test1", "compress_test")
+	errCh := make(chan error)
+	op.Exec(ctx, errCh)
+
+	cases := []any{
+		[]byte("{\"a\":1,\"b\":2}"),
+		errors.New("go through error"),
+		"invalid",
+	}
+	expects := [][]any{
+		{[]byte("mock compress")},
+		{errors.New("go through error")},
+		{errors.New("unsupported data received: invalid")},
+	}
+
+	for i, c := range cases {
+		op.input <- c
+		for _, e := range expects[i] {
+			r := <-out
+			switch tr := r.(type) {
+			case error:
+				assert.EqualError(t, e.(error), tr.Error())
+			default:
+				assert.Equal(t, e, r)
+			}
+		}
+	}
+}
+
+func (m *MockCompresser) Compress(_ []byte) ([]byte, error) {
+	return []byte("mock compress"), nil
+}
diff --git a/internal/topo/planner/planner_sink_test.go b/internal/topo/planner/planner_sink_test.go
index 3007658a..8c53c86a 100644
--- a/internal/topo/planner/planner_sink_test.go
+++ b/internal/topo/planner/planner_sink_test.go
@@ -86,6 +86,40 @@ func TestSinkPlan(t *testing.T) {
 				},
 			},
 		},
+		{
+			name: "batch and compress sink plan",
+			rule: &def.Rule{
+				Actions: []map[string]any{
+					{
+						"log": map[string]any{
+							"batchSize":   10,
+							"compression": "gzip",
+						},
+					},
+				},
+				Options: defaultOption,
+			},
+			topo: &def.PrintableTopo{
+				Sources: []string{"source_src1"},
+				Edges: map[string][]any{
+					"source_src1": {
+						"op_log_0_0_batch",
+					},
+					"op_log_0_0_batch": {
+						"op_log_0_1_transform",
+					},
+					"op_log_0_1_transform": {
+						"op_log_0_2_encode",
+					},
+					"op_log_0_2_encode": {
+						"op_log_0_3_compress",
+					},
+					"op_log_0_3_compress": {
+						"sink_log_0",
+					},
+				},
+			},
+		},
 	}
 	for _, c := range tc {
 		tp, err := topo.NewWithNameAndOptions("test", c.rule.Options)