feat: customize error codec

pkg/builder.go

@@ -18,6 +18,7 @@ import (
 	"github.com/topfreegames/pitaya/v3/pkg/serialize"
 	"github.com/topfreegames/pitaya/v3/pkg/service"
 	"github.com/topfreegames/pitaya/v3/pkg/session"
+	"github.com/topfreegames/pitaya/v3/pkg/util"
 	"github.com/topfreegames/pitaya/v3/pkg/worker"
 )
 
@@ -29,7 +30,8 @@ type Builder struct {
 	DieChan          chan bool
 	PacketDecoder    codec.PacketDecoder
 	PacketEncoder    codec.PacketEncoder
-	MessageEncoder   *message.MessagesEncoder
+	MessageEncoder   message.Encoder
+	MessageDecoder   message.Decoder
 	Serializer       serialize.Serializer
 	Router           *router.Router
 	RPCClient        cluster.RPCClient
@@ -43,6 +45,7 @@ type Builder struct {
 	Worker           *worker.Worker
 	RemoteHooks      *pipeline.RemoteHooks
 	HandlerHooks     *pipeline.HandlerHooks
+	ErrWrapper       serialize.ErrorWrapper
 }
 
 // PitayaBuilder Builder interface
@@ -147,13 +150,16 @@ func NewBuilder(isFrontend bool,
 	}
 
 	return &Builder{
-		acceptors:        []acceptor.Acceptor{},
-		postBuildHooks:   make([]func(app Pitaya), 0),
-		Config:           config,
-		DieChan:          dieChan,
-		PacketDecoder:    codec.NewPomeloPacketDecoder(),
-		PacketEncoder:    codec.NewPomeloPacketEncoder(),
-		MessageEncoder:   message.NewMessagesEncoder(config.Handler.Messages.Compression),
+		acceptors:      []acceptor.Acceptor{},
+		postBuildHooks: make([]func(app Pitaya), 0),
+		Config:         config,
+		DieChan:        dieChan,
+		PacketDecoder:  codec.NewPomeloPacketDecoder(),
+		PacketEncoder:  codec.NewPomeloPacketEncoder(),
+		MessageEncoder: message.NewMessagesEncoder(config.Handler.Messages.Compression),
+		MessageDecoder: message.NewMessagesDecoder(config.Handler.Messages.Compression),
+		//Serializer:       json.NewSerializer(),
+		//MessageEncoder:   message.NewMessagesEncoder(config.Handler.Messages.Compression),
 		Serializer:       serializer,
 		Router:           router.New(),
 		RPCClient:        rpcClient,
@@ -239,6 +245,7 @@ func (builder *Builder) Build() Pitaya {
 		builder.MetricsReporters,
 		builder.HandlerHooks,
 		handlerPool,
+		builder.MessageDecoder,
 	)
 
 	app := NewApp(
@@ -260,6 +267,10 @@ func (builder *Builder) Build() Pitaya {
 		builder.Config,
 	)
 
+	if builder.ErrWrapper != nil {
+		util.DefaultErrWrapper = builder.ErrWrapper
+	}
+
 	for _, postBuildHook := range builder.postBuildHooks {
 		postBuildHook(app)
 	}

pkg/component/method.go

@@ -114,7 +114,7 @@ func isHandlerMethod(method reflect.Method) bool {
 	return true
 }
 
-func suitableRemoteMethods(typ reflect.Type, nameFunc func(string) string) map[string]*Remote {
+func suitableRemoteMethods(typ reflect.Type, nameFunc func(string) string, proxyFunc ...*MethodProxyFuncOption) map[string]*Remote {
 	methods := make(map[string]*Remote)
 	for m := 0; m < typ.NumMethod(); m++ {
 		method := typ.Method(m)
@@ -125,9 +125,13 @@ func suitableRemoteMethods(typ reflect.Type, nameFunc func(string) string) map[s
 			if nameFunc != nil {
 				mn = nameFunc(mn)
 			}
+
+			// setup proxy
+			proxy := setUpProxy(proxyFunc, mn)
 			methods[mn] = &Remote{
-				Method:  method,
-				HasArgs: method.Type.NumIn() == 3,
+				Method:      method,
+				HasArgs:     method.Type.NumIn() == 3,
+				MethodProxy: proxy,
 			}
 			if mt.NumIn() == 3 {
 				methods[mn].Type = mt.In(2)
@@ -137,7 +141,7 @@ func suitableRemoteMethods(typ reflect.Type, nameFunc func(string) string) map[s
 	return methods
 }
 
-func suitableHandlerMethods(typ reflect.Type, nameFunc func(string) string) map[string]*Handler {
+func suitableHandlerMethods(typ reflect.Type, nameFunc func(string) string, proxyFunc ...*MethodProxyFuncOption) map[string]*Handler {
 	methods := make(map[string]*Handler)
 	for m := 0; m < typ.NumMethod(); m++ {
 		method := typ.Method(m)
@@ -158,10 +162,14 @@ func suitableHandlerMethods(typ reflect.Type, nameFunc func(string) string) map[
 			} else {
 				msgType = message.Request
 			}
+
+			// setup proxy
+			proxy := setUpProxy(proxyFunc, mn)
 			handler := &Handler{
 				Method:      method,
 				IsRawArg:    raw,
 				MessageType: msgType,
+				MethodProxy: proxy,
 			}
 			if mt.NumIn() == 3 {
 				handler.Type = mt.In(2)

pkg/component/options.go

@@ -20,10 +20,15 @@
 
 package component
 
+import (
+	"reflect"
+)
+
 type (
 	options struct {
-		name     string              // component name
-		nameFunc func(string) string // rename handler name
+		name       string              // component name
+		nameFunc   func(string) string // rename handler name
+		proxyChain *MethodProxyFuncOption
 	}
 
 	// Option used to customize handler
@@ -44,3 +49,26 @@ func WithNameFunc(fn func(string) string) Option {
 		opt.nameFunc = fn
 	}
 }
+
+// WithProxyFunc used to customize proxy function. To calling original method,
+// use `component.DefaultHandlerMethodInvoke` for shortcut
+func WithProxyFunc(methodOrName any, fn ...MethodProxyFunc) Option {
+	return func(opt *options) {
+		name := ""
+		typ := reflect.TypeOf(methodOrName)
+		switch typ.Kind() {
+		case reflect.String:
+			name = methodOrName.(string)
+		case reflect.Func:
+			name = reflect.TypeOf(methodOrName).Name()
+		default:
+			return
+		}
+		if opt.proxyChain == nil {
+			opt.proxyChain = &MethodProxyFuncOption{
+				Chains: make(map[string][]MethodProxyFunc),
+			}
+		}
+		opt.proxyChain.Chains[name] = append(opt.proxyChain.Chains[name], fn...)
+	}
+}

pkg/component/proxy.go

@@ -0,0 +1,161 @@
+package component
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"github.com/topfreegames/pitaya/v3/pkg/constants"
+	"github.com/topfreegames/pitaya/v3/pkg/logger"
+	"github.com/topfreegames/pitaya/v3/pkg/logger/interfaces"
+	"github.com/topfreegames/pitaya/v3/pkg/util"
+	"math"
+	"reflect"
+	"runtime/debug"
+	"strconv"
+)
+
+type (
+	// MethodProxy is a function that will be called when a handler is invoked.
+	// origin is the original method of the handler.
+	// args are the arguments passed to the method,including:
+	// method Receiver at args[0],
+	// context.Context at args[1],
+	// request struct or raw bytes at args[2] if exists.
+	// return proxy func accepts same arguments and returns same values as origin.
+	MethodProxy        func(origin reflect.Method, args []reflect.Value) (rets any, err error)
+	MethodProxyFunc    func(ctx *MethodProxyContext)
+	MethodProxyContext struct {
+		index          int8
+		handlers       []MethodProxyFunc
+		OriginArgs     []reflect.Value
+		OriginReceiver reflect.Value
+		OriginMethod   reflect.Method
+		InCtx          context.Context
+		InMsg          any
+		OutMsg         any
+		OutErr         error
+		data           map[string]any
+	}
+	MethodProxyFuncOption struct {
+		Chains map[string][]MethodProxyFunc // method proxy chains, method name -> middleware chain
+	}
+)
+
+const abortIndex = math.MaxInt8 >> 1
+
+var DefaultHandlerMethodInvoke = util.Pcall
+
+func setUpProxy(proxyFunc []*MethodProxyFuncOption, mn string) (proxy MethodProxy) {
+	if len(proxyFunc) == 0 || proxyFunc[0] == nil || len(proxyFunc[0].Chains) == 0 {
+		return
+	}
+	chain := proxyFunc[0].Chains[mn]
+	if len(chain) == 0 {
+		return
+	}
+	proxy = func(origin reflect.Method, args []reflect.Value) (rets any, err error) {
+		defer func() { // TODO use util.PcallFunc
+			if rec := recover(); rec != nil {
+				// Try to use logger from context here to help trace error cause
+				stackTrace := debug.Stack()
+				stackTraceAsRawStringLiteral := strconv.Quote(string(stackTrace))
+				log := getLoggerFromArgs(args)
+				log.Errorf("panic - pitaya/dispatch/proxy: methodName=%s panicData=%v stackTrace=%s", origin.Name, rec, stackTraceAsRawStringLiteral)
+
+				if s, ok := rec.(string); ok {
+					err = errors.New(s)
+				} else {
+					err = fmt.Errorf("rpc call internal error - %s: %v", origin.Name, rec)
+				}
+			}
+		}()
+
+		ctx := NewMethodProxyContext(origin, args, chain)
+		ctx.Next()
+		if ctx.IsDone() {
+			rets = ctx.OutMsg
+			err = ctx.OutErr
+			return
+		}
+		return
+	}
+	return
+}
+
+func getLoggerFromArgs(args []reflect.Value) interfaces.Logger {
+	for _, a := range args {
+		if !a.IsValid() {
+			continue
+		}
+		if ctx, ok := a.Interface().(context.Context); ok {
+			logVal := ctx.Value(constants.LoggerCtxKey)
+			if logVal != nil {
+				log := logVal.(interfaces.Logger)
+				return log
+			}
+		}
+	}
+	return logger.Log
+}
+
+func NewMethodProxyContext(origin reflect.Method, args []reflect.Value, fns []MethodProxyFunc) *MethodProxyContext {
+	ret := &MethodProxyContext{
+		OriginMethod: origin,
+		OriginArgs:   args,
+		data:         make(map[string]any),
+		handlers:     fns,
+		index:        -1,
+	}
+	if len(args) > 0 {
+		ret.OriginReceiver = args[0]
+	}
+	if len(args) > 1 {
+		ret.InCtx = args[1].Interface().(context.Context)
+	}
+	if len(args) > 2 {
+		ret.InMsg = args[2].Interface()
+	}
+	return ret
+}
+
+func (m *MethodProxyContext) Next() {
+	m.index++
+	for m.index < int8(len(m.handlers)) {
+		m.handlers[m.index](m)
+		m.index++
+	}
+}
+
+func (m *MethodProxyContext) IsDone() bool {
+	return m.index >= int8(len(m.handlers)) || m.index >= abortIndex
+}
+
+func (m *MethodProxyContext) Abort() {
+	m.index = abortIndex
+}
+
+func (m *MethodProxyContext) SetData(key string, value any) {
+	m.data[key] = value
+}
+
+func (m *MethodProxyContext) GetData(key string) any {
+	return m.data[key]
+}
+
+func (m *MethodProxyContext) GetRequest() any {
+	return m.InMsg
+}
+
+func (m *MethodProxyContext) GetContext() any {
+	return m.InCtx
+}
+
+func (m *MethodProxyContext) ResponseWithMsg(msg any) {
+	m.OutMsg = msg
+	m.Abort()
+}
+
+func (m *MethodProxyContext) ResponseWithError(err error) {
+	m.OutErr = err
+	m.Abort()
+}

pkg/component/proxy_test.go

@@ -0,0 +1,40 @@
+package component
+
+import (
+	"context"
+	"fmt"
+	"reflect"
+	"testing"
+)
+
+type testHandler struct {
+}
+
+type testReq struct {
+}
+
+type testResp struct {
+}
+
+type testNotify struct {
+}
+
+func (t *testHandler) TestFunc(ctx context.Context, req *testReq) (*testResp, error) {
+	fmt.Println("call TestFunc")
+	return nil, nil
+}
+
+func (t *testHandler) TestNotify(ctx context.Context, notice *testNotify) {
+	fmt.Println("call TestNotify")
+}
+
+func TestNewMethodProxyContext(t *testing.T) {
+	h := &testHandler{}
+	typ := reflect.TypeOf(h)
+	for i := range typ.NumMethod() {
+		f := typ.Method(i)
+
+		ctx := NewMethodProxyContext(f, nil, nil)
+		ctx.Next()
+	}
+}