Refactor nexus completion handler

Author: stephanos

chasm/lib/nexusoperation/completion_handler.go

@@ -0,0 +1,18 @@
+package nexusoperation
+
+import (
+	"context"
+
+	"github.com/nexus-rpc/sdk-go/nexus"
+	"go.temporal.io/server/common/nexus/nexusrpc"
+)
+
+type CompletionHandler struct{}
+
+func NewCompletionHandler() *CompletionHandler {
+	return &CompletionHandler{}
+}
+
+func (h *CompletionHandler) CompleteOperation(context.Context, *nexusrpc.CompletionRequest) error {
+	return nexus.NewHandlerErrorf(nexus.HandlerErrorTypeInternal, "CHASM nexus completion is not implemented")
+}

chasm/lib/nexusoperation/fx.go

@@ -32,6 +32,7 @@ var Module = fx.Module(
 	fx.Invoke(endpointRegistryLifetimeHooks),
 	fx.Provide(defaultNexusTransportProvider),
 	fx.Provide(clientProviderFactory),
+	fx.Provide(NewCompletionHandler),
 	fx.Provide(newCancellationBackoffTaskHandler),
 	fx.Provide(newCancellationInvocationTaskHandler),
 	fx.Provide(newOperationBackoffTaskHandler),

components/nexusoperations/frontend/fx.go

@@ -1,16 +1,8 @@
 package frontend
 
 import (
-	"net/http"
-
-	"github.com/gorilla/mux"
 	"go.temporal.io/server/common/dynamicconfig"
-	"go.temporal.io/server/common/headers"
-	"go.temporal.io/server/common/log"
-	"go.temporal.io/server/common/metrics"
 	commonnexus "go.temporal.io/server/common/nexus"
-	"go.temporal.io/server/common/nexus/nexusrpc"
-	"go.temporal.io/server/common/rpc"
 	"go.temporal.io/server/components/nexusoperations"
 	"go.uber.org/fx"
 )
@@ -19,7 +11,7 @@ var Module = fx.Module(
 	"component.nexusoperations.frontend",
 	fx.Provide(ConfigProvider),
 	fx.Provide(commonnexus.NewCallbackTokenGenerator),
-	fx.Invoke(RegisterHTTPHandler),
+	fx.Provide(NewCompletionHandler),
 )
 
 func ConfigProvider(coll *dynamicconfig.Collection) *Config {
@@ -29,29 +21,3 @@ func ConfigProvider(coll *dynamicconfig.Collection) *Config {
 		MaxOperationTokenLength:       nexusoperations.MaxOperationTokenLength.Get(coll),
 	}
 }
-
-func RegisterHTTPHandler(options HandlerOptions, logger log.Logger, router *mux.Router) {
-	h := nexusrpc.NewCompletionHTTPHandler(nexusrpc.CompletionHandlerOptions{
-		Handler: &completionHandler{
-			options,
-			headers.NewDefaultVersionChecker(),
-			options.MetricsHandler.Counter(metrics.NexusCompletionRequestPreProcessErrors.Name()),
-		},
-		Logger:     log.NewSlogLogger(logger),
-		Serializer: commonnexus.PayloadSerializer,
-	})
-	router.Path("/" + commonnexus.RouteCompletionCallback.Representation()).HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		// Limit the request body to max allowed Payload size.
-		// Content headers are transformed to Payload metadata and contribute to the Payload size as well. A separate
-		// limit is enforced on top of this in the CompleteOperation method.
-		r.Body = http.MaxBytesReader(w, r.Body, rpc.MaxNexusAPIRequestBodyBytes)
-		h.ServeHTTP(w, r)
-	})
-	router.Path(commonnexus.PathCompletionCallbackNoIdentifier).HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		// Limit the request body to max allowed Payload size.
-		// Content headers are transformed to Payload metadata and contribute to the Payload size as well. A separate
-		// limit is enforced on top of this in the CompleteOperation method.
-		r.Body = http.MaxBytesReader(w, r.Body, rpc.MaxNexusAPIRequestBodyBytes)
-		h.ServeHTTP(w, r)
-	})
-}

components/nexusoperations/frontend/handler.go

@@ -80,6 +80,20 @@ type completionHandler struct {
 	preProcessErrorsCounter metrics.CounterIface
 }
 
+type CompletionHandler struct {
+	*completionHandler
+}
+
+func NewCompletionHandler(options HandlerOptions) *CompletionHandler {
+	return &CompletionHandler{
+		completionHandler: &completionHandler{
+			HandlerOptions:          options,
+			clientVersionChecker:    headers.NewDefaultVersionChecker(),
+			preProcessErrorsCounter: options.MetricsHandler.Counter(metrics.NexusCompletionRequestPreProcessErrors.Name()),
+		},
+	}
+}
+
 // CompleteOperation implements nexus.CompletionHandler.
 // nolint:revive // (cyclomatic complexity) This function is long but the complexity is justified.
 func (h *completionHandler) CompleteOperation(ctx context.Context, r *nexusrpc.CompletionRequest) (retErr error) {
@@ -89,63 +103,16 @@ func (h *completionHandler) CompleteOperation(ctx context.Context, r *nexusrpc.C
 		h.Logger.Error("failed to decode callback token", tag.Error(err))
 		return nexus.NewHandlerErrorf(nexus.HandlerErrorTypeBadRequest, "invalid callback token")
 	}
-
 	completion, err := h.CallbackTokenGenerator.DecodeCompletion(token)
 	if err != nil {
 		h.Logger.Error("failed to decode completion from token", tag.Error(err))
 		return nexus.NewHandlerErrorf(nexus.HandlerErrorTypeBadRequest, "invalid callback token")
 	}
-	ns, err := h.NamespaceRegistry.GetNamespaceByID(namespace.ID(completion.NamespaceId))
+	rCtx, logger, ctx, err := h.newRequestContext(ctx, r, startTime, completion.GetNamespaceId(), completion.GetWorkflowId(), completion.GetRunId())
 	if err != nil {
-		h.Logger.Error("failed to get namespace for nexus completion request", tag.WorkflowNamespaceID(completion.NamespaceId), tag.Error(err))
-		h.preProcessErrorsCounter.Record(1)
-		var nfe *serviceerror.NamespaceNotFound
-		if errors.As(err, &nfe) {
-			return nexus.NewHandlerErrorf(nexus.HandlerErrorTypeNotFound, "namespace %q not found", completion.NamespaceId)
-		}
-		return commonnexus.ConvertGRPCError(err, false)
-	}
-	logger := log.With(
-		h.Logger,
-		tag.WorkflowNamespace(ns.Name().String()),
-		tag.WorkflowID(completion.GetWorkflowId()),
-		tag.WorkflowRunID(completion.GetRunId()),
-	)
-	rCtx := &requestContext{
-		completionHandler: h,
-		namespace:         ns,
-		workflowID:        completion.GetWorkflowId(),
-		logger:            log.With(h.Logger, tag.WorkflowNamespace(ns.Name().String())),
-		metricsHandler:    h.MetricsHandler.WithTags(metrics.NamespaceTag(ns.Name().String())),
-		metricsHandlerForInterceptors: h.MetricsHandler.WithTags(
-			metrics.OperationTag(methodNameForMetrics),
-			metrics.NamespaceTag(ns.Name().String()),
-		),
-		requestStartTime: startTime,
-	}
-	if r.HTTPRequest.Header != nil {
-		rCtx.originalHeaders = r.HTTPRequest.Header.Clone()
+		return err
 	}
-	ctx = rCtx.augmentContext(ctx, r.HTTPRequest.Header)
 	defer rCtx.capturePanicAndRecordMetrics(&ctx, &retErr)
-	if r.HTTPRequest.URL.Path != commonnexus.PathCompletionCallbackNoIdentifier {
-		nsNameEscaped := commonnexus.RouteCompletionCallback.Deserialize(mux.Vars(r.HTTPRequest))
-		nsName, err := url.PathUnescape(nsNameEscaped)
-		if err != nil {
-			h.Logger.Error("failed to extract namespace from request", tag.Error(err))
-			h.preProcessErrorsCounter.Record(1)
-			return nexus.NewHandlerErrorf(nexus.HandlerErrorTypeBadRequest, "invalid URL")
-		}
-		if nsName != ns.Name().String() {
-			logger.Error(
-				"namespace ID in token doesn't match the token",
-				tag.WorkflowNamespaceID(ns.ID().String()),
-				tag.Error(err),
-				tag.String("completion-namespace-id", completion.GetNamespaceId()),
-			)
-			return nexus.NewHandlerErrorf(nexus.HandlerErrorTypeBadRequest, "invalid callback token")
-		}
-	}
 
 	if err := rCtx.interceptRequest(ctx, r); err != nil {
 		var notActiveErr *serviceerror.NamespaceNotActive
@@ -154,7 +121,7 @@ func (h *completionHandler) CompleteOperation(ctx context.Context, r *nexusrpc.C
 		}
 		return err
 	}
-	tokenLimit := h.Config.MaxOperationTokenLength(ns.Name().String())
+	tokenLimit := h.Config.MaxOperationTokenLength(rCtx.namespace.Name().String())
 	if len(r.OperationToken) > tokenLimit {
 		return nexus.NewHandlerErrorf(nexus.HandlerErrorTypeBadRequest, "operation token length exceeds allowed limit (%d/%d)", len(r.OperationToken), tokenLimit)
 	}
@@ -201,8 +168,8 @@ func (h *completionHandler) CompleteOperation(ctx context.Context, r *nexusrpc.C
 			logger.Error("cannot deserialize payload from completion result", tag.Error(err))
 			return nexus.NewHandlerErrorf(nexus.HandlerErrorTypeBadRequest, "invalid result content")
 		}
-		if result.Size() > h.Config.PayloadSizeLimit(ns.Name().String()) {
-			logger.Error("payload size exceeds error limit for Nexus CompleteOperation request", tag.WorkflowNamespace(ns.Name().String()))
+		if result.Size() > h.Config.PayloadSizeLimit(rCtx.namespace.Name().String()) {
+			logger.Error("payload size exceeds error limit for Nexus CompleteOperation request", tag.WorkflowNamespace(rCtx.namespace.Name().String()))
 			return nexus.NewHandlerErrorf(nexus.HandlerErrorTypeBadRequest, "result exceeds size limit")
 		}
 		hr.Outcome = &historyservice.CompleteNexusOperationRequest_Success{
@@ -229,6 +196,70 @@ func (h *completionHandler) CompleteOperation(ctx context.Context, r *nexusrpc.C
 	return nil
 }
 
+func (h *completionHandler) newRequestContext(
+	ctx context.Context,
+	r *nexusrpc.CompletionRequest,
+	startTime time.Time,
+	namespaceID string,
+	workflowID string,
+	runID string,
+) (*requestContext, log.Logger, context.Context, error) {
+	ns, err := h.NamespaceRegistry.GetNamespaceByID(namespace.ID(namespaceID))
+	if err != nil {
+		h.Logger.Error("failed to get namespace for nexus completion request", tag.WorkflowNamespaceID(namespaceID), tag.Error(err))
+		h.preProcessErrorsCounter.Record(1)
+		var nfe *serviceerror.NamespaceNotFound
+		if errors.As(err, &nfe) {
+			return nil, nil, nil, nexus.NewHandlerErrorf(nexus.HandlerErrorTypeNotFound, "namespace %q not found", namespaceID)
+		}
+		return nil, nil, nil, commonnexus.ConvertGRPCError(err, false)
+	}
+
+	logger := log.With(
+		h.Logger,
+		tag.WorkflowNamespace(ns.Name().String()),
+		tag.WorkflowID(workflowID),
+		tag.WorkflowRunID(runID),
+	)
+	rCtx := &requestContext{
+		completionHandler: h,
+		namespace:         ns,
+		workflowID:        workflowID,
+		logger:            log.With(h.Logger, tag.WorkflowNamespace(ns.Name().String())),
+		metricsHandler:    h.MetricsHandler.WithTags(metrics.NamespaceTag(ns.Name().String())),
+		metricsHandlerForInterceptors: h.MetricsHandler.WithTags(
+			metrics.OperationTag(methodNameForMetrics),
+			metrics.NamespaceTag(ns.Name().String()),
+		),
+		requestStartTime: startTime,
+	}
+	if r.HTTPRequest.Header != nil {
+		rCtx.originalHeaders = r.HTTPRequest.Header.Clone()
+	}
+	ctx = rCtx.augmentContext(ctx, r.HTTPRequest.Header)
+
+	if r.HTTPRequest.URL.Path != commonnexus.PathCompletionCallbackNoIdentifier {
+		nsNameEscaped := commonnexus.RouteCompletionCallback.Deserialize(mux.Vars(r.HTTPRequest))
+		nsName, err := url.PathUnescape(nsNameEscaped)
+		if err != nil {
+			h.Logger.Error("failed to extract namespace from request", tag.Error(err))
+			h.preProcessErrorsCounter.Record(1)
+			return nil, nil, nil, nexus.NewHandlerErrorf(nexus.HandlerErrorTypeBadRequest, "invalid URL")
+		}
+		if nsName != ns.Name().String() {
+			logger.Error(
+				"namespace ID in token doesn't match the token",
+				tag.WorkflowNamespaceID(ns.ID().String()),
+				tag.Error(err),
+				tag.String("completion-namespace-id", namespaceID),
+			)
+			return nil, nil, nil, nexus.NewHandlerErrorf(nexus.HandlerErrorTypeBadRequest, "invalid callback token")
+		}
+	}
+
+	return rCtx, logger, ctx, nil
+}
+
 func (h *completionHandler) forwardCompleteOperation(ctx context.Context, r *nexusrpc.CompletionRequest, rCtx *requestContext) error {
 	client, err := h.ForwardingClients.Get(rCtx.namespace.ActiveClusterName(rCtx.workflowID))
 	if err != nil {
@@ -320,6 +351,10 @@ type requestContext struct {
 	originalHeaders               http.Header
 }
 
+type RequestContext struct {
+	*requestContext
+}
+
 func (c *requestContext) augmentContext(ctx context.Context, header http.Header) context.Context {
 	ctx = metrics.AddMetricsContext(ctx)
 	ctx = interceptor.AddTelemetryContext(ctx, c.metricsHandlerForInterceptors)

service/frontend/fx.go

@@ -107,7 +107,8 @@ var Module = fx.Options(
 	fx.Provide(OperatorHandlerProvider),
 	fx.Provide(NewVersionChecker),
 	fx.Provide(ServiceResolverProvider),
-	fx.Invoke(RegisterNexusHTTPHandler),
+	fx.Invoke(RegisterNexusOperationHTTPHandler),
+	fx.Invoke(RegisterNexusCompletionHTTPHandler),
 	fx.Invoke(RegisterOpenAPIHTTPHandler),
 	fx.Provide(HTTPAPIServerProvider),
 	fx.Provide(NewServiceProvider),
@@ -862,7 +863,7 @@ func HandlerProvider(
 	return wfHandler
 }
 
-func RegisterNexusHTTPHandler(
+func RegisterNexusOperationHTTPHandler(
 	serviceConfig *Config,
 	serviceName primitives.ServiceName,
 	matchingClient resource.MatchingClient,
@@ -883,7 +884,7 @@ func RegisterNexusHTTPHandler(
 	router *mux.Router,
 	httpTraceProvider nexus.HTTPClientTraceProvider,
 ) {
-	h := NewNexusHTTPHandler(
+	h := NewNexusOperationHTTPHandler(
 		serviceConfig,
 		matchingClient,
 		metricsHandler,

service/frontend/nexus_completion_http_handler.go

@@ -0,0 +1,78 @@
+package frontend
+
+import (
+	"context"
+	"net/http"
+
+	"github.com/gorilla/mux"
+	"github.com/nexus-rpc/sdk-go/nexus"
+	persistencespb "go.temporal.io/server/api/persistence/v1"
+	chasmnexusoperation "go.temporal.io/server/chasm/lib/nexusoperation"
+	"go.temporal.io/server/common/log"
+	commonnexus "go.temporal.io/server/common/nexus"
+	"go.temporal.io/server/common/nexus/nexusrpc"
+	"go.temporal.io/server/common/rpc"
+	nexusfrontend "go.temporal.io/server/components/nexusoperations/frontend"
+)
+
+type completionRoutingHandler struct {
+	callbackTokenGenerator *commonnexus.CallbackTokenGenerator
+	hsm                    nexusrpc.CompletionHandler
+	chasm                  nexusrpc.CompletionHandler
+}
+
+func (h *completionRoutingHandler) CompleteOperation(ctx context.Context, r *nexusrpc.CompletionRequest) error {
+	token, err := commonnexus.DecodeCallbackToken(r.HTTPRequest.Header.Get(commonnexus.CallbackTokenHeader))
+	if err != nil {
+		return nexus.NewHandlerErrorf(nexus.HandlerErrorTypeBadRequest, "invalid callback token")
+	}
+
+	completion, err := h.callbackTokenGenerator.DecodeCompletion(token)
+	if err != nil {
+		return nexus.NewHandlerErrorf(nexus.HandlerErrorTypeBadRequest, "invalid callback token")
+	}
+
+	// TODO: need to handle migration between HSM and CHASM
+
+	if len(completion.GetComponentRef()) > 0 {
+		ref := &persistencespb.ChasmComponentRef{}
+		if err := ref.Unmarshal(completion.GetComponentRef()); err != nil {
+			return nexus.NewHandlerErrorf(nexus.HandlerErrorTypeBadRequest, "invalid callback token")
+		}
+		return h.chasm.CompleteOperation(ctx, r)
+	}
+	return h.hsm.CompleteOperation(ctx, r)
+}
+
+func RegisterNexusCompletionHTTPHandler(
+	hsmHandler *nexusfrontend.CompletionHandler,
+	chasmHandler *chasmnexusoperation.CompletionHandler,
+	callbackTokenGenerator *commonnexus.CallbackTokenGenerator,
+	logger log.Logger,
+	router *mux.Router,
+) {
+	h := nexusrpc.NewCompletionHTTPHandler(nexusrpc.CompletionHandlerOptions{
+		Handler: &completionRoutingHandler{
+			callbackTokenGenerator: callbackTokenGenerator,
+			hsm:                    hsmHandler,
+			chasm:                  chasmHandler,
+		},
+		Logger:     log.NewSlogLogger(logger),
+		Serializer: commonnexus.PayloadSerializer,
+	})
+
+	router.Path("/" + commonnexus.RouteCompletionCallback.Representation()).HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// Limit the request body to max allowed Payload size.
+		// Content headers are transformed to Payload metadata and contribute to the Payload size as well. A separate
+		// limit is enforced on top of this in the CompleteOperation method.
+		r.Body = http.MaxBytesReader(w, r.Body, rpc.MaxNexusAPIRequestBodyBytes)
+		h.ServeHTTP(w, r)
+	})
+	router.Path(commonnexus.PathCompletionCallbackNoIdentifier).HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// Limit the request body to max allowed Payload size.
+		// Content headers are transformed to Payload metadata and contribute to the Payload size as well. A separate
+		// limit is enforced on top of this in the CompleteOperation method.
+		r.Body = http.MaxBytesReader(w, r.Body, rpc.MaxNexusAPIRequestBodyBytes)
+		h.ServeHTTP(w, r)
+	})
+}