Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 7 commits into from
Jan 20, 2024
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion cmd/osmosisd/cmd/stargate-query.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down
5 changes: 4 additions & 1 deletion wasmbinding/query_plugin.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,10 +17,13 @@ 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, but we must return
// the response type to the pool.
defer returnStargateResponseToPool(request.Path, protoResponseType)
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we need to return the sync.Pool object, so it does not leak


route := queryRouter.Route(request.Path)
if route == nil {
Expand Down
59 changes: 37 additions & 22 deletions wasmbinding/stargate_whitelist.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 map[string]*sync.Pool
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we use a sync.Pool of multiple proto responses type so we do not allocate every time, this should provide relief to the GC in moments of high traffic.


// 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
Expand Down Expand Up @@ -184,34 +185,48 @@ 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 {
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Exposed this method in place of getWhitelistedQuery to avoid unexported usage of a function that can leak memory if not used properly

_, 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)
if !ok {
return nil, wasmvmtypes.Unknown{}
return protoResponseAny.Get().(codec.ProtoMarshaler), nil
}

type protoTypeG[T any] interface {
*T
codec.ProtoMarshaler
}

func setWhitelistedQuery[T any, PT protoTypeG[T]](queryPath string, _ PT) {
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this creates a sync.Pool for the given protobuf object, we use generics so we can properly instantiate an object that queryPath expects as response.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you comment in the code with this context please?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added comment here 801eae8

stargateResponsePools[queryPath] = &sync.Pool{
New: func() any {
return PT(new(T))
},
}
return protoResponseType, nil
}

func setWhitelistedQuery(queryPath string, protoType codec.ProtoMarshaler) {
stargateWhitelist.Store(queryPath, protoType)
func returnStargateResponseToPool(queryPath string, pb codec.ProtoMarshaler) {
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this returns the protobuf object to its appropriate pool (based on the queryPath)

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Transferring this into a comment would also be helpful IMO 🙏

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added comment here 2192082

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
}