fix(StargateQueries): use a sync pool when unmarshalling responses of protobuf objects (backport #7346) (#7349)

* fix(StargateQueries): use a sync pool when unmarshalling responses of protobuf objects  (#7346)

* use a sync pool when unmarshalling responses of protobuf objects in StargateQueries

* fix uninitted pool

* type assertion and lints

* changelog

* add comment for returnStargateResponseToPool

* add setWhitelistedQuery comment

* lint

---------

Co-authored-by: unknown unknown <unknown@unknown>
Co-authored-by: Adam Tucker <[email protected]>
(cherry picked from commit 2caa5c6)

# Conflicts:
#	CHANGELOG.md

* changelog

---------

Co-authored-by: testinginprod <[email protected]>
Co-authored-by: Adam Tucker <[email protected]>

Author: 3 people

CHANGELOG.md

@@ -40,6 +40,12 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## v22.0.1
+
+### Bug Fixes
+
+* [#7346](https://github.com/osmosis-labs/osmosis/pull/7346) Prevent heavy gRPC load from app hashing nodes
+
 ## v22.0.0
 
 ### Fee Market Parameter Updates

cmd/osmosisd/cmd/stargate-query.go

@@ -78,7 +78,7 @@ Example:
 //nolint:staticcheck
 func GetStructAndFill(queryPath, module, structName string, structArguments ...string) (interface{}, error) {
 	const ParamRequest = "QueryParamsRequest"
-	_, err := wasmbinding.GetWhitelistedQuery(queryPath)
+	err := wasmbinding.IsWhitelistedQuery(queryPath)
 	if err != nil {
 		return nil, err
 	}

wasmbinding/export_test.go

@@ -2,6 +2,10 @@ package wasmbinding
 
 import "github.com/cosmos/cosmos-sdk/codec"
 
-func SetWhitelistedQuery(queryPath string, protoType codec.ProtoMarshaler) {
+func SetWhitelistedQuery[T any, PT protoTypeG[T]](queryPath string, protoType PT) {
 	setWhitelistedQuery(queryPath, protoType)
 }
+
+func GetWhitelistedQuery(queryPath string) (codec.ProtoMarshaler, error) {
+	return getWhitelistedQuery(queryPath)
+}

wasmbinding/query_plugin.go

@@ -17,11 +17,15 @@ import (
 // StargateQuerier dispatches whitelisted stargate queries
 func StargateQuerier(queryRouter baseapp.GRPCQueryRouter, cdc codec.Codec) func(ctx sdk.Context, request *wasmvmtypes.StargateQuery) ([]byte, error) {
 	return func(ctx sdk.Context, request *wasmvmtypes.StargateQuery) ([]byte, error) {
-		protoResponseType, err := GetWhitelistedQuery(request.Path)
+		protoResponseType, err := getWhitelistedQuery(request.Path)
 		if err != nil {
 			return nil, err
 		}
 
+		// no matter what happens after this point, we must return
+		// the response type to prevent sync.Pool from leaking.
+		defer returnStargateResponseToPool(request.Path, protoResponseType)
+
 		route := queryRouter.Route(request.Path)
 		if route == nil {
 			return nil, wasmvmtypes.UnsupportedRequest{Kind: fmt.Sprintf("No route to query '%s'", request.Path)}

wasmbinding/stargate_whitelist.go

@@ -33,14 +33,15 @@ import (
 	epochtypes "github.com/osmosis-labs/osmosis/x/epochs/types"
 )
 
-// stargateWhitelist keeps whitelist and its deterministic
+// stargateResponsePools keeps whitelist and its deterministic
 // response binding for stargate queries.
 // CONTRACT: since results of queries go into blocks, queries being added here should always be
 // deterministic or can cause non-determinism in the state machine.
 //
-// The query can be multi-thread, so we have to use
-// thread safe sync.Map.
-var stargateWhitelist sync.Map
+// The query is multi-threaded so we're using a sync.Pool
+// to manage the allocation and de-allocation of newly created
+// pb objects.
+var stargateResponsePools = make(map[string]*sync.Pool)
 
 // Note: When adding a migration here, we should also add it to the Async ICQ params in the upgrade.
 // In the future we may want to find a better way to keep these in sync
@@ -184,34 +185,57 @@ func init() {
 	setWhitelistedQuery("/osmosis.concentratedliquidity.v1beta1.Query/CFMMPoolIdLinkFromConcentratedPoolId", &concentratedliquidityquery.CFMMPoolIdLinkFromConcentratedPoolIdResponse{})
 }
 
-// GetWhitelistedQuery returns the whitelisted query at the provided path.
+// IsWhitelistedQuery returns if the query is not whitelisted.
+func IsWhitelistedQuery(queryPath string) error {
+	_, isWhitelisted := stargateResponsePools[queryPath]
+	if !isWhitelisted {
+		return wasmvmtypes.UnsupportedRequest{Kind: fmt.Sprintf("'%s' path is not allowed from the contract", queryPath)}
+	}
+	return nil
+}
+
+// getWhitelistedQuery returns the whitelisted query at the provided path.
 // If the query does not exist, or it was setup wrong by the chain, this returns an error.
-func GetWhitelistedQuery(queryPath string) (codec.ProtoMarshaler, error) {
-	protoResponseAny, isWhitelisted := stargateWhitelist.Load(queryPath)
+// CONTRACT: must call returnStargateResponseToPool in order to avoid pointless allocs.
+func getWhitelistedQuery(queryPath string) (codec.ProtoMarshaler, error) {
+	protoResponseAny, isWhitelisted := stargateResponsePools[queryPath]
 	if !isWhitelisted {
 		return nil, wasmvmtypes.UnsupportedRequest{Kind: fmt.Sprintf("'%s' path is not allowed from the contract", queryPath)}
 	}
-	protoResponseType, ok := protoResponseAny.(codec.ProtoMarshaler)
+	protoMarshaler, ok := protoResponseAny.Get().(codec.ProtoMarshaler)
 	if !ok {
-		return nil, wasmvmtypes.Unknown{}
+		return nil, fmt.Errorf("failed to assert type to codec.ProtoMarshaler")
 	}
-	return protoResponseType, nil
+	return protoMarshaler, nil
 }
 
-func setWhitelistedQuery(queryPath string, protoType codec.ProtoMarshaler) {
-	stargateWhitelist.Store(queryPath, protoType)
+type protoTypeG[T any] interface {
+	*T
+	codec.ProtoMarshaler
+}
+
+// setWhitelistedQuery sets the whitelisted query at the provided path.
+// This method also creates a sync.Pool for the provided protoMarshaler.
+// We use generics so we can properly instantiate an object that the
+// queryPath expects as a response.
+func setWhitelistedQuery[T any, PT protoTypeG[T]](queryPath string, _ PT) {
+	stargateResponsePools[queryPath] = &sync.Pool{
+		New: func() any {
+			return PT(new(T))
+		},
+	}
+}
+
+// returnStargateResponseToPool returns the provided protoMarshaler to the appropriate pool based on it's query path.
+func returnStargateResponseToPool(queryPath string, pb codec.ProtoMarshaler) {
+	stargateResponsePools[queryPath].Put(pb)
 }
 
 func GetStargateWhitelistedPaths() (keys []string) {
 	// Iterate over the map and collect the keys
-	stargateWhitelist.Range(func(key, value interface{}) bool {
-		keyStr, ok := key.(string)
-		if !ok {
-			panic("key is not a string")
-		}
-		keys = append(keys, keyStr)
-		return true
-	})
-
+	keys = make([]string, 0, len(stargateResponsePools))
+	for k := range stargateResponsePools {
+		keys = append(keys, k)
+	}
 	return keys
 }