fix: use LA pool for the local activities execution

aggregatedpool/handler.go

@@ -146,7 +146,7 @@ func (wp *Workflow) handleMessage(msg *internal.Message) error {
 
 	case *internal.ExecuteLocalActivity:
 		wp.log.Debug("local activity request", zap.Uint64("ID", msg.ID))
-		params := command.LocalActivityParams(wp.env, NewLocalActivityFn(msg.Header, wp.codec, wp.pool, wp.log).execute, msg.Payloads, msg.Header)
+		params := command.LocalActivityParams(wp.env, wp.la, msg.Payloads, msg.Header)
 		activityID := wp.env.ExecuteLocalActivity(params, wp.createLocalActivityCallback(msg.ID))
 		wp.canceller.Register(msg.ID, func() error {
 			wp.log.Debug("registering local activity canceller", zap.String("activityID", activityID.String()))

aggregatedpool/local_activity.go

@@ -18,23 +18,21 @@ import (
 )
 
 type LocalActivityFn struct {
-	header *commonpb.Header
-	codec  common.Codec
-	pool   common.Pool
-	log    *zap.Logger
-	seqID  uint64
+	codec common.Codec
+	pool  common.Pool
+	log   *zap.Logger
+	seqID uint64
 }
 
-func NewLocalActivityFn(header *commonpb.Header, codec common.Codec, pool common.Pool, log *zap.Logger) *LocalActivityFn {
+func NewLocalActivityFn(codec common.Codec, pool common.Pool, log *zap.Logger) *LocalActivityFn {
 	return &LocalActivityFn{
-		header: header,
-		codec:  codec,
-		pool:   pool,
-		log:    log,
+		codec: codec,
+		pool:  pool,
+		log:   log,
 	}
 }
 
-func (la *LocalActivityFn) execute(ctx context.Context, args *commonpb.Payloads) (*commonpb.Payloads, error) {
+func (la *LocalActivityFn) ExecuteLA(ctx context.Context, hdr *commonpb.Header, args *commonpb.Payloads) (*commonpb.Payloads, error) {
 	const op = errors.Op("activity_pool_execute_activity")
 
 	var info = tActivity.GetInfo(ctx)
@@ -53,7 +51,7 @@ func (la *LocalActivityFn) execute(ctx context.Context, args *commonpb.Payloads)
 			Info: info,
 		},
 		Payloads: args,
-		Header:   la.header,
+		Header:   hdr,
 	}
 
 	la.log.Debug("executing local activity fn", zap.Uint64("ID", msg.ID), zap.String("task-queue", info.TaskQueue), zap.String("la ID", info.ActivityID))

aggregatedpool/workflow.go

@@ -1,6 +1,7 @@
 package aggregatedpool
 
 import (
+	"context"
 	"fmt"
 	"sync"
 	"sync/atomic"
@@ -32,6 +33,7 @@ import (
 */
 
 type Callback func() error
+type LaFn func(ctx context.Context, hdr *commonpb.Header, args *commonpb.Payloads) (*commonpb.Payloads, error)
 
 // seqID is global sequence ID
 var seqID uint64 //nolint:gochecknoglobals
@@ -45,6 +47,9 @@ type Workflow struct {
 	pool  common.Pool
 	rrID  string
 
+	// LocalActivityFn
+	la LaFn
+
 	env          bindings.WorkflowEnvironment
 	header       *commonpb.Header
 	mq           *queue.MessageQueue
@@ -67,10 +72,12 @@ type Workflow struct {
 	pldPool *sync.Pool
 }
 
-func NewWorkflowDefinition(codec common.Codec, pool common.Pool, log *zap.Logger) *Workflow {
+// NewWorkflowDefinition ... WorkflowDefinition Constructor
+func NewWorkflowDefinition(codec common.Codec, la LaFn, pool common.Pool, log *zap.Logger) *Workflow {
 	return &Workflow{
 		rrID:  uuid.NewString(),
 		log:   log,
+		la:    la,
 		codec: codec,
 		pool:  pool,
 		pldPool: &sync.Pool{
@@ -83,9 +90,12 @@ func NewWorkflowDefinition(codec common.Codec, pool common.Pool, log *zap.Logger
 
 // NewWorkflowDefinition ... Workflow should match the WorkflowDefinitionFactory interface (sdk-go/internal/internal_worker.go:463, RegisterWorkflowWithOptions func)
 // DO NOT USE THIS FUNCTION DIRECTLY!!!!
+// This function called after the constructor above, it is safe to assign fields like that
 func (wp *Workflow) NewWorkflowDefinition() bindings.WorkflowDefinition {
 	return &Workflow{
 		rrID: uuid.NewString(),
+		// LocalActivity
+		la: wp.la,
 		// updates logic
 		updateCompleteCb: make(map[string]func(res *internal.Message)),
 		updateValidateCb: make(map[string]func(res *internal.Message)),
@@ -143,7 +153,7 @@ func (wp *Workflow) Execute(env bindings.WorkflowEnvironment, header *commonpb.H
 	)
 }
 
-// OnWorkflowTaskStarted is called for each non timed out startWorkflowTask event.
+// OnWorkflowTaskStarted is called for each non-timed out startWorkflowTask event.
 // Executed after all history events since the previous commands are applied to WorkflowDefinition
 // Application level code must be executed from this function only.
 // Execute call as well as callbacks called from WorkflowEnvironment functions can only schedule callbacks

internal.go

@@ -28,7 +28,10 @@ func (p *Plugin) initPool() error {
 	dc := data_converter.NewDataConverter(converter.GetDefaultDataConverter())
 	codec := proto.NewCodec(p.log, dc)
 
+	// LA + A definitions
 	actDef := aggregatedpool.NewActivityDefinition(codec, ap, p.log)
+	laDef := aggregatedpool.NewLocalActivityFn(codec, ap, p.log)
+	// ------------------
 
 	// ---------- WORKFLOW POOL -------------
 	wp, err := p.server.NewPool(
@@ -48,7 +51,7 @@ func (p *Plugin) initPool() error {
 		return err
 	}
 
-	wfDef := aggregatedpool.NewWorkflowDefinition(codec, wp, p.log)
+	wfDef := aggregatedpool.NewWorkflowDefinition(codec, laDef.ExecuteLA, wp, p.log)
 
 	// get worker information
 	wi, err := WorkerInfo(codec, wp, p.rrVersion)

internal/protocol.go

@@ -347,11 +347,12 @@ func (cmd ExecuteLocalActivity) LocalActivityParams(env bindings.WorkflowEnviron
 		truTemOptions.RetryPolicy = rp
 	}
 
+	// TODO: should be careful here: header + inputArgs header are pointers and might be changed independently which will cause race
 	params := bindings.ExecuteLocalActivityParams{
 		ExecuteLocalActivityOptions: truTemOptions,
 		ActivityFn:                  fn,
 		ActivityType:                cmd.Name,
-		InputArgs:                   []any{payloads},
+		InputArgs:                   []any{header, payloads},
 		WorkflowInfo:                env.WorkflowInfo(),
 		ScheduledTime:               time.Now(),
 		Header:                      header,