Close and return immediately if the context is already canceled (#1158)

Signed-off-by: Clifton Kaznocha <ckaznocha@users.noreply.github.com>
Co-authored-by: Clifton Kaznocha <ckaznocha@users.noreply.github.com>

internal/wasm/call_context.go

@@ -94,18 +94,27 @@ func (m *CallContext) closeModuleOnCanceledOrTimeout(ctx context.Context, cancel
 			// case go will randomize which branch of the outer select to enter
 			// and we don't want to close the module.
 		default:
-			if errors.Is(ctx.Err(), context.Canceled) {
-				// TODO: figure out how to report error here.
-				_ = m.CloseWithExitCode(ctx, sys.ExitCodeContextCanceled)
-			} else if errors.Is(ctx.Err(), context.DeadlineExceeded) {
-				// TODO: figure out how to report error here.
-				_ = m.CloseWithExitCode(ctx, sys.ExitCodeDeadlineExceeded)
-			}
+			m.CloseWithCtxErr(ctx)
 		}
 	case <-cancelChan:
 	}
 }
 
+// CloseWithExitCode closes the module with an exit code based on the type of
+// error reported by the context.
+//
+// If the context's error is unknown or nil, the module does not close.
+func (m *CallContext) CloseWithCtxErr(ctx context.Context) {
+	switch {
+	case errors.Is(ctx.Err(), context.Canceled):
+		// TODO: figure out how to report error here.
+		_ = m.CloseWithExitCode(ctx, sys.ExitCodeContextCanceled)
+	case errors.Is(ctx.Err(), context.DeadlineExceeded):
+		// TODO: figure out how to report error here.
+		_ = m.CloseWithExitCode(ctx, sys.ExitCodeDeadlineExceeded)
+	}
+}
+
 // Name implements the same method as documented on api.Module
 func (m *CallContext) Name() string {
 	return m.module.Name

internal/wasm/call_context_test.go

@@ -388,6 +388,44 @@ func TestCallContext_CloseModuleOnCanceledOrTimeout(t *testing.T) {
 	})
 }
 
+func TestCallContext_CloseWithCtxErr(t *testing.T) {
+	s := newStore()
+
+	t.Run("context canceled", func(t *testing.T) {
+		cc := &CallContext{Closed: new(uint64), module: &ModuleInstance{Name: "test"}, s: s}
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+
+		cc.CloseWithCtxErr(ctx)
+
+		err := cc.FailIfClosed()
+		require.EqualError(t, err, "module \"test\" closed with context canceled")
+	})
+
+	t.Run("context timeout", func(t *testing.T) {
+		cc := &CallContext{Closed: new(uint64), module: &ModuleInstance{Name: "test"}, s: s}
+		duration := time.Second
+		ctx, cancel := context.WithTimeout(context.Background(), duration)
+		defer cancel()
+
+		time.Sleep(duration * 2)
+
+		cc.CloseWithCtxErr(ctx)
+
+		err := cc.FailIfClosed()
+		require.EqualError(t, err, "module \"test\" closed with context deadline exceeded")
+	})
+
+	t.Run("no error", func(t *testing.T) {
+		cc := &CallContext{Closed: new(uint64), module: &ModuleInstance{Name: "test"}, s: s}
+
+		cc.CloseWithCtxErr(context.Background())
+
+		err := cc.FailIfClosed()
+		require.Nil(t, err)
+	})
+}
+
 type mockCloser struct{ called int }
 
 func (m *mockCloser) Close(context.Context) error {