feat(orchestration): add limites on parallel activities

formancehq · Jul 17, 2024 · bb87c5c · bb87c5c
1 parent edbf248
commit bb87c5c
Show file tree

Hide file tree

Showing 8 changed files with 20 additions and 8 deletions.
diff --git a/ee/orchestration/cmd/root.go b/ee/orchestration/cmd/root.go
@@ -51,6 +51,7 @@ const (
 	temporalSSLClientCertFlag    = "temporal-ssl-client-cert"
 	temporalTaskQueueFlag        = "temporal-task-queue"
 	temporalInitSearchAttributes = "temporal-init-search-attributes"
+	temporalMaxActivitiesPerSecond = "temporal-max-activities-per-second"
 	topicsFlag                   = "topics"
 	listenFlag                   = "listen"
 	workerFlag                   = "worker"

diff --git a/ee/orchestration/cmd/serve.go b/ee/orchestration/cmd/serve.go
@@ -70,6 +70,7 @@ func newServeCommand() *cobra.Command {
 
 	cmd.Flags().Bool(workerFlag, false, "Enable worker mode")
 	cmd.Flags().String(listenFlag, ":8080", "Listening address")
+	cmd.Flags().Float64(temporalMaxActivitiesPerSecond, 10, "Maximum number of parallel activities")
 	service.BindFlags(cmd)
 
 	return cmd

diff --git a/ee/orchestration/cmd/worker.go b/ee/orchestration/cmd/worker.go
@@ -1,6 +1,7 @@
 package cmd
 
 import (
+	"go.temporal.io/sdk/worker"
 	"net/http"
 
 	"github.com/formancehq/orchestration/internal/triggers"
@@ -27,7 +28,9 @@ func stackClientModule() fx.Option {
 func workerOptions() fx.Option {
 	return fx.Options(
 		stackClientModule(),
-		temporalworker.NewWorkerModule(viper.GetString(temporalTaskQueueFlag)),
+		temporalworker.NewWorkerModule(viper.GetString(temporalTaskQueueFlag), worker.Options{
+			TaskQueueActivitiesPerSecond: viper.GetFloat64(temporalMaxActivitiesPerSecond),
+		}),
 		triggers.NewListenerModule(
 			viper.GetString(stackFlag),
 			viper.GetString(temporalTaskQueueFlag),

diff --git a/ee/orchestration/internal/api/v1/main_test.go b/ee/orchestration/internal/api/v1/main_test.go
@@ -2,6 +2,7 @@ package v1
 
 import (
 	"context"
+	"go.temporal.io/sdk/worker"
 	"log"
 	"net/http"
 	"os"
@@ -53,6 +54,7 @@ func test(t *testing.T, fn func(router *chi.Mux, backend api.Backend, db *bun.DB
 		[]temporalworker.DefinitionSet{
 			workflow.NewActivities(publish.NoOpPublisher, db).DefinitionSet(),
 		},
+		worker.Options{},
 	)
 
 	require.NoError(t, worker.Start())

diff --git a/ee/orchestration/internal/api/v2/main_test.go b/ee/orchestration/internal/api/v2/main_test.go
@@ -2,6 +2,7 @@ package v2
 
 import (
 	"context"
+	"go.temporal.io/sdk/worker"
 	"log"
 	"net/http"
 	"os"
@@ -11,7 +12,7 @@ import (
 	"github.com/formancehq/orchestration/internal/workflow/stages"
 	"github.com/formancehq/stack/libs/go-libs/logging"
 	"github.com/formancehq/stack/libs/go-libs/publish"
-	chi "github.com/go-chi/chi/v5"
+	"github.com/go-chi/chi/v5"
 	"github.com/google/uuid"
 	"go.temporal.io/sdk/testsuite"
 
@@ -53,6 +54,7 @@ func test(t *testing.T, fn func(router *chi.Mux, backend api.Backend, db *bun.DB
 		[]temporalworker.DefinitionSet{
 			workflow.NewActivities(publish.NoOpPublisher, db).DefinitionSet(),
 		},
+		worker.Options{},
 	)
 	require.NoError(t, worker.Start())
 	t.Cleanup(worker.Stop)

diff --git a/ee/orchestration/internal/temporalworker/module.go b/ee/orchestration/internal/temporalworker/module.go
@@ -30,10 +30,9 @@ func (d DefinitionSet) Append(definition Definition) DefinitionSet {
 	return d
 }
 
-func New(logger logging.Logger, c client.Client, taskQueue string, workflows, activities []DefinitionSet) worker.Worker {
-	worker := worker.New(c, taskQueue, worker.Options{
-		BackgroundActivityContext: logging.ContextWithLogger(context.Background(), logger),
-	})
+func New(logger logging.Logger, c client.Client, taskQueue string, workflows, activities []DefinitionSet, options worker.Options) worker.Worker {
+	options.BackgroundActivityContext = logging.ContextWithLogger(context.Background(), logger)
+	worker := worker.New(c, taskQueue, options)
 
 	for _, set := range workflows {
 		for _, workflow := range set {
@@ -54,11 +53,11 @@ func New(logger logging.Logger, c client.Client, taskQueue string, workflows, ac
 	return worker
 }
 
-func NewWorkerModule(taskQueue string) fx.Option {
+func NewWorkerModule(taskQueue string, options worker.Options) fx.Option {
 	return fx.Options(
 		fx.Provide(
 			fx.Annotate(func(logger logging.Logger, c client.Client, workflows, activities []DefinitionSet) worker.Worker {
-				return New(logger, c, taskQueue, workflows, activities)
+				return New(logger, c, taskQueue, workflows, activities, options)
 			}, fx.ParamTags(``, ``, `group:"workflows"`, `group:"activities"`)),
 		),
 		fx.Invoke(func(lc fx.Lifecycle, w worker.Worker) {

diff --git a/ee/orchestration/internal/triggers/workflow_trigger_test.go b/ee/orchestration/internal/triggers/workflow_trigger_test.go
@@ -1,6 +1,7 @@
 package triggers
 
 import (
+	worker "go.temporal.io/sdk/worker"
 	"testing"
 	"time"
 
@@ -47,6 +48,7 @@ func TestWorkflow(t *testing.T) {
 			workflow.NewActivities(publish.NoOpPublisher, db).DefinitionSet(),
 			NewActivities(db, workflowManager, NewDefaultExpressionEvaluator(), publish.NoOpPublisher).DefinitionSet(),
 		},
+		worker.Options{},
 	)
 	require.NoError(t, worker.Start())
 	t.Cleanup(worker.Stop)

diff --git a/ee/orchestration/internal/workflow/manager_test.go b/ee/orchestration/internal/workflow/manager_test.go
@@ -1,6 +1,7 @@
 package workflow
 
 import (
+	worker "go.temporal.io/sdk/worker"
 	"testing"
 	"time"
 
@@ -42,6 +43,7 @@ func TestConfig(t *testing.T) {
 		[]temporalworker.DefinitionSet{
 			NewActivities(publish.NoOpPublisher, db).DefinitionSet(),
 		},
+		worker.Options{},
 	)
 	require.NoError(t, worker.Start())
 	t.Cleanup(worker.Stop)