chore: Inject InstrumentationCollector in engine invocation (#370)

Co-authored-by: Jason Luong <[email protected]>

Author: PeterSchafer

pkg/analytics/analytics.go

@@ -315,6 +315,10 @@ func (a *AnalyticsImpl) GetInstrumentation() InstrumentationCollector {
 	return a.instrumentor
 }
 
+func (a *AnalyticsImpl) SetInstrumentation(ic InstrumentationCollector) {
+	a.instrumentor = ic
+}
+
 var DisabledInFedrampErr = errors.New("analytics are disabled in FedRAMP environments") //nolint:errname // breaking API change
 
 // This method sanitizes the given content by searching for key-value mappings. It thereby replaces all keys defined in keysToFilter by the replacement string

pkg/mocks/workflow.go

@@ -126,6 +126,69 @@ func Test_EngineBasics(t *testing.T) {
 	assert.Equal(t, "callback2", extension["cmd2::hello"])
 }
 
+func Test_EngineBasics_InvokeWithCollector(t *testing.T) {
+	engine := NewWorkFlowEngine(configuration.NewWithOpts())
+	collector := analytics.NewInstrumentationCollector()
+
+	expectedKey := "some"
+	expectedValue := "value"
+
+	test1Id := NewWorkflowIdentifier("test1")
+	test2Id := NewWorkflowIdentifier("test2")
+	noOpWorkflowOptions := ConfigurationOptionsFromFlagset(pflag.NewFlagSet("", pflag.ContinueOnError))
+	input := []Data{NewData(NewTypeIdentifier(test1Id, "input"), "random", nil)}
+
+	_, err := engine.Register(test1Id, noOpWorkflowOptions, func(invocation InvocationContext, input []Data) ([]Data, error) {
+		invocation.GetAnalytics().AddExtensionStringValue(expectedKey, expectedValue)
+		return invocation.GetEngine().Invoke(test2Id, WithInput(input)) // without specifying the collector, Invoke will use the default
+	})
+	assert.NoError(t, err)
+
+	_, err = engine.Register(test2Id, noOpWorkflowOptions, func(invocation InvocationContext, input []Data) ([]Data, error) {
+		invocation.GetAnalytics().AddExtensionStringValue(expectedKey, expectedValue)
+		return input, nil
+	})
+	assert.NoError(t, err)
+
+	err = engine.Init()
+	assert.NoError(t, err)
+
+	testcases := []struct {
+		name            string
+		invokeOptions   []EngineInvokeOption
+		assertCollector analytics.InstrumentationCollector
+	}{
+		{
+			name:            "Use specified Instrumentation Collector",
+			assertCollector: collector,
+			invokeOptions:   []EngineInvokeOption{WithInput(input), WithInstrumentationCollector(collector)},
+		},
+		{
+			name:            "Use default Instrumentation Collector",
+			assertCollector: engine.GetAnalytics().GetInstrumentation(),
+			invokeOptions:   []EngineInvokeOption{WithInput(input)},
+		},
+	}
+
+	for _, tc := range testcases {
+		t.Run(tc.name, func(t *testing.T) {
+			// invoke method under test, case collector is specified
+			output, err := engine.Invoke(test1Id, tc.invokeOptions...)
+			assert.NoError(t, err)
+			assert.Equal(t, input, output)
+
+			// ensure the provided collector has the expected data
+			instrumentationData, err := analytics.GetV2InstrumentationObject(tc.assertCollector)
+			assert.NoError(t, err)
+			assert.NotNil(t, instrumentationData)
+
+			extension := *instrumentationData.Data.Attributes.Interaction.Extension
+			assert.Equal(t, expectedValue, extension["test1::some"])
+			assert.Equal(t, expectedValue, extension["test2::some"])
+		})
+	}
+}
+
 func Test_EngineRegisterErrorHandling(t *testing.T) {
 	configuration := configuration.New()
 	engine := NewWorkFlowEngine(configuration)

pkg/workflow/engine_test.go

@@ -36,6 +36,32 @@ type EngineImpl struct {
 
 var _ Engine = (*EngineImpl)(nil)
 
+type engineRuntimeConfig struct {
+	config configuration.Configuration
+	input  []Data
+	ic     analytics.InstrumentationCollector
+}
+
+type EngineInvokeOption func(*engineRuntimeConfig)
+
+func WithConfig(config configuration.Configuration) EngineInvokeOption {
+	return func(e *engineRuntimeConfig) {
+		e.config = config
+	}
+}
+
+func WithInput(input []Data) EngineInvokeOption {
+	return func(e *engineRuntimeConfig) {
+		e.input = input
+	}
+}
+
+func WithInstrumentationCollector(ic analytics.InstrumentationCollector) EngineInvokeOption {
+	return func(e *engineRuntimeConfig) {
+		e.ic = ic
+	}
+}
+
 func (e *EngineImpl) GetLogger() *zerolog.Logger {
 	return e.logger
 }
@@ -214,26 +240,35 @@ func (e *EngineImpl) GetWorkflow(id Identifier) (Entry, bool) {
 	return workflow, ok
 }
 
-// Invoke invokes the workflow with the given identifier.
-func (e *EngineImpl) Invoke(id Identifier) ([]Data, error) {
-	return e.InvokeWithInputAndConfig(id, []Data{}, nil)
-}
-
+// Deprecated: Use Invoke() with WithInput() instead
+//
 // InvokeWithInput invokes the workflow with the given identifier and input data.
 func (e *EngineImpl) InvokeWithInput(id Identifier, input []Data) ([]Data, error) {
-	return e.InvokeWithInputAndConfig(id, input, nil)
+	return e.Invoke(id, WithInput(input))
 }
 
+// Deprecated: Use Invoke() with WithConfig() instead
+//
 // InvokeWithConfig invokes the workflow with the given identifier and configuration.
 func (e *EngineImpl) InvokeWithConfig(id Identifier, config configuration.Configuration) ([]Data, error) {
-	return e.InvokeWithInputAndConfig(id, []Data{}, config)
+	return e.Invoke(id, WithConfig(config))
 }
 
-// InvokeWithInputAndConfig invokes the workflow with the given identifier, input data and configuration.
+// Deprecated: Use Invoke() with WithInput() and WithConfig() instead
+//
+// InvokeWithInputAndConfig invokes the workflow with the given identifier, input data, and configuration.
 func (e *EngineImpl) InvokeWithInputAndConfig(
 	id Identifier,
 	input []Data,
 	config configuration.Configuration,
+) ([]Data, error) {
+	return e.Invoke(id, WithConfig(config), WithInput(input))
+}
+
+// Invoke invokes the workflow with the given identifier.
+func (e *EngineImpl) Invoke(
+	id Identifier,
+	opts ...EngineInvokeOption,
 ) ([]Data, error) {
 	var output []Data
 	var err error
@@ -249,30 +284,49 @@ func (e *EngineImpl) InvokeWithInputAndConfig(
 			e.mu.Lock()
 			e.invocationCounter++
 
+			// create default options
+			options := engineRuntimeConfig{
+				config: e.config.Clone(),
+				input:  []Data{},
+			}
+
+			// override default options based on optional parameters
+			for _, opt := range opts {
+				opt(&options)
+			}
+
 			// prepare logger
 			prefix := fmt.Sprintf("%s:%d", id.Host, e.invocationCounter)
-			zlogger := e.logger.With().Str("ext", prefix).Logger()
+			localLogger := e.logger.With().Str("ext", prefix).Logger()
 
 			localUi := e.ui
-			localAnalytics := NewAnalyticsWrapper(e.analytics, id.Host)
 
-			// prepare configuration
-			if config == nil {
-				config = e.config.Clone()
+			var localAnalytics analytics.Analytics
+			if options.ic != nil {
+				tmpAnalytics := &analytics.AnalyticsImpl{}
+				tmpAnalytics.SetInstrumentation(options.ic)
+				localAnalytics = NewAnalyticsWrapper(tmpAnalytics, id.Host)
+			} else {
+				localAnalytics = NewAnalyticsWrapper(e.analytics, id.Host)
 			}
 
 			// prepare networkAccess
-			networkAccess := e.networkAccess.Clone()
-			networkAccess.SetConfiguration(config)
+			localNetworkAccess := e.networkAccess.Clone()
+			localNetworkAccess.SetConfiguration(options.config)
+
+			localEngine := &engineWrapper{
+				WrappedEngine:                   e,
+				defaultInstrumentationCollector: options.ic,
+			}
 			e.mu.Unlock()
 
 			// create a context object for the invocation
-			context := NewInvocationContext(id, config, e, networkAccess, zlogger, localAnalytics, localUi)
+			context := NewInvocationContext(id, options.config, localEngine, localNetworkAccess, localLogger, localAnalytics, localUi)
 
 			// invoke workflow through its callback
-			zlogger.Printf("Workflow Start")
-			output, err = callback(context, input)
-			zlogger.Printf("Workflow End")
+			localLogger.Printf("Workflow Start")
+			output, err = callback(context, options.input)
+			localLogger.Printf("Workflow End")
 		}
 	} else {
 		err = fmt.Errorf("workflow '%v' not found", id)

pkg/workflow/engineimpl.go

@@ -0,0 +1,114 @@
+package workflow
+
+import (
+	"github.com/rs/zerolog"
+
+	"github.com/snyk/go-application-framework/pkg/analytics"
+	"github.com/snyk/go-application-framework/pkg/configuration"
+	"github.com/snyk/go-application-framework/pkg/networking"
+	"github.com/snyk/go-application-framework/pkg/runtimeinfo"
+	"github.com/snyk/go-application-framework/pkg/ui"
+)
+
+type engineWrapper struct {
+	WrappedEngine                   Engine
+	defaultInstrumentationCollector analytics.InstrumentationCollector
+}
+
+var _ Engine = (*engineWrapper)(nil)
+
+func (e *engineWrapper) Init() error {
+	return e.WrappedEngine.Init()
+}
+
+func (e *engineWrapper) AddExtensionInitializer(initializer ExtensionInit) {
+	e.WrappedEngine.AddExtensionInitializer(initializer)
+}
+
+func (e *engineWrapper) Register(id Identifier, config ConfigurationOptions, callback Callback) (Entry, error) {
+	return e.WrappedEngine.Register(id, config, callback)
+}
+
+func (e *engineWrapper) GetWorkflows() []Identifier {
+	return e.WrappedEngine.GetWorkflows()
+}
+
+func (e *engineWrapper) GetWorkflow(id Identifier) (Entry, bool) {
+	return e.WrappedEngine.GetWorkflow(id)
+}
+
+// Invoke invokes the workflow with the given identifier.
+func (e *engineWrapper) Invoke(id Identifier, opts ...EngineInvokeOption) ([]Data, error) {
+	options := &engineRuntimeConfig{}
+	for _, opt := range opts {
+		opt(options)
+	}
+
+	// if no InstrumentationCollector is specified, and a default is available, the default be used
+	if options.ic == nil && e.defaultInstrumentationCollector != nil {
+		opts = append(opts, WithInstrumentationCollector(e.defaultInstrumentationCollector))
+	}
+
+	return e.WrappedEngine.Invoke(id, opts...)
+}
+
+// Deprecated: Use Invoke() with WithInput() instead
+//
+// InvokeWithInput invokes the workflow with the given identifier and input data.
+func (e *engineWrapper) InvokeWithInput(id Identifier, input []Data) ([]Data, error) {
+	return e.Invoke(id, WithInput(input))
+}
+
+// Deprecated: Use Invoke() with WithConfig() instead
+//
+// InvokeWithConfig invokes the workflow with the given identifier and configuration.
+func (e *engineWrapper) InvokeWithConfig(id Identifier, config configuration.Configuration) ([]Data, error) {
+	return e.Invoke(id, WithConfig(config))
+}
+
+// Deprecated: Use Invoke() with WithInput() and WithConfig() instead
+//
+// InvokeWithInputAndConfig invokes the workflow with the given identifier, input data, and configuration.
+func (e *engineWrapper) InvokeWithInputAndConfig(id Identifier, input []Data, config configuration.Configuration) ([]Data, error) {
+	return e.Invoke(id, WithInput(input), WithConfig(config))
+}
+
+func (e *engineWrapper) GetAnalytics() analytics.Analytics {
+	return e.WrappedEngine.GetAnalytics()
+}
+
+func (e *engineWrapper) GetNetworkAccess() networking.NetworkAccess {
+	return e.WrappedEngine.GetNetworkAccess()
+}
+
+func (e *engineWrapper) GetConfiguration() configuration.Configuration {
+	return e.WrappedEngine.GetConfiguration()
+}
+
+func (e *engineWrapper) SetLogger(logger *zerolog.Logger) {
+	e.WrappedEngine.SetLogger(logger)
+}
+
+func (e *engineWrapper) SetConfiguration(config configuration.Configuration) {
+	e.WrappedEngine.SetConfiguration(config)
+}
+
+func (e *engineWrapper) GetLogger() *zerolog.Logger {
+	return e.WrappedEngine.GetLogger()
+}
+
+func (e *engineWrapper) GetUserInterface() ui.UserInterface {
+	return e.WrappedEngine.GetUserInterface()
+}
+
+func (e *engineWrapper) SetUserInterface(ui ui.UserInterface) {
+	e.WrappedEngine.SetUserInterface(ui)
+}
+
+func (e *engineWrapper) GetRuntimeInfo() runtimeinfo.RuntimeInfo {
+	return e.WrappedEngine.GetRuntimeInfo()
+}
+
+func (e *engineWrapper) SetRuntimeInfo(ri runtimeinfo.RuntimeInfo) {
+	e.WrappedEngine.SetRuntimeInfo(ri)
+}