tasks/chasm/context.go

package chasm

import (
	"context"
	"errors"
	"time"

	persistencespb "go.temporal.io/server/api/persistence/v1"
	"go.temporal.io/server/common/log"
	"go.temporal.io/server/common/metrics"
	"go.temporal.io/server/common/namespace"
	"google.golang.org/grpc/metadata"
)

type Context interface {
	// Context is not bound to any component,
	// so all methods needs to take in component as a parameter

	// NOTE: component created in the current transaction won't have a ref
	// this is a Ref to the component state at the start of the transition
	Ref(Component) ([]byte, error)
	// Now returns the current time in the context of the given component.
	// In a context of a transaction, this time must be used to allow for framework support of pause and time skipping.
	Now(Component) time.Time
	// ExecutionKey returns the execution key for the execution the context is operating on.
	ExecutionKey() ExecutionKey
	// ExecutionInfo returns metadata information about the execution.
	ExecutionInfo() ExecutionInfo
	// Logger returns a logger tagged with execution key and other chasm framework internal information.
	Logger() log.Logger
	// NamespaceEntry returns the namespace entry for the execution.
	NamespaceEntry() *namespace.Namespace
	// EndpointByName resolves a nexus endpoint entry.
	EndpointByName(endpointName string) (*persistencespb.NexusEndpointEntry, error)
	// MetricsHandler returns a metrics handler with namespace tag.
	MetricsHandler() metrics.Handler
	// Value returns the value associated with this context for key. The behavior is the same as context.Context.Value().
	// Use WithContextValues RegistrableComponentOption to set key values pair for a component upon registration.
	// Registered key-value pairs will automatically be added to the Context whenever framework accesses the component.
	// Alternatively, use ContextWithValue() to manually set values on Context which will take precedence over registered ones.
	Value(key any) any
	// RequestHeader returns the first value of the named gRPC metadata header from the inbound request context, or ""
	// if absent.
	//
	// Only available when this Context was constructed from an inbound gRPC request, i.e. inside the start/update/read
	// callbacks invoked by the chasm engine. In other contexts, such as pure tasks executed at the end of a transaction
	// or background task handlers, the underlying ctx has no gRPC metadata and this method always returns "".
	RequestHeader(key string) string

	// Intent() OperationIntent
	// ComponentOptions(Component) []ComponentOption

	// withValue should only be used by ContextWithValue() function, do NOT call it directly.
	// For structs implementing this method, although the returned value has type Context,
	// the concrete type MUST be the same concrete type as the receiver.
	withValue(key any, value any) Context
	structuredRef(Component) (ComponentRef, error)
	goContext() context.Context
}

type ExecutionInfo struct {
	// StateTransitionCount is the number of create/update transactions in the history of this execution.
	StateTransitionCount int64
	// ApproximateStateSize is the approximate size in bytes of the persisted execution state of this execution.
	ApproximateStateSize int
	// CloseTime is the time when the execution was closed.
	// An execution is closed when its root component reaches a terminal state in its lifecycle.
	// If the component is still running (not yet closed), it returns a zero time.Time value.
	CloseTime time.Time
}

type EndpointRegistry interface {
	GetByName(ctx context.Context, namespaceID namespace.ID, endpointName string) (*persistencespb.NexusEndpointEntry, error)
}

type MutableContext interface {
	Context

	// AddTask adds a task to be emitted as part of the current transaction.
	// The task is associated with the given component and will be invoked via the registered handler for the given task
	// referencing the component.
	AddTask(Component, TaskAttributes, any)

	// Get a Ref for the component
	// This ref to the component state at the end of the transition
	// Same as Ref(Component) method in Context,
	// this only works for components that already exists at the start of the transition
	//
	// If we provide this method, then the method on the engine doesn't need to
	// return a Ref
	// NewRef(Component) (ComponentRef, bool)
}

type immutableCtx struct {
	// The context here is not really used today.
	// But it will be when we support partial loading later,
	// and the framework potentially needs to go to persistence to load some fields.
	ctx context.Context

	executionKey ExecutionKey

	// Not embedding the Node here to avoid exposing AddTask() method on Node,
	// so that ContextImpl won't implement MutableContext interface.
	root *Node
}

type mutableCtx struct {
	*immutableCtx
}

// NewContext creates a new Context from an existing Context and root Node.
//
// NOTE: Library authors should not invoke this constructor directly, and instead use [ReadComponent].
func NewContext(
	ctx context.Context,
	node *Node,
) Context {
	return newContext(ctx, node)
}

// newContext creates a new immutableCtx from an existing Context and root Node.
// This is similar to NewContext, but returns *immutableCtx instead of Context interface.
func newContext(
	ctx context.Context,
	node *Node,
) *immutableCtx {
	workflowKey := node.backend.GetWorkflowKey()
	return &immutableCtx{
		ctx:  ctx,
		root: node.root(),
		executionKey: ExecutionKey{
			NamespaceID: workflowKey.NamespaceID,
			BusinessID:  workflowKey.WorkflowID,
			RunID:       workflowKey.RunID,
		},
	}
}

func (c *immutableCtx) Ref(component Component) ([]byte, error) {
	return c.root.Ref(component)
}

func (c *immutableCtx) Now(component Component) time.Time {
	return c.root.Now(component)
}

func (c *immutableCtx) ExecutionKey() ExecutionKey {
	return c.executionKey
}

func (c *immutableCtx) ExecutionInfo() ExecutionInfo {
	executionInfo := c.root.backend.GetExecutionInfo()

	var closeTime time.Time
	closeTimestamp := executionInfo.GetCloseTime()
	if closeTimestamp != nil {
		closeTime = closeTimestamp.AsTime()
	}

	return ExecutionInfo{
		StateTransitionCount: executionInfo.GetStateTransitionCount(),
		ApproximateStateSize: c.root.backend.GetApproximatePersistedSize(),
		CloseTime:            closeTime,
	}
}

func (c *immutableCtx) Logger() log.Logger {
	return c.root.logger
}

func (c *immutableCtx) MetricsHandler() metrics.Handler {
	return c.root.metricsHandler
}

func (c *immutableCtx) Value(key any) any {
	if v := c.goContext().Value(key); v != nil {
		return v
	}

	return c.root.registry.componentContextValue(key)
}

func (c *immutableCtx) withValue(key any, value any) Context {
	return &immutableCtx{
		ctx:          context.WithValue(c.goContext(), key, value),
		root:         c.root,
		executionKey: c.executionKey,
	}
}

func (c *immutableCtx) structuredRef(component Component) (ComponentRef, error) {
	return c.root.structuredRef(component)
}

func (c *immutableCtx) NamespaceEntry() *namespace.Namespace {
	return c.root.backend.GetNamespaceEntry()
}

func (c *immutableCtx) goContext() context.Context {
	return c.ctx
}

func (c *immutableCtx) RequestHeader(key string) string {
	if values := metadata.ValueFromIncomingContext(c.ctx, key); len(values) > 0 {
		return values[0]
	}
	return ""
}

func (c *immutableCtx) EndpointByName(name string) (*persistencespb.NexusEndpointEntry, error) {
	reg := c.root.backend.EndpointRegistry()
	if reg == nil {
		return nil, errors.New("endpoint registry not available")
	}
	return reg.GetByName(c.ctx, c.NamespaceEntry().ID(), name)
}

// NewMutableContext creates a new MutableContext from an existing Context and root Node.
//
// NOTE: Library authors should not invoke this constructor directly, and instead use the [UpdateComponent],
// [UpdateWithStartExecution], or [StartExecution] APIs.
func NewMutableContext(
	ctx context.Context,
	node *Node,
) MutableContext {
	return &mutableCtx{
		immutableCtx: newContext(ctx, node),
	}
}

func (c *mutableCtx) AddTask(
	component Component,
	attributes TaskAttributes,
	payload any,
) {
	c.root.AddTask(component, attributes, payload)
}

func (c *mutableCtx) withValue(key any, value any) Context {
	return &mutableCtx{
		immutableCtx: ContextWithValue(c.immutableCtx, key, value),
	}
}

// ContextWithValue returns a new Context with the given key-value pair added.
// Added key-value pairs will be accessible via the Value() method on the returned Context,
// and the behavior of the key-value pair is the same as context.Context.WithValue().
func ContextWithValue[C Context](c C, key any, value any) C {
	//nolint:revive // unchecked-type-assertion
	return any(c.withValue(key, value)).(C)
}