refactor: abstract contube-go impl (#131)

Signed-off-by: Zike Yang <zike@apache.org>

Author: RobertIndie

benchmark/bench_test.go

@@ -22,6 +22,7 @@ import (
 	"github.com/apache/pulsar-client-go/pulsaradmin/pkg/utils"
 	"github.com/functionstream/functionstream/common"
 	"github.com/functionstream/functionstream/lib"
+	"github.com/functionstream/functionstream/lib/contube"
 	"github.com/functionstream/functionstream/perf"
 	"github.com/functionstream/functionstream/restclient"
 	"github.com/functionstream/functionstream/server"
@@ -35,12 +36,10 @@ import (
 
 func BenchmarkStressForBasicFunc(b *testing.B) {
 	s := server.New(server.LoadConfigFromEnv())
-	go s.Run()
+	svrCtx, svrCancel := context.WithCancel(context.Background())
+	go s.Run(svrCtx)
 	defer func() {
-		err := s.Close()
-		if err != nil {
-			b.Fatal(err)
-		}
+		svrCancel()
 	}()
 
 	inputTopic := "test-input-" + strconv.Itoa(rand.Int())
@@ -78,7 +77,7 @@ func BenchmarkStressForBasicFunc(b *testing.B) {
 
 	b.ReportAllocs()
 
-	ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(3*time.Second))
+	ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(30*time.Second))
 	defer cancel()
 
 	profile := b.Name() + ".pprof"
@@ -95,28 +94,27 @@ func BenchmarkStressForBasicFunc(b *testing.B) {
 		b.Fatal(err)
 	}
 
+	<-s.WaitForReady(context.Background())
 	perf.New(pConfig).Run(ctx)
 
 	pprof.StopCPUProfile()
 }
 
 func BenchmarkStressForBasicFuncWithMemoryQueue(b *testing.B) {
-	memoryQueueFactory := lib.NewMemoryQueueFactory(context.Background())
+	memoryQueueFactory := contube.NewMemoryQueueFactory(context.Background())
 
 	svrConf := &lib.Config{
 		ListenAddr: common.DefaultAddr,
-		QueueBuilder: func(ctx context.Context, config *lib.Config) (lib.EventQueueFactory, error) {
+		QueueBuilder: func(ctx context.Context, config *lib.Config) (contube.TubeFactory, error) {
 			return memoryQueueFactory, nil
 		},
 	}
 
 	s := server.New(svrConf)
-	go s.Run()
+	svrCtx, svrCancel := context.WithCancel(context.Background())
+	go s.Run(svrCtx)
 	defer func() {
-		err := s.Close()
-		if err != nil {
-			b.Fatal(err)
-		}
+		svrCancel()
 	}()
 
 	inputTopic := "test-input-" + strconv.Itoa(rand.Int())
@@ -132,14 +130,14 @@ func BenchmarkStressForBasicFuncWithMemoryQueue(b *testing.B) {
 			Output:   outputTopic,
 			Replicas: &replicas,
 		},
-		QueueBuilder: func(ctx context.Context, c *lib.Config) (lib.EventQueueFactory, error) {
+		QueueBuilder: func(ctx context.Context, c *lib.Config) (contube.TubeFactory, error) {
 			return memoryQueueFactory, nil
 		},
 	}
 
 	b.ReportAllocs()
 
-	ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(3*time.Second))
+	ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(30*time.Second))
 	defer cancel()
 
 	profile := b.Name() + ".pprof"
@@ -156,6 +154,7 @@ func BenchmarkStressForBasicFuncWithMemoryQueue(b *testing.B) {
 		b.Fatal(err)
 	}
 
+	<-s.WaitForReady(context.Background())
 	perf.New(pConfig).Run(ctx)
 
 	pprof.StopCPUProfile()

cmd/server/cmd.go

@@ -17,6 +17,7 @@
 package server
 
 import (
+	"context"
 	"github.com/functionstream/functionstream/common"
 	"github.com/functionstream/functionstream/server"
 	"github.com/spf13/cobra"
@@ -35,7 +36,7 @@ var (
 func exec(*cobra.Command, []string) {
 	common.RunProcess(func() (io.Closer, error) {
 		s := server.New(server.LoadConfigFromEnv())
-		go s.Run()
+		go s.Run(context.Background())
 		return s, nil
 	})
 }

cmd/standalone/cmd.go

@@ -17,6 +17,7 @@
 package standalone
 
 import (
+	"context"
 	"github.com/functionstream/functionstream/common"
 	"github.com/functionstream/functionstream/server"
 	"github.com/spf13/cobra"
@@ -35,7 +36,7 @@ var (
 func exec(*cobra.Command, []string) {
 	common.RunProcess(func() (io.Closer, error) {
 		s := server.New(server.LoadStandaloneConfigFromEnv())
-		go s.Run()
+		go s.Run(context.Background())
 		return s, nil
 	})
 }

lib/config.go

@@ -18,9 +18,10 @@ package lib
 
 import (
 	"context"
+	"github.com/functionstream/functionstream/lib/contube"
 )
 
-type QueueBuilder func(ctx context.Context, config *Config) (EventQueueFactory, error)
+type QueueBuilder func(ctx context.Context, config *Config) (contube.TubeFactory, error)
 
 type Config struct {
 	ListenAddr   string

lib/contube/event_tube.go

@@ -0,0 +1,105 @@
+/*
+ * Copyright 2024 Function Stream Org.
+ *
+ * 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 contube
+
+import (
+	"context"
+)
+
+type Record interface {
+	GetPayload() []byte
+	Commit()
+}
+
+type SourceQueueConfig struct {
+	Topics  []string
+	SubName string
+}
+
+type SinkQueueConfig struct {
+	Topic string
+}
+
+const (
+	TopicKey     = "topic"
+	TopicListKey = "topicList"
+	SubNameKey   = "subName"
+)
+
+func NewSourceQueueConfig(config ConfigMap) *SourceQueueConfig {
+	var result SourceQueueConfig
+	if topics, ok := config[TopicListKey].([]string); ok {
+		result.Topics = topics
+	}
+	if subName, ok := config[SubNameKey].(string); ok {
+		result.SubName = subName
+	}
+	return &result
+}
+
+func (c *SourceQueueConfig) ToConfigMap() ConfigMap {
+	return ConfigMap{
+		TopicListKey: c.Topics,
+		SubNameKey:   c.SubName,
+	}
+}
+
+func NewSinkQueueConfig(config ConfigMap) *SinkQueueConfig {
+	var result SinkQueueConfig
+	if topic, ok := config[TopicKey].(string); ok {
+		result.Topic = topic
+	}
+	return &result
+}
+
+func (c *SinkQueueConfig) ToConfigMap() ConfigMap {
+	return ConfigMap{
+		TopicKey: c.Topic,
+	}
+}
+
+type ConfigMap map[string]interface{}
+
+type TubeFactory interface {
+	// NewSourceTube returns a new channel that can be used to receive events
+	// The channel would be closed when the context is done
+	NewSourceTube(ctx context.Context, config ConfigMap) (<-chan Record, error)
+	// NewSinkTube returns a new channel that can be used to sink events
+	// The event.Commit() would be invoked after the event is sunk successfully
+	// The caller should close the channel when it is done
+	NewSinkTube(ctx context.Context, config ConfigMap) (chan<- Record, error)
+}
+
+type RecordImpl struct {
+	payload    []byte
+	commitFunc func()
+}
+
+func NewRecordImpl(payload []byte, ackFunc func()) *RecordImpl {
+	return &RecordImpl{
+		payload:    payload,
+		commitFunc: ackFunc,
+	}
+}
+
+func (e *RecordImpl) GetPayload() []byte {
+	return e.payload
+}
+
+func (e *RecordImpl) Commit() {
+	e.commitFunc()
+}

lib/memory_queue.go lib/contube/memory_tube.golib/memory_queue.go renamed to lib/contube/memory_tube.go

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package lib
+package contube
 
 import (
 	"context"
@@ -24,7 +24,7 @@ import (
 )
 
 type queue struct {
-	c      chan Event
+	c      chan Record
 	refCnt int32
 }
 
@@ -34,14 +34,14 @@ type MemoryQueueFactory struct {
 	queues map[string]*queue
 }
 
-func NewMemoryQueueFactory(ctx context.Context) EventQueueFactory {
+func NewMemoryQueueFactory(ctx context.Context) TubeFactory {
 	return &MemoryQueueFactory{
 		ctx:    ctx,
 		queues: make(map[string]*queue),
 	}
 }
 
-func (f *MemoryQueueFactory) getOrCreateChan(name string) chan Event {
+func (f *MemoryQueueFactory) getOrCreateChan(name string) chan Record {
 	f.mu.Lock()
 	defer f.mu.Unlock()
 	defer func() {
@@ -53,7 +53,7 @@ func (f *MemoryQueueFactory) getOrCreateChan(name string) chan Event {
 		atomic.AddInt32(&q.refCnt, 1)
 		return q.c
 	}
-	c := make(chan Event, 100)
+	c := make(chan Record, 100)
 	f.queues[name] = &queue{
 		c:      c,
 		refCnt: 1,
@@ -77,8 +77,9 @@ func (f *MemoryQueueFactory) release(name string) {
 		"name", name)
 }
 
-func (f *MemoryQueueFactory) NewSourceChan(ctx context.Context, config *SourceQueueConfig) (<-chan Event, error) {
-	result := make(chan Event)
+func (f *MemoryQueueFactory) NewSourceTube(ctx context.Context, configMap ConfigMap) (<-chan Record, error) {
+	config := NewSourceQueueConfig(configMap)
+	result := make(chan Record)
 	for _, topic := range config.Topics {
 		t := topic
 		go func() {
@@ -101,9 +102,10 @@ func (f *MemoryQueueFactory) NewSourceChan(ctx context.Context, config *SourceQu
 	return result, nil
 }
 
-func (f *MemoryQueueFactory) NewSinkChan(ctx context.Context, config *SinkQueueConfig) (chan<- Event, error) {
+func (f *MemoryQueueFactory) NewSinkTube(ctx context.Context, configMap ConfigMap) (chan<- Record, error) {
+	config := NewSinkQueueConfig(configMap)
 	c := f.getOrCreateChan(config.Topic)
-	wrapperC := make(chan Event)
+	wrapperC := make(chan Record)
 	go func() {
 		defer f.release(config.Topic)
 		for {
@@ -114,7 +116,7 @@ func (f *MemoryQueueFactory) NewSinkChan(ctx context.Context, config *SinkQueueC
 				if !ok {
 					return
 				}
-				event.Ack()
+				event.Commit()
 				c <- event
 			}
 		}