diff --git a/pkg/workflow/dataimpl.go b/pkg/workflow/dataimpl.go index 8346f05e0..a3b2efcc8 100644 --- a/pkg/workflow/dataimpl.go +++ b/pkg/workflow/dataimpl.go @@ -66,10 +66,18 @@ func withPayload(payload *interface{}) Option { } } +// WithConfiguration copies IN_MEMORY_THRESHOLD_BYTES and TEMP_DIR_PATH from +// config when each key resolves to a value (including AddDefaultValue). If +// TEMP_DIR_PATH is absent, d.tempDirPath is left unchanged (typically ""), and +// os.CreateTemp uses the process default temp directory. func WithConfiguration(config configuration.Configuration) Option { return func(d *DataImpl) { - d.inMemoryThreshold = config.GetInt(configuration.IN_MEMORY_THRESHOLD_BYTES) - d.tempDirPath = config.GetString(configuration.TEMP_DIR_PATH) + if v := config.Get(configuration.IN_MEMORY_THRESHOLD_BYTES); v != nil { + d.inMemoryThreshold = config.GetInt(configuration.IN_MEMORY_THRESHOLD_BYTES) + } + if v := config.Get(configuration.TEMP_DIR_PATH); v != nil { + d.tempDirPath = config.GetString(configuration.TEMP_DIR_PATH) + } } } @@ -256,6 +264,28 @@ func (d *DataImpl) AddError(err snyk_errors.Error) { d.errors = append(d.errors, err) } +// applyConfiguration re-evaluates the payload location using the given +// configuration. Field updates use the same rules as WithConfiguration. +// If the payload is currently in memory and exceeds the +// configured threshold, it is written to disk under the configured temp +// directory. This allows the engine to apply its configuration to Data +// objects that were created without WithConfiguration. +func (d *DataImpl) applyConfiguration(config configuration.Configuration) { + if config.Get(configuration.IN_MEMORY_THRESHOLD_BYTES) == nil { + return + } + + WithConfiguration(config)(d) + + if d.payloadLocation.Type == InMemory && d.payload != nil { + d.payloadLocation = setPayloadLocation(d.identifier, d.inMemoryThreshold, d.tempDirPath, d.payload, d.logger) + if d.payloadLocation.Type == OnDisk { + d.logger.Debug().Msg("payload relocated to disk after applyConfiguration") + d.payload = nil + } + } +} + func setPayloadLocation(id Identifier, inMemoryThreshold int, tempDirPath string, payload interface{}, logger *zerolog.Logger) Location { payloadLocation := Location{ Path: "", diff --git a/pkg/workflow/dataimpl_test.go b/pkg/workflow/dataimpl_test.go index 848298146..1cb072945 100644 --- a/pkg/workflow/dataimpl_test.go +++ b/pkg/workflow/dataimpl_test.go @@ -11,6 +11,7 @@ import ( "github.com/snyk/error-catalog-golang-public/snyk_errors" "github.com/snyk/go-application-framework/pkg/configuration" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func Test_NewDataFromInput(t *testing.T) { @@ -194,6 +195,140 @@ func Test_NewData(t *testing.T) { assert.Contains(t, actualFiles[0].Name(), expectedFileName) }) + t.Run("applyConfiguration relocates in-memory payload to disk", func(t *testing.T) { + tmpDir := t.TempDir() + logger := zerolog.Nop() + + id := NewTypeIdentifier(NewWorkflowIdentifier("cmd"), "applytest") + payloadBytes := []byte("payload that should end up on disk after applyConfiguration") + + data := NewData(id, "application/octet-stream", payloadBytes, WithLogger(&logger)) + + // Without WithConfiguration, payload stays in memory (threshold defaults to -1) + di, ok := data.(*DataImpl) + require.True(t, ok) + assert.Equal(t, InMemory, di.payloadLocation.Type) + assert.NotNil(t, di.payload) + + // Now apply a configuration with threshold=0 — everything spills to disk + cfg := configuration.NewInMemory() + cfg.Set(configuration.IN_MEMORY_THRESHOLD_BYTES, 0) + cfg.Set(configuration.TEMP_DIR_PATH, tmpDir) + di.applyConfiguration(cfg) + + assert.Equal(t, OnDisk, di.payloadLocation.Type) + assert.Nil(t, di.payload) + assert.Contains(t, di.payloadLocation.Path, tmpDir) + + // GetPayload still returns the original bytes + result := data.GetPayload() + assert.Equal(t, payloadBytes, result) + }) + + t.Run("applyConfiguration is no-op when threshold is disabled", func(t *testing.T) { + logger := zerolog.Nop() + id := NewTypeIdentifier(NewWorkflowIdentifier("cmd"), "noop") + payloadBytes := []byte("stays in memory") + + data := NewData(id, "text/plain", payloadBytes, WithLogger(&logger)) + di, ok := data.(*DataImpl) + require.True(t, ok) + + cfg := configuration.NewInMemory() + cfg.Set(configuration.IN_MEMORY_THRESHOLD_BYTES, -1) + di.applyConfiguration(cfg) + + assert.Equal(t, InMemory, di.payloadLocation.Type) + assert.NotNil(t, di.payload) + }) + + t.Run("applyConfiguration is no-op when threshold key is not set in config", func(t *testing.T) { + logger := zerolog.Nop() + id := NewTypeIdentifier(NewWorkflowIdentifier("cmd"), "unset") + payloadBytes := []byte("should stay in memory because key is unset") + + data := NewData(id, "text/plain", payloadBytes, WithLogger(&logger)) + di, ok := data.(*DataImpl) + require.True(t, ok) + assert.Equal(t, -1, di.inMemoryThreshold) + + cfg := configuration.NewInMemory() + require.Nil(t, cfg.Get(configuration.IN_MEMORY_THRESHOLD_BYTES)) + di.applyConfiguration(cfg) + + assert.Equal(t, -1, di.inMemoryThreshold) + assert.Equal(t, InMemory, di.payloadLocation.Type) + assert.NotNil(t, di.payload) + }) + + t.Run("applyConfiguration uses AddDefaultValue when key is not Set", func(t *testing.T) { + tmpDir := t.TempDir() + logger := zerolog.Nop() + id := NewTypeIdentifier(NewWorkflowIdentifier("cmd"), "adddefault") + payloadBytes := []byte("spill via default value functions") + + data := NewData(id, "application/octet-stream", payloadBytes, WithLogger(&logger)) + di, ok := data.(*DataImpl) + require.True(t, ok) + + cfg := configuration.NewInMemory() + cfg.AddDefaultValue(configuration.IN_MEMORY_THRESHOLD_BYTES, configuration.StandardDefaultValueFunction(0)) + cfg.AddDefaultValue(configuration.TEMP_DIR_PATH, configuration.StandardDefaultValueFunction(tmpDir)) + + assert.False(t, cfg.IsSet(configuration.IN_MEMORY_THRESHOLD_BYTES)) + assert.False(t, cfg.IsSet(configuration.TEMP_DIR_PATH)) + require.NotNil(t, cfg.Get(configuration.IN_MEMORY_THRESHOLD_BYTES)) + require.NotNil(t, cfg.Get(configuration.TEMP_DIR_PATH)) + + di.applyConfiguration(cfg) + + assert.Equal(t, OnDisk, di.payloadLocation.Type) + assert.Nil(t, di.payload) + assert.Contains(t, di.payloadLocation.Path, tmpDir) + }) + + t.Run("applyConfiguration with threshold only uses system temp when TEMP_DIR_PATH absent", func(t *testing.T) { + logger := zerolog.Nop() + id := NewTypeIdentifier(NewWorkflowIdentifier("cmd"), "notemp") + payloadBytes := []byte("spill with no temp path in config") + + data := NewData(id, "application/octet-stream", payloadBytes, WithLogger(&logger)) + di, ok := data.(*DataImpl) + require.True(t, ok) + + cfg := configuration.NewInMemory() + cfg.Set(configuration.IN_MEMORY_THRESHOLD_BYTES, 0) + require.Nil(t, cfg.Get(configuration.TEMP_DIR_PATH)) + + di.applyConfiguration(cfg) + + assert.Equal(t, OnDisk, di.payloadLocation.Type) + assert.Contains(t, di.payloadLocation.Path, os.TempDir()) + }) + + t.Run("applyConfiguration is no-op when payload is already on disk", func(t *testing.T) { + tmpDir := t.TempDir() + logger := zerolog.Nop() + + cfg := configuration.NewInMemory() + cfg.Set(configuration.IN_MEMORY_THRESHOLD_BYTES, 0) + cfg.Set(configuration.TEMP_DIR_PATH, tmpDir) + + id := NewTypeIdentifier(NewWorkflowIdentifier("cmd"), "alreadyondisk") + payloadBytes := []byte("on disk from the start") + + data := NewData(id, "application/octet-stream", payloadBytes, WithConfiguration(cfg), WithLogger(&logger)) + di, ok := data.(*DataImpl) + require.True(t, ok) + assert.Equal(t, OnDisk, di.payloadLocation.Type) + originalPath := di.payloadLocation.Path + + // applyConfiguration again should not relocate + di.applyConfiguration(cfg) + assert.Equal(t, OnDisk, di.payloadLocation.Type) + assert.Equal(t, originalPath, di.payloadLocation.Path) + }) + t.Run("when configuration is not provided, filesystem cache is not used", func(t *testing.T) { expectedConfig := configuration.NewInMemory() expectedLogger := zerolog.Logger{} diff --git a/pkg/workflow/engine_test.go b/pkg/workflow/engine_test.go index 7217b3e76..ace6d63ee 100644 --- a/pkg/workflow/engine_test.go +++ b/pkg/workflow/engine_test.go @@ -598,3 +598,42 @@ func Test_EngineImpl_InvokeWithContext_DefaultContext(t *testing.T) { assert.NoError(t, err) assert.NotNil(t, receivedCtx) } + +func Test_Invoke_AppliesConfigurationToOutput(t *testing.T) { + tmpDir := t.TempDir() + + config := configuration.New() + config.Set(configuration.IN_MEMORY_THRESHOLD_BYTES, 0) + config.Set(configuration.TEMP_DIR_PATH, tmpDir) + + engine := NewWorkFlowEngine(config) + + wfId := NewWorkflowIdentifier("threshold-test") + flagset := pflag.NewFlagSet("tt", pflag.ContinueOnError) + + payload := []byte("this payload should be relocated to disk by the engine after Invoke") + + _, err := engine.Register(wfId, ConfigurationOptionsFromFlagset(flagset), func(invocation InvocationContext, input []Data) ([]Data, error) { + // Create Data WITHOUT WithConfiguration — simulates what most callers do + id := NewTypeIdentifier(invocation.GetWorkflowIdentifier(), "testdata") + d := NewData(id, "application/octet-stream", payload) + return []Data{d}, nil + }) + assert.NoError(t, err) + assert.NoError(t, engine.Init()) + + output, err := engine.Invoke(wfId) + assert.NoError(t, err) + assert.Len(t, output, 1) + + // The engine should have applied its config, relocating the payload to disk + di, ok := output[0].(*DataImpl) + assert.True(t, ok) + assert.Equal(t, OnDisk, di.payloadLocation.Type) + assert.Contains(t, di.payloadLocation.Path, tmpDir) + assert.Nil(t, di.payload) + + // GetPayload still returns the original bytes from disk + result := output[0].GetPayload() + assert.Equal(t, payload, result) +} diff --git a/pkg/workflow/engineimpl.go b/pkg/workflow/engineimpl.go index 351c5572e..d762281af 100644 --- a/pkg/workflow/engineimpl.go +++ b/pkg/workflow/engineimpl.go @@ -338,6 +338,15 @@ func (e *EngineImpl) Invoke( localLogger.Printf("Workflow Start") output, err = callback(invocationCtx, options.input) localLogger.Printf("Workflow End") + + // Apply the engine's configuration to output data so that + // IN_MEMORY_THRESHOLD_BYTES and TEMP_DIR_PATH are respected + // even when workflows create Data without WithConfiguration. + for _, d := range output { + if di, ok := d.(*DataImpl); ok { + di.applyConfiguration(options.config) + } + } } } else { err = fmt.Errorf("workflow '%v' not found", id)