diff --git a/cmd/proxygenerator/interceptor.go b/cmd/proxygenerator/interceptor.go index 787b8944..1a6f8830 100644 --- a/cmd/proxygenerator/interceptor.go +++ b/cmd/proxygenerator/interceptor.go @@ -106,19 +106,12 @@ func VisitPayloads(ctx context.Context, msg proto.Message, options VisitPayloads return fmt.Errorf("ConcurrencyLimit must be 0 or greater, got %d", options.ConcurrencyLimit) } visitCtx := VisitPayloadsContext{Context: ctx, Parent: msg} - if options.ConcurrencyLimit <= 1 { - return visitPayloads(&visitCtx, &options, nil, nil, msg) - } - c := &payloadConcurrencyState{sem: make(chan struct{}, options.ConcurrencyLimit)} - err := visitPayloads(&visitCtx, &options, nil, c, msg) - c.wg.Wait() - if err != nil { - return err - } - if errPtr := c.firstErr.Load(); errPtr != nil { - return *errPtr + + var c *payloadConcurrencyState + if options.ConcurrencyLimit > 1 { + c = &payloadConcurrencyState{sem: make(chan struct{}, options.ConcurrencyLimit)} } - return nil + return visitPayloadsAndWait(&visitCtx, &options, nil, c, msg) } // PayloadVisitorInterceptorOptions configures outbound/inbound interception of Payloads within msgs. @@ -275,32 +268,20 @@ func (o *VisitFailuresOptions) defaultWellKnownAnyVisitor(ctx *VisitFailuresCont } func (o *VisitPayloadsOptions) defaultWellKnownAnyVisitor(ctx *VisitPayloadsContext, concState *payloadConcurrencyState, p *anypb.Any) error { + // We choose to visit and re-marshal instead of cloning, visiting, + // and checking if anything changed before re-marshaling. It is assumed the + // clone + equality check is not much cheaper than re-marshal. child, err := p.UnmarshalNew() if err != nil { return fmt.Errorf("failed to unmarshal any: %w", err) } - // Sub-state shares the semaphore but has its own WaitGroup so we can wait - // for goroutines writing into child's fields before re-marshaling. - var anyConcState *payloadConcurrencyState - if concState != nil { - anyConcState = &payloadConcurrencyState{sem: concState.sem} - } - - // We choose to visit and re-marshal always instead of cloning, visiting, - // and checking if anything changed before re-marshaling. It is assumed the - // clone + equality check is not much cheaper than re-marshal. - if err := visitPayloads(ctx, o, p, anyConcState, child); err != nil { + // Visit payloads and wait for goroutines to finish writing child fields + // before re-marshaling. + if err := visitPayloadsAndWait(ctx, o, p, concState, child); err != nil { return err } - if anyConcState != nil { - anyConcState.wg.Wait() - if errPtr := anyConcState.firstErr.Load(); errPtr != nil { - return *errPtr - } - } - // Confirmed this replaces both Any fields on non-error, there is nothing // left over if err := p.MarshalFrom(child); err != nil { @@ -367,6 +348,34 @@ func visitPayload( return nil } +// visitPayloadsAndWait invokes visitPayloads and waits for its child +// goroutines to finish before returning. +func visitPayloadsAndWait( + ctx *VisitPayloadsContext, + options *VisitPayloadsOptions, + parent proto.Message, + concState *payloadConcurrencyState, + objs ...interface{}, +) error { + if concState == nil { + return visitPayloads(ctx, options, parent, concState, objs...) + } + + // Create a new concurrency state with the same semaphore for concurrency + // control, but a new WaitGroup so we can wait for just the child goroutines. + miniConcState := &payloadConcurrencyState{sem: concState.sem} + err := visitPayloads(ctx, options, parent, miniConcState, objs...) + miniConcState.wg.Wait() + + if err != nil { + return err + } + if errPtr := miniConcState.firstErr.Load(); errPtr != nil { + return *errPtr + } + return nil +} + func visitPayloads( ctx *VisitPayloadsContext, options *VisitPayloadsOptions, @@ -550,6 +559,14 @@ func visitPayloads( if err := visitPayload(ctx, options, o, concState, &result); err != nil { return err } o.Outcome = &workflowservice.PollNexusOperationExecutionResponse_Result{Result: result} } + {{else if and (eq $type "*command.ScheduleNexusOperationCommandAttributes") (eq . "Input")}} + if o.Input != nil { + if o.GetEndpoint() == "__temporal_system" { + if err := visitSystemNexusEnvelope(ctx, options, concState, o); err != nil { return err } + } else { + if err := visitPayload(ctx, options, o, concState, &o.Input); err != nil { return err } + } + } {{else}} if o.{{.}} != nil { if err := visitPayload(ctx, options, o, concState, &o.{{.}}); err != nil { return err } diff --git a/proxy/interceptor.go b/proxy/interceptor.go index b4519197..edad35ec 100644 --- a/proxy/interceptor.go +++ b/proxy/interceptor.go @@ -103,19 +103,12 @@ func VisitPayloads(ctx context.Context, msg proto.Message, options VisitPayloads return fmt.Errorf("ConcurrencyLimit must be 0 or greater, got %d", options.ConcurrencyLimit) } visitCtx := VisitPayloadsContext{Context: ctx, Parent: msg} - if options.ConcurrencyLimit <= 1 { - return visitPayloads(&visitCtx, &options, nil, nil, msg) - } - c := &payloadConcurrencyState{sem: make(chan struct{}, options.ConcurrencyLimit)} - err := visitPayloads(&visitCtx, &options, nil, c, msg) - c.wg.Wait() - if err != nil { - return err - } - if errPtr := c.firstErr.Load(); errPtr != nil { - return *errPtr + + var c *payloadConcurrencyState + if options.ConcurrencyLimit > 1 { + c = &payloadConcurrencyState{sem: make(chan struct{}, options.ConcurrencyLimit)} } - return nil + return visitPayloadsAndWait(&visitCtx, &options, nil, c, msg) } // PayloadVisitorInterceptorOptions configures outbound/inbound interception of Payloads within msgs. @@ -271,32 +264,20 @@ func (o *VisitFailuresOptions) defaultWellKnownAnyVisitor(ctx *VisitFailuresCont } func (o *VisitPayloadsOptions) defaultWellKnownAnyVisitor(ctx *VisitPayloadsContext, concState *payloadConcurrencyState, p *anypb.Any) error { + // We choose to visit and re-marshal instead of cloning, visiting, + // and checking if anything changed before re-marshaling. It is assumed the + // clone + equality check is not much cheaper than re-marshal. child, err := p.UnmarshalNew() if err != nil { return fmt.Errorf("failed to unmarshal any: %w", err) } - // Sub-state shares the semaphore but has its own WaitGroup so we can wait - // for goroutines writing into child's fields before re-marshaling. - var anyConcState *payloadConcurrencyState - if concState != nil { - anyConcState = &payloadConcurrencyState{sem: concState.sem} - } - - // We choose to visit and re-marshal always instead of cloning, visiting, - // and checking if anything changed before re-marshaling. It is assumed the - // clone + equality check is not much cheaper than re-marshal. - if err := visitPayloads(ctx, o, p, anyConcState, child); err != nil { + // Visit payloads and wait for goroutines to finish writing child fields + // before re-marshaling. + if err := visitPayloadsAndWait(ctx, o, p, concState, child); err != nil { return err } - if anyConcState != nil { - anyConcState.wg.Wait() - if errPtr := anyConcState.firstErr.Load(); errPtr != nil { - return *errPtr - } - } - // Confirmed this replaces both Any fields on non-error, there is nothing // left over if err := p.MarshalFrom(child); err != nil { @@ -363,6 +344,34 @@ func visitPayload( return nil } +// visitPayloadsAndWait invokes visitPayloads and waits for its child +// goroutines to finish before returning. +func visitPayloadsAndWait( + ctx *VisitPayloadsContext, + options *VisitPayloadsOptions, + parent proto.Message, + concState *payloadConcurrencyState, + objs ...interface{}, +) error { + if concState == nil { + return visitPayloads(ctx, options, parent, concState, objs...) + } + + // Create a new concurrency state with the same semaphore for concurrency + // control, but a new WaitGroup so we can wait for just the child goroutines. + miniConcState := &payloadConcurrencyState{sem: concState.sem} + err := visitPayloads(ctx, options, parent, miniConcState, objs...) + miniConcState.wg.Wait() + + if err != nil { + return err + } + if errPtr := miniConcState.firstErr.Load(); errPtr != nil { + return *errPtr + } + return nil +} + func visitPayloads( ctx *VisitPayloadsContext, options *VisitPayloadsOptions, @@ -989,8 +998,14 @@ func visitPayloads( } if o.Input != nil { - if err := visitPayload(ctx, options, o, concState, &o.Input); err != nil { - return err + if o.GetEndpoint() == "__temporal_system" { + if err := visitSystemNexusEnvelope(ctx, options, concState, o); err != nil { + return err + } + } else { + if err := visitPayload(ctx, options, o, concState, &o.Input); err != nil { + return err + } } } diff --git a/proxy/system_nexus.go b/proxy/system_nexus.go new file mode 100644 index 00000000..81350e53 --- /dev/null +++ b/proxy/system_nexus.go @@ -0,0 +1,72 @@ +package proxy + +import ( + "fmt" + + "go.temporal.io/api/command/v1" + "google.golang.org/protobuf/proto" + "google.golang.org/protobuf/reflect/protoreflect" + "google.golang.org/protobuf/reflect/protoregistry" +) + +// visitSystemNexusEnvelope decodes the system Nexus envelope in attrs.Input, +// visits the payloads inside the decoded request message, and re-encodes it. +// +// The envelope's proto message type is taken from the payload's "messageType" +// metadata, so no operation registry is required. The envelope must be encoded +// as binary/protobuf. The inner payloads (and only those) are passed to the +// visitor, so external storage and codecs apply to them and not to the envelope +// itself, which is never offloaded or codec-encoded. +func visitSystemNexusEnvelope( + ctx *VisitPayloadsContext, + options *VisitPayloadsOptions, + concState *payloadConcurrencyState, + attrs *command.ScheduleNexusOperationCommandAttributes, +) error { + input := attrs.Input + + if encoding := string(input.GetMetadata()["encoding"]); encoding != "binary/protobuf" { + return fmt.Errorf( + "system nexus envelope for operation %q must be encoded as binary/protobuf but got %q", + attrs.GetOperation(), encoding, + ) + } + + messageType := string(input.GetMetadata()["messageType"]) + if messageType == "" { + return fmt.Errorf( + "system nexus envelope for operation %q is missing the messageType metadata", + attrs.GetOperation(), + ) + } + + mt, err := protoregistry.GlobalTypes.FindMessageByName(protoreflect.FullName(messageType)) + if err != nil { + return fmt.Errorf( + "system nexus envelope for operation %q references unknown message type %q: %w", + attrs.GetOperation(), messageType, err, + ) + } + msg := mt.New().Interface() + + if err := proto.Unmarshal(input.GetData(), msg); err != nil { + return fmt.Errorf( + "failed to unmarshal system nexus envelope for operation %q: %w", + attrs.GetOperation(), err, + ) + } + + if err := visitPayloadsAndWait(ctx, options, msg, concState, msg); err != nil { + return err + } + + data, err := proto.Marshal(msg) + if err != nil { + return fmt.Errorf( + "failed to marshal system nexus envelope for operation %q: %w", + attrs.GetOperation(), err, + ) + } + input.Data = data + return nil +} diff --git a/proxy/system_nexus_test.go b/proxy/system_nexus_test.go new file mode 100644 index 00000000..5fa3cc10 --- /dev/null +++ b/proxy/system_nexus_test.go @@ -0,0 +1,198 @@ +package proxy + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" + command "go.temporal.io/api/command/v1" + common "go.temporal.io/api/common/v1" + sdk "go.temporal.io/api/sdk/v1" + workflowservice "go.temporal.io/api/workflowservice/v1" + "google.golang.org/protobuf/proto" +) + +const signalWithStartType = "temporal.api.workflowservice.v1.SignalWithStartWorkflowExecutionRequest" + +var signalWithStartRequest = &workflowservice.SignalWithStartWorkflowExecutionRequest{ + Namespace: "default", + WorkflowId: "wf-id", + SignalName: "my-signal", + Input: &common.Payloads{Payloads: []*common.Payload{ + {Data: []byte("workflow-input")}, + }}, + SignalInput: &common.Payloads{Payloads: []*common.Payload{ + {Data: []byte("signal-input")}, + }}, + Memo: &common.Memo{Fields: map[string]*common.Payload{ + "memo-key": {Data: []byte("memo-value")}, + }}, + Header: &common.Header{Fields: map[string]*common.Payload{ + "header-key": {Data: []byte("header-value")}, + }}, + SearchAttributes: &common.SearchAttributes{IndexedFields: map[string]*common.Payload{ + "sa-key": {Data: []byte("sa-value")}, + }}, + UserMetadata: &sdk.UserMetadata{ + Summary: &common.Payload{Data: []byte("summary-value")}, + Details: &common.Payload{Data: []byte("details-value")}, + }, +} + +var visitedSignalWithStartRequest = &workflowservice.SignalWithStartWorkflowExecutionRequest{ + Namespace: signalWithStartRequest.Namespace, + WorkflowId: signalWithStartRequest.WorkflowId, + SignalName: signalWithStartRequest.SignalName, + Input: &common.Payloads{Payloads: []*common.Payload{ + {Data: []byte("visited-workflow-input")}, + }}, + SignalInput: &common.Payloads{Payloads: []*common.Payload{ + {Data: []byte("visited-signal-input")}, + }}, + Memo: &common.Memo{Fields: map[string]*common.Payload{ + "memo-key": {Data: []byte("visited-memo-value")}, + }}, + Header: &common.Header{Fields: map[string]*common.Payload{ + "header-key": {Data: []byte("visited-header-value")}, + }}, + SearchAttributes: &common.SearchAttributes{IndexedFields: map[string]*common.Payload{ + "sa-key": {Data: []byte("visited-sa-value")}, + }}, + UserMetadata: &sdk.UserMetadata{ + Summary: &common.Payload{Data: []byte("visited-summary-value")}, + Details: &common.Payload{Data: []byte("visited-details-value")}, + }, +} + +// buildSystemNexusCommand builds a system Nexus command whose Input is +// a SignalWithStartWorkflowExecutionRequest. +// The encoding and messageType parameters represent the "encoding" and +// "messageType" metadata of the Payload. +func buildSystemNexusCommand(t *testing.T, encoding, messageType string) *command.Command { + t.Helper() + data, err := proto.Marshal(signalWithStartRequest) + require.NoError(t, err) + input := &common.Payload{ + Data: data, + Metadata: map[string][]byte{ + "encoding": []byte(encoding), + "messageType": []byte(messageType), + }, + } + attrs := &command.ScheduleNexusOperationCommandAttributes{ + Endpoint: "__temporal_system", + Service: "temporal.api.workflowservice.v1.WorkflowService", + Operation: "SignalWithStartWorkflowExecution", + Input: input, + } + return &command.Command{ + Attributes: &command.Command_ScheduleNexusOperationCommandAttributes{ + ScheduleNexusOperationCommandAttributes: attrs, + }, + } +} + +// collectVisitor rewrites each payload to "visited-" so callers can +// confirm write-back. +func collectVisitor(_ *VisitPayloadsContext, p []*common.Payload) ([]*common.Payload, error) { + out := make([]*common.Payload, len(p)) + for i, pl := range p { + out[i] = &common.Payload{Data: append([]byte("visited-"), pl.Data...)} + } + return out, nil +} + +func trivialVisitor(_ *VisitPayloadsContext, p []*common.Payload) ([]*common.Payload, error) { + return p, nil +} + +func decodeEnvelope(t *testing.T, input *common.Payload) *workflowservice.SignalWithStartWorkflowExecutionRequest { + t.Helper() + require.Equal(t, []byte("binary/protobuf"), input.GetMetadata()["encoding"]) + var req workflowservice.SignalWithStartWorkflowExecutionRequest + require.NoError(t, proto.Unmarshal(input.GetData(), &req)) + return &req +} + +//////////////////////// TESTS //////////////////////// + +func TestSystemNexusEnvelopeVisitsInnerPayloads(t *testing.T) { + cmd := buildSystemNexusCommand(t, "binary/protobuf", signalWithStartType) + + err := VisitPayloads(context.Background(), cmd, VisitPayloadsOptions{ + Visitor: collectVisitor, + }) + require.NoError(t, err) + + req := decodeEnvelope(t, cmd.GetScheduleNexusOperationCommandAttributes().Input) + require.True(t, proto.Equal(visitedSignalWithStartRequest, req)) +} + +func TestSystemNexusEnvelopeVisitsInnerPayloadsConcurrent(t *testing.T) { + cmd := buildSystemNexusCommand(t, "binary/protobuf", signalWithStartType) + + err := VisitPayloads(context.Background(), cmd, VisitPayloadsOptions{ + ConcurrencyLimit: 4, + Visitor: collectVisitor, + }) + require.NoError(t, err) + + req := decodeEnvelope(t, cmd.GetScheduleNexusOperationCommandAttributes().Input) + require.True(t, proto.Equal(visitedSignalWithStartRequest, req)) +} + +func TestSystemNexusEnvelopeRejectsNonProtoBinaryEncoding(t *testing.T) { + cmd := buildSystemNexusCommand(t, "json/protobuf", signalWithStartType) + + err := VisitPayloads(context.Background(), cmd, VisitPayloadsOptions{ + Visitor: trivialVisitor, + }) + require.Error(t, err) + require.ErrorContains(t, err, "binary/protobuf") +} + +func TestSystemNexusEnvelopeRejectsUnknownMessageType(t *testing.T) { + cmd := buildSystemNexusCommand(t, "binary/protobuf", "this isn't a valid message type") + + err := VisitPayloads(context.Background(), cmd, VisitPayloadsOptions{ + Visitor: trivialVisitor, + }) + require.Error(t, err) + require.ErrorContains(t, err, "unknown message type") +} + +func TestSystemNexusEnvelopeRejectsMissingMessageType(t *testing.T) { + cmd := buildSystemNexusCommand(t, "binary/protobuf", "") + + err := VisitPayloads(context.Background(), cmd, VisitPayloadsOptions{ + Visitor: trivialVisitor, + }) + require.Error(t, err) + require.ErrorContains(t, err, "missing") +} + +func TestNonSystemNexusInput(t *testing.T) { + cmd := &command.Command{ + Attributes: &command.Command_ScheduleNexusOperationCommandAttributes{ + ScheduleNexusOperationCommandAttributes: + &command.ScheduleNexusOperationCommandAttributes{ + Endpoint: "my-endpoint", + Service: "my-service", + Operation: "DoThing", + Input: &common.Payload{Data: []byte("user-payload")}, + }, + }, + } + + var seen []string + err := VisitPayloads(context.Background(), cmd, VisitPayloadsOptions{ + Visitor: func(ctx *VisitPayloadsContext, p []*common.Payload) ([]*common.Payload, error) { + for _, pl := range p { + seen = append(seen, string(pl.Data)) + } + return p, nil + }, + }) + require.NoError(t, err) + require.Equal(t, []string{"user-payload"}, seen) +}