feat: state store (#153)

Signed-off-by: Zike Yang <[email protected]>

Author: RobertIndie

.github/workflows/ci.yaml

@@ -57,7 +57,6 @@ jobs:
       - run: make build_all
       - name: Wait for Pulsar service
         run: until curl http://localhost:8080/metrics > /dev/null 2>&1 ; do sleep 1; done
-      - run: nohup bin/function-stream server > nohup.out &
       - run: make test
       - name: Collect Docker Compose logs
         if: failure()

benchmark/bench_test.go

@@ -35,7 +35,7 @@ import (
 )
 
 func BenchmarkStressForBasicFunc(b *testing.B) {
-	s, err := server.NewServer(server.LoadConfigFromEnv())
+	s, err := server.NewDefaultServer()
 	if err != nil {
 		b.Fatal(err)
 	}
@@ -110,15 +110,7 @@ func BenchmarkStressForBasicFunc(b *testing.B) {
 func BenchmarkStressForBasicFuncWithMemoryQueue(b *testing.B) {
 	memoryQueueFactory := contube.NewMemoryQueueFactory(context.Background())
 
-	svrConf := &common.Config{
-		ListenAddr: common.DefaultAddr,
-	}
-
-	fm, err := fs.NewFunctionManager(fs.WithDefaultTubeFactory(memoryQueueFactory))
-	if err != nil {
-		b.Fatal(err)
-	}
-	s, err := server.NewServer(svrConf, server.WithFunctionManager(fm))
+	s, err := server.NewServer(server.WithFunctionManager(fs.WithDefaultTubeFactory(memoryQueueFactory)))
 	if err != nil {
 		b.Fatal(err)
 	}

cmd/server/cmd.go

@@ -35,7 +35,7 @@ var (
 
 func exec(*cobra.Command, []string) {
 	common.RunProcess(func() (io.Closer, error) {
-		s, err := server.NewServer(server.LoadConfigFromEnv())
+		s, err := server.NewServer()
 		if err != nil {
 			return nil, err
 		}

cmd/standalone/cmd.go

@@ -35,7 +35,7 @@ var (
 
 func exec(*cobra.Command, []string) {
 	common.RunProcess(func() (io.Closer, error) {
-		s, err := server.NewServer(server.LoadStandaloneConfigFromEnv())
+		s, err := server.NewServerWithConfig(server.LoadStandaloneConfigFromEnv())
 		if err != nil {
 			return nil, err
 		}

fs/api/func_ctx.go

@@ -0,0 +1,22 @@
+/*
+ * 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 api
+
+type FunctionContext interface {
+	PutState(key string, value []byte) error
+	GetState(key string) ([]byte, error)
+}

fs/api/instance.go

@@ -25,6 +25,7 @@ import (
 
 type FunctionInstance interface {
 	Context() context.Context
+	FunctionContext() FunctionContext
 	Definition() *model.Function
 	Index() int32
 	Stop()
@@ -34,5 +35,5 @@ type FunctionInstance interface {
 }
 
 type FunctionInstanceFactory interface {
-	NewFunctionInstance(f *model.Function, sourceFactory contube.SourceTubeFactory, sinkFactory contube.SinkTubeFactory, i int32, logger *slog.Logger) FunctionInstance
+	NewFunctionInstance(f *model.Function, funcCtx FunctionContext, sourceFactory contube.SourceTubeFactory, sinkFactory contube.SinkTubeFactory, i int32, logger *slog.Logger) FunctionInstance
 }

fs/api/statestore.go

@@ -0,0 +1,27 @@
+/*
+ * 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 api
+
+import "github.com/pkg/errors"
+
+var ErrNotFound = errors.New("key not found")
+
+type StateStore interface {
+	PutState(key string, value []byte) error
+	GetState(key string) ([]byte, error)
+	Close() error
+}

fs/contube/http.go

@@ -56,7 +56,7 @@ type HttpTubeFactory struct {
 	endpoints map[string]*endpointHandler
 }
 
-func NewHttpTubeFactory(ctx context.Context) TubeFactory {
+func NewHttpTubeFactory(ctx context.Context) *HttpTubeFactory {
 	return &HttpTubeFactory{
 		ctx:       ctx,
 		endpoints: make(map[string]*endpointHandler),

fs/contube/http_test.go

@@ -24,7 +24,7 @@ import (
 
 func TestHttpTubeHandleRecord(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
-	f := NewHttpTubeFactory(ctx).(*HttpTubeFactory)
+	f := NewHttpTubeFactory(ctx)
 
 	endpoint := "test"
 	err := f.Handle(ctx, "test", []byte("test"))
@@ -51,7 +51,7 @@ func TestHttpTubeHandleRecord(t *testing.T) {
 }
 
 func TestHttpTubeSinkTubeNotImplement(t *testing.T) {
-	f := NewHttpTubeFactory(context.Background()).(*HttpTubeFactory)
+	f := NewHttpTubeFactory(context.Background())
 	_, err := f.NewSinkTube(context.Background(), make(ConfigMap))
 	assert.ErrorIs(t, err, ErrSinkTubeNotImplemented)
 }

fs/func_ctx_impl.go

@@ -0,0 +1,55 @@
+/*
+ * 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 fs
+
+import (
+	"github.com/functionstream/function-stream/fs/api"
+	"github.com/pkg/errors"
+)
+
+var ErrStateStoreNotLoaded = errors.New("state store not loaded")
+
+type FuncCtxImpl struct {
+	api.FunctionContext
+	store        api.StateStore
+	putStateFunc func(key string, value []byte) error
+	getStateFunc func(key string) ([]byte, error)
+}
+
+func NewFuncCtxImpl(store api.StateStore) *FuncCtxImpl {
+	putStateFunc := func(key string, value []byte) error {
+		return ErrStateStoreNotLoaded
+	}
+	getStateFunc := func(key string) ([]byte, error) {
+		return nil, ErrStateStoreNotLoaded
+	}
+	if store != nil {
+		putStateFunc = store.PutState
+		getStateFunc = store.GetState
+	}
+	return &FuncCtxImpl{store: store,
+		putStateFunc: putStateFunc,
+		getStateFunc: getStateFunc}
+}
+
+func (f *FuncCtxImpl) PutState(key string, value []byte) error {
+	return f.putStateFunc(key, value)
+}
+
+func (f *FuncCtxImpl) GetState(key string) ([]byte, error) {
+	return f.getStateFunc(key)
+}