chore: add WithContext() option for engine.Invoke()

Author: PeterSchafer

pkg/workflow/engineimpl.go

@@ -1,6 +1,7 @@
 package workflow
 
 import (
+	"context"
 	"fmt"
 	"net/http"
 	"net/url"
@@ -37,9 +38,10 @@ type EngineImpl struct {
 var _ Engine = (*EngineImpl)(nil)
 
 type engineRuntimeConfig struct {
-	config configuration.Configuration
-	input  []Data
-	ic     analytics.InstrumentationCollector
+	config  configuration.Configuration
+	input   []Data
+	ic      analytics.InstrumentationCollector
+	ctxFunc func() context.Context
 }
 
 type EngineInvokeOption func(*engineRuntimeConfig)
@@ -62,6 +64,12 @@ func WithInstrumentationCollector(ic analytics.InstrumentationCollector) EngineI
 	}
 }
 
+func WithContext(ctx context.Context) EngineInvokeOption {
+	return func(e *engineRuntimeConfig) {
+		e.ctxFunc = func() context.Context { return ctx }
+	}
+}
+
 func (e *EngineImpl) GetLogger() *zerolog.Logger {
 	return e.logger
 }
@@ -321,11 +329,11 @@ func (e *EngineImpl) Invoke(
 			e.mu.Unlock()
 
 			// create a context object for the invocation
-			context := NewInvocationContext(id, options.config, localEngine, localNetworkAccess, localLogger, localAnalytics, localUi)
+			invocationCtx := NewInvocationContext(options.ctxFunc, id, options.config, localEngine, localNetworkAccess, localLogger, localAnalytics, localUi)
 
 			// invoke workflow through its callback
 			localLogger.Printf("Workflow Start")
-			output, err = callback(context, options.input)
+			output, err = callback(invocationCtx, options.input)
 			localLogger.Printf("Workflow End")
 		}
 	} else {

pkg/workflow/invocationcontextimpl.go

@@ -14,6 +14,7 @@ import (
 )
 
 func NewInvocationContext(
+	ctxFunc func() context.Context,
 	id Identifier,
 	config configuration.Configuration,
 	engine Engine,
@@ -22,7 +23,11 @@ func NewInvocationContext(
 	analyticsImpl analytics.Analytics,
 	ui ui.UserInterface,
 ) InvocationContext {
+	if ctxFunc == nil {
+		ctxFunc = context.Background
+	}
 	return &invocationContextImpl{
+		ctxFunc:        ctxFunc,
 		WorkflowID:     id,
 		Configuration:  config,
 		WorkflowEngine: engine,
@@ -35,6 +40,7 @@ func NewInvocationContext(
 
 // invocationContextImpl is the default implementation of the InvocationContext interface.
 type invocationContextImpl struct {
+	ctxFunc        func() context.Context
 	WorkflowID     Identifier
 	WorkflowEngine Engine
 	Configuration  configuration.Configuration
@@ -47,10 +53,8 @@ type invocationContextImpl struct {
 var _ InvocationContext = (*invocationContextImpl)(nil)
 
 // Context returns the context of the workflow that is being invoked.
-func (*invocationContextImpl) Context() context.Context {
-	// TODO: This is using context.Background() as a placeholder. Ideally this returns
-	// the context representing the lifecycle of the workflow that is being invoked.
-	return context.Background()
+func (ici *invocationContextImpl) Context() context.Context {
+	return ici.ctxFunc()
 }
 
 // GetWorkflowIdentifier returns the identifier of the workflow that is being invoked.