feat: state store (#153)

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

.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)
+}

fs/func_ctx_impl_test.go

@@ -0,0 +1,29 @@
+/*
+ * 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/stretchr/testify/assert"
+	"testing"
+)
+
+func TestFuncCtx_NilStore(t *testing.T) {
+	f := NewFuncCtxImpl(nil)
+	assert.ErrorIs(t, f.PutState("key", []byte("value")), ErrStateStoreNotLoaded)
+	_, err := f.GetState("key")
+	assert.ErrorIs(t, err, ErrStateStoreNotLoaded)
+}

fs/instance_impl.go

@@ -28,8 +28,8 @@ import (
 )
 
 type FunctionInstanceImpl struct {
-	api.FunctionInstance
 	ctx           context.Context
+	funcCtx       api.FunctionContext
 	cancelFunc    context.CancelFunc
 	definition    *model.Function
 	sourceFactory contube.SourceTubeFactory
@@ -47,18 +47,21 @@ const (
 	CtxKeyInstanceIndex CtxKey = "instance-index"
 )
 
-type DefaultInstanceFactory struct{}
+type DefaultInstanceFactory struct {
+	api.FunctionInstanceFactory
+}
 
 func NewDefaultInstanceFactory() api.FunctionInstanceFactory {
 	return &DefaultInstanceFactory{}
 }
 
-func (f *DefaultInstanceFactory) NewFunctionInstance(definition *model.Function, sourceFactory contube.SourceTubeFactory, sinkFactory contube.SinkTubeFactory, index int32, logger *slog.Logger) api.FunctionInstance {
+func (f *DefaultInstanceFactory) NewFunctionInstance(definition *model.Function, funcCtx api.FunctionContext, sourceFactory contube.SourceTubeFactory, sinkFactory contube.SinkTubeFactory, index int32, logger *slog.Logger) api.FunctionInstance {
 	ctx, cancelFunc := context.WithCancel(context.Background())
 	ctx = context.WithValue(ctx, CtxKeyFunctionName, definition.Name)
 	ctx = context.WithValue(ctx, CtxKeyInstanceIndex, index)
 	return &FunctionInstanceImpl{
 		ctx:           ctx,
+		funcCtx:       funcCtx,
 		cancelFunc:    cancelFunc,
 		definition:    definition,
 		sourceFactory: sourceFactory,
@@ -136,6 +139,10 @@ func (instance *FunctionInstanceImpl) Context() context.Context {
 	return instance.ctx
 }
 
+func (instance *FunctionInstanceImpl) FunctionContext() api.FunctionContext {
+	return instance.funcCtx
+}
+
 func (instance *FunctionInstanceImpl) Definition() *model.Function {
 	return instance.definition
 }

fs/instance_impl_test.go

@@ -28,7 +28,7 @@ func TestFunctionInstanceContextSetting(t *testing.T) {
 		Name: "test-function",
 	}
 	index := int32(1)
-	instance := defaultInstanceFactory.NewFunctionInstance(definition, nil, nil, index, slog.Default())
+	instance := defaultInstanceFactory.NewFunctionInstance(definition, nil, nil, nil, index, slog.Default())
 
 	if instance == nil {
 		t.Error("FunctionInstance should not be nil")

fs/manager.go

@@ -23,6 +23,7 @@ import (
 	"github.com/functionstream/function-stream/fs/api"
 	"github.com/functionstream/function-stream/fs/contube"
 	"github.com/functionstream/function-stream/fs/runtime/wazero"
+	"github.com/functionstream/function-stream/fs/statestore"
 	"github.com/pkg/errors"
 	"log/slog"
 	"math/rand"
@@ -38,9 +39,11 @@ type FunctionManager struct {
 }
 
 type managerOptions struct {
-	tubeFactoryMap    map[string]contube.TubeFactory
-	runtimeFactoryMap map[string]api.FunctionRuntimeFactory
-	instanceFactory   api.FunctionInstanceFactory
+	tubeFactoryMap           map[string]contube.TubeFactory
+	runtimeFactoryMap        map[string]api.FunctionRuntimeFactory
+	instanceFactory          api.FunctionInstanceFactory
+	stateStore               api.StateStore
+	dontUseDefaultStateStore bool
 }
 
 type ManagerOption interface {
@@ -82,6 +85,14 @@ func WithInstanceFactory(factory api.FunctionInstanceFactory) ManagerOption {
 	})
 }
 
+func WithStateStore(store api.StateStore) ManagerOption {
+	return managerOptionFunc(func(c *managerOptions) (*managerOptions, error) {
+		c.dontUseDefaultStateStore = true
+		c.stateStore = store
+		return c, nil
+	})
+}
+
 func NewFunctionManager(opts ...ManagerOption) (*FunctionManager, error) {
 	options := &managerOptions{
 		tubeFactoryMap:    make(map[string]contube.TubeFactory),
@@ -96,6 +107,13 @@ func NewFunctionManager(opts ...ManagerOption) (*FunctionManager, error) {
 			return nil, err
 		}
 	}
+	if !options.dontUseDefaultStateStore {
+		var err error
+		options.stateStore, err = statestore.NewTmpPebbleStateStore()
+		if err != nil {
+			return nil, errors.Wrap(err, "failed to create default state store")
+		}
+	}
 	log := slog.With()
 	loadedRuntimeFact := make([]string, 0, len(options.runtimeFactoryMap))
 	for k := range options.runtimeFactoryMap {
@@ -145,6 +163,10 @@ func (fm *FunctionManager) getRuntimeFactory(t string) (api.FunctionRuntimeFacto
 	return factory, nil
 }
 
+func (fm *FunctionManager) createFuncCtx(f *model.Function) api.FunctionContext {
+	return NewFuncCtxImpl(fm.options.stateStore)
+}
+
 func (fm *FunctionManager) StartFunction(f *model.Function) error {
 	fm.functionsLock.Lock()
 	if _, exist := fm.functions[f.Name]; exist {
@@ -156,6 +178,7 @@ func (fm *FunctionManager) StartFunction(f *model.Function) error {
 	if f.Replicas <= 0 {
 		return errors.New("replicas should be greater than 0")
 	}
+	funcCtx := fm.createFuncCtx(f)
 	for i := int32(0); i < f.Replicas; i++ {
 		sourceFactory, err := fm.getTubeFactory(f.Source)
 		if err != nil {
@@ -167,7 +190,7 @@ func (fm *FunctionManager) StartFunction(f *model.Function) error {
 		}
 		runtimeType := fm.getRuntimeType(f.Runtime)
 
-		instance := fm.options.instanceFactory.NewFunctionInstance(f, sourceFactory, sinkFactory, i, slog.With(
+		instance := fm.options.instanceFactory.NewFunctionInstance(f, funcCtx, sourceFactory, sinkFactory, i, slog.With(
 			slog.String("name", f.Name),
 			slog.Int("index", int(i)),
 			slog.String("runtime", runtimeType),
@@ -181,7 +204,7 @@ func (fm *FunctionManager) StartFunction(f *model.Function) error {
 		}
 		go instance.Run(runtimeFactory)
 		select {
-		case <-instance.WaitForReady():
+		case err := <-instance.WaitForReady():
 			if err != nil {
 				fm.log.ErrorContext(instance.Context(), "Error starting function instance", slog.Any("error", err.Error()))
 				instance.Stop()
@@ -241,3 +264,24 @@ func (fm *FunctionManager) ConsumeEvent(name string) (contube.Record, error) {
 	}
 	return <-c, nil
 }
+
+// GetStateStore returns the state store used by the function manager
+// Return nil if no state store is configured
+func (fm *FunctionManager) GetStateStore() api.StateStore {
+	return fm.options.stateStore
+}
+
+func (fm *FunctionManager) Close() error {
+	fm.functionsLock.Lock()
+	defer fm.functionsLock.Unlock()
+	for _, instances := range fm.functions {
+		for _, instance := range instances {
+			instance.Stop()
+		}
+	}
+	err := fm.options.stateStore.Close()
+	if err != nil {
+		return err
+	}
+	return nil
+}