#817 use PreCall and PostCall plugins for RegisterFunction

smallnest · Sep 3, 2023 · e43aa14 · e43aa14
1 parent 883cc59
commit e43aa14
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 11 deletions.
diff --git a/server/plugin.go b/server/plugin.go
@@ -31,7 +31,7 @@ type PluginContainer interface {
 
 	DoPreHandleRequest(ctx context.Context, req *protocol.Message) error
 	DoPreCall(ctx context.Context, serviceName, methodName string, args interface{}) (interface{}, error)
-	DoPostCall(ctx context.Context, serviceName, methodName string, args, reply interface{}) (interface{}, error)
+	DoPostCall(ctx context.Context, serviceName, methodName string, args, reply interface{}, err error) (interface{}, error)
 
 	DoPreWriteResponse(context.Context, *protocol.Message, *protocol.Message, error) error
 	DoPostWriteResponse(context.Context, *protocol.Message, *protocol.Message, error) error
@@ -96,7 +96,7 @@ type (
 	}
 
 	PostCallPlugin interface {
-		PostCall(ctx context.Context, serviceName, methodName string, args, reply interface{}) (interface{}, error)
+		PostCall(ctx context.Context, serviceName, methodName string, args, reply interface{}, err error) (interface{}, error)
 	}
 
 	// PreWriteResponsePlugin represents .
@@ -314,18 +314,18 @@ func (p *pluginContainer) DoPreCall(ctx context.Context, serviceName, methodName
 }
 
 // DoPostCall invokes PostCallPlugin plugin.
-func (p *pluginContainer) DoPostCall(ctx context.Context, serviceName, methodName string, args, reply interface{}) (interface{}, error) {
-	var err error
+func (p *pluginContainer) DoPostCall(ctx context.Context, serviceName, methodName string, args, reply interface{}, err error) (interface{}, error) {
+	var e error
 	for i := range p.plugins {
 		if plugin, ok := p.plugins[i].(PostCallPlugin); ok {
-			reply, err = plugin.PostCall(ctx, serviceName, methodName, args, reply)
-			if err != nil {
-				return reply, err
+			reply, e = plugin.PostCall(ctx, serviceName, methodName, args, reply, err)
+			if e != nil {
+				return reply, e
 			}
 		}
 	}
 
-	return reply, err
+	return reply, e
 }
 
 // DoPreWriteResponse invokes PreWriteResponse plugin.

diff --git a/server/server.go b/server/server.go
@@ -722,9 +722,7 @@ func (s *Server) handleRequest(ctx context.Context, req *protocol.Message) (res
 		err = service.call(ctx, mtype, reflect.ValueOf(argv), reflect.ValueOf(replyv))
 	}
 
-	if err == nil {
-		replyv, err = s.Plugins.DoPostCall(ctx, serviceName, methodName, argv, replyv)
-	}
+	replyv, err = s.Plugins.DoPostCall(ctx, serviceName, methodName, argv, replyv, err)
 
 	// return argc to object pool
 	reflectTypePools.Put(mtype.ArgType, argv)
@@ -795,13 +793,21 @@ func (s *Server) handleRequestForFunction(ctx context.Context, req *protocol.Mes
 	}
 
 	replyv := reflectTypePools.Get(mtype.ReplyType)
+	argv, err = s.Plugins.DoPreCall(ctx, serviceName, methodName, argv)
+	if err != nil {
+		// return reply to object pool
+		reflectTypePools.Put(mtype.ReplyType, replyv)
+		return s.handleError(res, err)
+	}
 
 	if mtype.ArgType.Kind() != reflect.Ptr {
 		err = service.callForFunction(ctx, mtype, reflect.ValueOf(argv).Elem(), reflect.ValueOf(replyv))
 	} else {
 		err = service.callForFunction(ctx, mtype, reflect.ValueOf(argv), reflect.ValueOf(replyv))
 	}
 
+	replyv, err = s.Plugins.DoPostCall(ctx, serviceName, methodName, argv, replyv, err)
+
 	reflectTypePools.Put(mtype.ArgType, argv)
 
 	if err != nil {