Running a module respects the timeout on the provided context (#1999)

* Running a module respects the timeout on the provided context

* AI feedback

* Fix module timeout detection and rename test module build so it's ignored. Remove shadowing of variable

* Updated the wrong file before for the test binary name...

Author: nolag

pkg/workflows/wasm/host/module.go

@@ -50,10 +50,10 @@ var (
 	defaultMaxLogCountNodeMode       = 10_000
 	ResponseBufferTooSmall           = "response buffer too small"
 
-	defaultMaxUserMetricPayloadBytes      = uint32(4096) // 4 KB
-	defaultMaxUserMetricNameLength        = uint32(128)
-	defaultMaxUserMetricLabelsPerMetric   = uint32(10)
-	defaultMaxUserMetricLabelValueLength  = uint32(256)
+	defaultMaxUserMetricPayloadBytes     = uint32(4096) // 4 KB
+	defaultMaxUserMetricNameLength       = uint32(128)
+	defaultMaxUserMetricLabelsPerMetric  = uint32(10)
+	defaultMaxUserMetricLabelValueLength = uint32(256)
 )
 
 type DeterminismConfig struct {
@@ -82,7 +82,7 @@ type ModuleConfig struct {
 	MaxLogCountDONMode  uint32
 	MaxLogCountNodeMode uint32
 
-	EnableUserMetricsLimiter limits.GateLimiter
+	EnableUserMetricsLimiter             limits.GateLimiter
 	MaxUserMetricPayloadBytes            uint32
 	MaxUserMetricPayloadLimiter          limits.BoundLimiter[config.Size] // supersedes MaxUserMetricPayloadBytes if set
 	MaxUserMetricNameLength              uint32
@@ -453,15 +453,15 @@ func linkNoDAG(m *module, store *wasmtime.Store, exec *execution[*sdkpb.Executio
 		return nil, fmt.Errorf("error wrapping await_secrets func: %w", err)
 	}
 
-	if err := linker.FuncWrap(
+	if err = linker.FuncWrap(
 		"env",
 		"log",
 		exec.log,
 	); err != nil {
 		return nil, fmt.Errorf("error wrapping log func: %w", err)
 	}
 
-	if err := linker.FuncWrap(
+	if err = linker.FuncWrap(
 		"env",
 		"emit_metric",
 		exec.emitMetric,
@@ -623,7 +623,16 @@ func runWasm[I, O proto.Message](
 
 	var o O
 
-	ctxWithTimeout, cancel := context.WithTimeout(ctx, *m.cfg.Timeout)
+	// No reason to run the WASM longer if the outer ctx will cancel.
+	ctxDeadline, hasDeadline := ctx.Deadline()
+	var ctxWithTimeout context.Context
+	var cancel func()
+	if hasDeadline && ctxDeadline.Before(time.Now().Add(*m.cfg.Timeout)) {
+		ctxWithTimeout, cancel = context.WithCancel(ctx)
+	} else {
+		ctxWithTimeout, cancel = context.WithTimeout(ctx, *m.cfg.Timeout)
+	}
+
 	defer cancel()
 
 	store := wasmtime.NewStore(m.engine)
@@ -731,7 +740,7 @@ func runWasm[I, O proto.Message](
 	// Note - there is no other reliable signal on the error that can be used to infer it is due to epoch deadline
 	// being reached, so if an error is returned after the deadline it is assumed it is due to that and return
 	// context.DeadlineExceeded.
-	if err != nil && executionDuration >= *m.cfg.Timeout-m.cfg.TickInterval { // As start could be called just before epoch update 1 tick interval is deducted to account for this
+	if err != nil && ((executionDuration >= *m.cfg.Timeout-m.cfg.TickInterval) || ctx.Err() != nil) { // As start could be called just before epoch update 1 tick interval is deducted to account for this
 		m.cfg.Logger.Errorw("start function returned error after deadline reached, returning deadline exceeded error", "errFromStartFunction", err)
 		return o, context.DeadlineExceeded
 	}

pkg/workflows/wasm/host/wasm_nodag_test.go

@@ -60,6 +60,104 @@ func Test_Sleep_Timeout(t *testing.T) {
 	require.Less(t, duration.Seconds(), 3.0, "execution should be interrupted quickly")
 }
 
+func Test_Execute_CtxTimeout(t *testing.T) {
+	t.Parallel()
+	// different build location so it doesn't clash with the other test when building
+	binary := createTestBinary(sleepBinaryCmd, sleepBinaryLocation2, true, t)
+	t.Run("timeout from module is first", func(t *testing.T) {
+		t.Parallel()
+		mc := defaultNoDAGModCfg(t)
+		timeout := time.Second
+		mc.Timeout = &timeout
+		m, err := NewModule(t.Context(), mc, binary)
+		require.NoError(t, err)
+
+		m.v2ImportName = "test"
+		m.Start()
+		defer m.Close()
+
+		mockExecutionHelper := NewMockExecutionHelper(t)
+		mockExecutionHelper.EXPECT().GetWorkflowExecutionID().Return("id")
+		mockExecutionHelper.EXPECT().GetNodeTime().RunAndReturn(func() time.Time {
+			return time.Now()
+		})
+
+		req := &sdk.ExecuteRequest{
+			Request: &sdk.ExecuteRequest_Trigger{},
+		}
+
+		start := time.Now()
+		timeoutCtx, cancel := context.WithTimeout(t.Context(), time.Minute)
+		defer cancel()
+		_, err = m.Execute(timeoutCtx, req, mockExecutionHelper)
+		duration := time.Since(start)
+		require.ErrorIs(t, err, context.DeadlineExceeded)
+		require.Less(t, duration.Seconds(), 3.0, "execution should be interrupted quickly")
+	})
+
+	t.Run("no context timeout", func(t *testing.T) {
+		t.Parallel()
+
+		mc := defaultNoDAGModCfg(t)
+		timeout := time.Second
+		mc.Timeout = &timeout
+		m, err := NewModule(t.Context(), mc, binary)
+		require.NoError(t, err)
+
+		m.v2ImportName = "test"
+		m.Start()
+		defer m.Close()
+
+		mockExecutionHelper := NewMockExecutionHelper(t)
+		mockExecutionHelper.EXPECT().GetWorkflowExecutionID().Return("id")
+		mockExecutionHelper.EXPECT().GetNodeTime().RunAndReturn(func() time.Time {
+			return time.Now()
+		})
+
+		req := &sdk.ExecuteRequest{
+			Request: &sdk.ExecuteRequest_Trigger{},
+		}
+
+		start := time.Now()
+		_, err = m.Execute(t.Context(), req, mockExecutionHelper)
+		duration := time.Since(start)
+		require.ErrorIs(t, err, context.DeadlineExceeded)
+		require.Less(t, duration.Seconds(), 3.0, "execution should be interrupted quickly")
+	})
+
+	t.Run("timeout from context is first", func(t *testing.T) {
+		t.Parallel()
+
+		mc := defaultNoDAGModCfg(t)
+		timeout := time.Minute
+		mc.Timeout = &timeout
+		m, err := NewModule(t.Context(), mc, binary)
+		require.NoError(t, err)
+
+		m.v2ImportName = "test"
+		m.Start()
+		defer m.Close()
+
+		mockExecutionHelper := NewMockExecutionHelper(t)
+		mockExecutionHelper.EXPECT().GetWorkflowExecutionID().Return("id")
+		mockExecutionHelper.EXPECT().GetNodeTime().RunAndReturn(func() time.Time {
+			return time.Now()
+		})
+
+		req := &sdk.ExecuteRequest{
+			Request: &sdk.ExecuteRequest_Trigger{},
+		}
+
+		start := time.Now()
+		timeoutCtx, cancel := context.WithTimeout(t.Context(), time.Second)
+		defer cancel()
+		_, err = m.Execute(timeoutCtx, req, mockExecutionHelper)
+		duration := time.Since(start)
+		require.ErrorIs(t, err, context.DeadlineExceeded)
+		require.Less(t, duration.Seconds(), 3.0, "execution should be interrupted quickly")
+	})
+}
+
 func Test_NoDag_Run(t *testing.T) {
 	t.Parallel()
 

pkg/workflows/wasm/host/wasm_test.go

@@ -35,6 +35,7 @@ const (
 	oomBinaryLocation          = "test/oom/cmd/testmodule.wasm"
 	oomBinaryCmd               = "test/oom/cmd"
 	sleepBinaryLocation        = "test/sleep/cmd/testmodule.wasm"
+	sleepBinaryLocation2       = "test/sleep/cmd/testmodule_2.wasm" // used to avoid a build race between tests
 	sleepBinaryCmd             = "test/sleep/cmd"
 	filesBinaryLocation        = "test/files/cmd/testmodule.wasm"
 	filesBinaryCmd             = "test/files/cmd"