From 55f7cfcebae4511ba12049b510eb6539e490da2d Mon Sep 17 00:00:00 2001
From: Hieu Vu <72878483+hieuvubk@users.noreply.github.com>
Date: Tue, 5 Nov 2024 14:39:32 +0700
Subject: [PATCH] refactor(serevr/v2/cometbft): update `RegisterQueryHandlers`
 and GRPC queries (#22403)

---
 runtime/v2/manager.go      | 13 ++++++++++---
 server/v2/cometbft/abci.go |  9 ++++++---
 2 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/runtime/v2/manager.go b/runtime/v2/manager.go
index e2e90c27f808..9e99a2b08c4f 100644
--- a/runtime/v2/manager.go
+++ b/runtime/v2/manager.go
@@ -642,7 +642,7 @@ func registerServices[T transaction.Tx](s appmodulev2.AppModule, app *App[T], re
 
 	// if module implements register msg handlers
 	if module, ok := s.(appmodulev2.HasMsgHandlers); ok {
-		wrapper := stfRouterWrapper{stfRouter: app.msgRouterBuilder}
+		wrapper := newStfRouterWrapper(app.msgRouterBuilder)
 		module.RegisterMsgHandlers(&wrapper)
 		if wrapper.error != nil {
 			return fmt.Errorf("unable to register handlers: %w", wrapper.error)
@@ -651,7 +651,7 @@ func registerServices[T transaction.Tx](s appmodulev2.AppModule, app *App[T], re
 
 	// if module implements register query handlers
 	if module, ok := s.(appmodulev2.HasQueryHandlers); ok {
-		wrapper := stfRouterWrapper{stfRouter: app.queryRouterBuilder}
+		wrapper := newStfRouterWrapper(app.queryRouterBuilder)
 		module.RegisterQueryHandlers(&wrapper)
 
 		for path, handler := range wrapper.handlers {
@@ -842,6 +842,13 @@ type stfRouterWrapper struct {
 	handlers map[string]appmodulev2.Handler
 }
 
+func newStfRouterWrapper(stfRouterBuilder *stf.MsgRouterBuilder) stfRouterWrapper {
+	wrapper := stfRouterWrapper{stfRouter: stfRouterBuilder}
+	wrapper.error = nil
+	wrapper.handlers = map[string]appmodulev2.Handler{}
+	return wrapper
+}
+
 func (s *stfRouterWrapper) RegisterHandler(handler appmodulev2.Handler) {
 	req := handler.MakeMsg()
 	requestName := gogoproto.MessageName(req)
@@ -854,7 +861,7 @@ func (s *stfRouterWrapper) RegisterHandler(handler appmodulev2.Handler) {
 	s.error = errors.Join(s.error, err)
 
 	// also make the decoder
-	if s.error == nil {
+	if s.handlers == nil {
 		s.handlers = map[string]appmodulev2.Handler{}
 	}
 	s.handlers[requestName] = handler
diff --git a/server/v2/cometbft/abci.go b/server/v2/cometbft/abci.go
index 5daf0c103038..45a5be69da89 100644
--- a/server/v2/cometbft/abci.go
+++ b/server/v2/cometbft/abci.go
@@ -259,12 +259,15 @@ func (c *Consensus[T]) maybeRunGRPCQuery(ctx context.Context, req *abci.QueryReq
 		return nil, false, err
 	}
 
+	var handlerFullName string
 	md, isGRPC := desc.(protoreflect.MethodDescriptor)
 	if !isGRPC {
-		return nil, false, nil
+		handlerFullName = string(desc.FullName())
+	} else {
+		handlerFullName = string(md.Input().FullName())
 	}
 
-	handler, found := c.queryHandlersMap[string(md.Input().FullName())]
+	handler, found := c.queryHandlersMap[handlerFullName]
 	if !found {
 		return nil, true, fmt.Errorf("no query handler found for %s", req.Path)
 	}
@@ -282,7 +285,7 @@ func (c *Consensus[T]) maybeRunGRPCQuery(ctx context.Context, req *abci.QueryReq
 	}
 
 	resp, err = queryResponse(res, req.Height)
-	return resp, isGRPC, err
+	return resp, true, err
 }
 
 // InitChain implements types.Application.