Add Context Propagation support for Local Activities (#986)

Author: tfcace

internal/internal_activity.go

@@ -88,6 +88,7 @@ type (
 		DataConverter DataConverter
 		Attempt       int32
 		ScheduledTime time.Time
+		Header        *shared.Header
 	}
 
 	// asyncActivityClient for requesting activity execution

internal/internal_event_handlers.go

@@ -132,6 +132,7 @@ type (
 		attempt      int32 // attempt starting from 0
 		retryPolicy  *RetryPolicy
 		expireTime   time.Time
+		header       *shared.Header
 	}
 
 	localActivityMarkerData struct {
@@ -503,6 +504,7 @@ func newLocalActivityTask(params executeLocalActivityParams, callback laResultHa
 		callback:    callback,
 		retryPolicy: params.RetryPolicy,
 		attempt:     params.Attempt,
+		header:      params.Header,
 	}
 
 	if params.RetryPolicy != nil && params.RetryPolicy.ExpirationInterval > 0 {

internal/internal_task_pollers.go

@@ -482,6 +482,19 @@ func (lath *localActivityTaskHandler) executeLocalActivityTask(task *localActivi
 		attempt:           task.attempt,
 	})
 
+	// propagate context information into the local activity activity context from the headers
+	for _, ctxProp := range lath.contextPropagators {
+		var err error
+		if ctx, err = ctxProp.Extract(ctx, NewHeaderReader(task.header)); err != nil {
+			result = &localActivityResult{
+				task:   task,
+				result: nil,
+				err:    fmt.Errorf("unable to propagate context %v", err),
+			}
+			return result
+		}
+	}
+
 	// panic handler
 	defer func() {
 		if p := recover(); p != nil {

internal/workflow.go

@@ -518,13 +518,13 @@ func ExecuteLocalActivity(ctx Context, activity interface{}, args ...interface{}
 }
 
 func (wc *workflowEnvironmentInterceptor) ExecuteLocalActivity(ctx Context, activityType string, args ...interface{}) Future {
+	header := getHeadersFromContext(ctx)
 	activityFn := ctx.Value(localActivityFnContextKey)
 	if activityFn == nil {
 		panic("ExecuteLocalActivity: Expected context key " + localActivityFnContextKey + " is missing")
 	}
 
 	future, settable := newDecodeFuture(ctx, activityFn)
-
 	if err := validateFunctionArgs(activityFn, args, false); err != nil {
 		settable.Set(nil, err)
 		return future
@@ -542,6 +542,7 @@ func (wc *workflowEnvironmentInterceptor) ExecuteLocalActivity(ctx Context, acti
 		WorkflowInfo:         GetWorkflowInfo(ctx),
 		DataConverter:        getDataConverterFromWorkflowContext(ctx),
 		ScheduledTime:        Now(ctx), // initial scheduled time
+		Header:               header,
 	}
 
 	Go(ctx, func(ctx Context) {

test/activity_test.go

@@ -87,19 +87,12 @@ func (a *Activities) fail(_ context.Context) error {
 	return errFailOnPurpose
 }
 
-func (a *Activities) InspectActivityInfo(ctx context.Context, domain, taskList, wfType string) error {
-	a.append("inspectActivityInfo")
-	info := activity.GetInfo(ctx)
-	if info.WorkflowDomain != domain {
-		return fmt.Errorf("expected domainName %v but got %v", domain, info.WorkflowDomain)
-	}
-	if info.WorkflowType == nil || info.WorkflowType.Name != wfType {
-		return fmt.Errorf("expected workflowType %v but got %v", wfType, info.WorkflowType)
-	}
-	if info.TaskList != taskList {
-		return fmt.Errorf("expected taskList %v but got %v", taskList, info.TaskList)
+func (a *Activities) DuplicateStringInContext(ctx context.Context) (string, error) {
+	originalString := ctx.Value(contextKey(testContextKey))
+	if originalString == nil {
+		return "", fmt.Errorf("context did not propagate to activity")
 	}
-	return nil
+	return strings.Repeat(originalString.(string), 2), nil
 }
 
 func (a *Activities) append(name string) {
@@ -140,6 +133,21 @@ func (a *Activities) GetMemoAndSearchAttr(_ context.Context, memo, searchAttr st
 	return memo + ", " + searchAttr, nil
 }
 
+func (a *Activities) InspectActivityInfo(ctx context.Context, domain, taskList, wfType string) error {
+	a.append("inspectActivityInfo")
+	info := activity.GetInfo(ctx)
+	if info.WorkflowDomain != domain {
+		return fmt.Errorf("expected domainName %v but got %v", domain, info.WorkflowDomain)
+	}
+	if info.WorkflowType == nil || info.WorkflowType.Name != wfType {
+		return fmt.Errorf("expected workflowType %v but got %v", wfType, info.WorkflowType)
+	}
+	if info.TaskList != taskList {
+		return fmt.Errorf("expected taskList %v but got %v", taskList, info.TaskList)
+	}
+	return nil
+}
+
 func (a *Activities) register(worker worker.Worker) {
 	// Kept to verify backward compatibility of activity registration.
 	activity.RegisterWithOptions(a, activity.RegisterOptions{DisableAlreadyRegisteredCheck: true})

test/integration_test.go

@@ -61,6 +61,7 @@ const (
 	ctxTimeout                 = 15 * time.Second
 	domainName                 = "integration-test-domain"
 	domainCacheRefreshInterval = 20 * time.Second
+	testContextKey			   = "test-context-key"
 )
 
 func TestIntegrationSuite(t *testing.T) {
@@ -97,7 +98,10 @@ func (ts *IntegrationTestSuite) SetupSuite() {
 	rpcClient, err := newRPCClient(ts.config.ServiceName, ts.config.ServiceAddr)
 	ts.NoError(err)
 	ts.rpcClient = rpcClient
-	ts.libClient = client.NewClient(ts.rpcClient.Interface, domainName, &client.Options{})
+	ts.libClient = client.NewClient(ts.rpcClient.Interface, domainName,
+		&client.Options{
+			ContextPropagators: []workflow.ContextPropagator{NewStringMapPropagator([]string{testContextKey})},
+		})
 	ts.registerDomain()
 }
 
@@ -138,12 +142,14 @@ func (ts *IntegrationTestSuite) SetupTest() {
 	ts.worker = worker.New(ts.rpcClient.Interface, domainName, ts.taskListName, worker.Options{
 		DisableStickyExecution: ts.config.IsStickyOff,
 		Logger:                 logger,
+		ContextPropagators:     []workflow.ContextPropagator{NewStringMapPropagator([]string{testContextKey})},
 	})
 	ts.tracer = newtracingInterceptorFactory()
 	options := worker.Options{
 		DisableStickyExecution:            ts.config.IsStickyOff,
 		Logger:                            logger,
 		WorkflowInterceptorChainFactories: []interceptors.WorkflowInterceptorFactory{ts.tracer},
+		ContextPropagators:     []workflow.ContextPropagator{NewStringMapPropagator([]string{testContextKey})},
 	}
 	ts.worker = worker.New(ts.rpcClient.Interface, domainName, ts.taskListName, options)
 	ts.registerWorkflowsAndActivities(ts.worker)
@@ -405,6 +411,13 @@ func (ts *IntegrationTestSuite) TestActivityCancelRepro() {
 	ts.EqualValues(expected, ts.activities.invoked())
 }
 
+func (ts *IntegrationTestSuite) TestWorkflowWithLocalActivityCtxPropagation() {
+	var expected string
+	err := ts.executeWorkflow("test-wf-local-activity-ctx-prop", ts.workflows.WorkflowWithLocalActivityCtxPropagation, &expected)
+	ts.NoError(err)
+	ts.EqualValues(expected, "test-data-in-contexttest-data-in-context")
+}
+
 func (ts *IntegrationTestSuite) TestLargeQueryResultError() {
 	ctx, cancel := context.WithTimeout(context.Background(), ctxTimeout)
 	defer cancel()

test/test_utils.go

@@ -22,21 +22,30 @@
 package test
 
 import (
+	"context"
+	"fmt"
 	"os"
 	"strings"
 
-	"go.uber.org/cadence/.gen/go/cadence/workflowserviceclient"
 	"go.uber.org/yarpc"
 	"go.uber.org/yarpc/transport/tchannel"
+
+	"go.uber.org/cadence/.gen/go/cadence/workflowserviceclient"
+	"go.uber.org/cadence/workflow"
 )
 
-// Config contains the integration test configuration
-type Config struct {
-	ServiceAddr string
-	ServiceName string
-	IsStickyOff bool
-	Debug       bool
-}
+type (
+	// Config contains the integration test configuration
+	 Config struct {
+		ServiceAddr string
+		ServiceName string
+		IsStickyOff bool
+		Debug       bool
+	}
+
+	// context.WithValue need this type instead of basic type string to avoid lint error
+	contextKey string
+)
 
 func newConfig() Config {
 	cfg := Config{
@@ -107,3 +116,70 @@ func newRPCClient(
 	client := workflowserviceclient.New(dispatcher.ClientConfig(serviceName))
 	return &rpcClient{Interface: client, dispatcher: dispatcher}, nil
 }
+
+// stringMapPropagator propagates the list of keys across a workflow,
+// interpreting the payloads as strings.
+// BORROWED FROM 'internal' PACKAGE TESTS.
+type stringMapPropagator struct {
+	keys map[string]struct{}
+}
+
+// NewStringMapPropagator returns a context propagator that propagates a set of
+// string key-value pairs across a workflow
+func NewStringMapPropagator(keys []string) workflow.ContextPropagator {
+	keyMap := make(map[string]struct{}, len(keys))
+	for _, key := range keys {
+		keyMap[key] = struct{}{}
+	}
+	return &stringMapPropagator{keyMap}
+}
+
+// Inject injects values from context into headers for propagation
+func (s *stringMapPropagator) Inject(ctx context.Context, writer workflow.HeaderWriter) error {
+	for key := range s.keys {
+		value, ok := ctx.Value(contextKey(key)).(string)
+		if !ok {
+			return fmt.Errorf("unable to extract key from context %v", key)
+		}
+		writer.Set(key, []byte(value))
+	}
+	return nil
+}
+
+// InjectFromWorkflow injects values from context into headers for propagation
+func (s *stringMapPropagator) InjectFromWorkflow(ctx workflow.Context, writer workflow.HeaderWriter) error {
+	for key := range s.keys {
+		value, ok := ctx.Value(contextKey(key)).(string)
+		if !ok {
+			return fmt.Errorf("unable to extract key from context %v", key)
+		}
+		writer.Set(key, []byte(value))
+	}
+	return nil
+}
+
+// Extract extracts values from headers and puts them into context
+func (s *stringMapPropagator) Extract(ctx context.Context, reader workflow.HeaderReader) (context.Context, error) {
+	if err := reader.ForEachKey(func(key string, value []byte) error {
+		if _, ok := s.keys[key]; ok {
+			ctx = context.WithValue(ctx, contextKey(key), string(value))
+		}
+		return nil
+	}); err != nil {
+		return nil, err
+	}
+	return ctx, nil
+}
+
+// ExtractToWorkflow extracts values from headers and puts them into context
+func (s *stringMapPropagator) ExtractToWorkflow(ctx workflow.Context, reader workflow.HeaderReader) (workflow.Context, error) {
+	if err := reader.ForEachKey(func(key string, value []byte) error {
+		if _, ok := s.keys[key]; ok {
+			ctx = workflow.WithValue(ctx, contextKey(key), string(value))
+		}
+		return nil
+	}); err != nil {
+		return nil, err
+	}
+	return ctx, nil
+}

test/workflow_test.go

@@ -511,9 +511,21 @@ func (w *Workflows) InspectLocalActivityInfo(ctx workflow.Context) error {
 	wfType := info.WorkflowType.Name
 	taskList := info.TaskListName
 	ctx = workflow.WithLocalActivityOptions(ctx, w.defaultLocalActivityOptions())
-	activites := Activities{}
+	activities := Activities{}
 	return workflow.ExecuteLocalActivity(
-		ctx, activites.InspectActivityInfo, domain, taskList, wfType).Get(ctx, nil)
+		ctx, activities.InspectActivityInfo, domain, taskList, wfType).Get(ctx, nil)
+}
+
+func (w *Workflows) WorkflowWithLocalActivityCtxPropagation(ctx workflow.Context) (string, error) {
+	ctx = workflow.WithLocalActivityOptions(ctx, w.defaultLocalActivityOptions())
+	ctx = workflow.WithValue(ctx, contextKey(testContextKey), "test-data-in-context")
+	activities := Activities{}
+	var result string
+	err := workflow.ExecuteLocalActivity(ctx, activities.DuplicateStringInContext).Get(ctx, &result)
+	if err != nil {
+		return "", err
+	}
+	return result, nil
 }
 
 func (w *Workflows) register(worker worker.Worker) {
@@ -541,6 +553,7 @@ func (w *Workflows) register(worker worker.Worker) {
 	worker.RegisterWorkflow(w.LargeQueryResultWorkflow)
 	worker.RegisterWorkflow(w.RetryTimeoutStableErrorWorkflow)
 	worker.RegisterWorkflow(w.ConsistentQueryWorkflow)
+	worker.RegisterWorkflow(w.WorkflowWithLocalActivityCtxPropagation)
 
 }