Total liquidity store (osmosis-labs#194)

* basic changes for total liquidity store * add query for total liquidity * add liquidity amount modifications, e2e tests, cli command & tests * fix query total liquidity cli command and test
CrueLu · May 25, 2021 · 1218e92 · 1218e92
1 parent 682b2a2
commit 1218e92
Show file tree

Hide file tree

Showing 11 changed files with 649 additions and 79 deletions.
diff --git a/proto/osmosis/gamm/v1beta1/query.proto b/proto/osmosis/gamm/v1beta1/query.proto
@@ -34,6 +34,9 @@ service Query {
   rpc SpotPrice(QuerySpotPriceRequest) returns (QuerySpotPriceResponse) {
     option (google.api.http).get = "/osmosis/gamm/v1beta1/{poolId}/prices";
   }
+  rpc TotalLiquidity(QueryTotalLiquidityRequest) returns (QueryTotalLiquidityResponse) {
+    option (google.api.http).get = "/osmosis/gamm/v1beta1/total_liquidity";
+  }
 
   // Estimate the swap.
   rpc EstimateSwapExactAmountIn(QuerySwapExactAmountInRequest)
@@ -153,3 +156,12 @@ message QuerySwapExactAmountOutResponse {
     (gogoproto.nullable) = false
   ];
 }
+
+message QueryTotalLiquidityRequest {}
+
+message QueryTotalLiquidityResponse {
+  repeated cosmos.base.v1beta1.Coin liquidity = 1 [
+    (gogoproto.moretags) = "yaml:\"liquidity\"",
+    (gogoproto.nullable) = false
+  ];
+}
diff --git a/x/gamm/client/cli/cli_test.go b/x/gamm/client/cli/cli_test.go
@@ -934,6 +934,42 @@ func (s *IntegrationTestSuite) TestGetCmdTotalShare() {
 	}
 }
 
+func (s *IntegrationTestSuite) TestGetCmdTotalLiquidity() {
+	val := s.network.Validators[0]
+
+	testCases := []struct {
+		name      string
+		args      []string
+		expectErr bool
+	}{
+		{
+			"query total liquidity", // osmosisd query gamm total-liquidity
+			[]string{
+				fmt.Sprintf("--%s=%s", tmcli.OutputFlag, "json"),
+			},
+			false,
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+
+		s.Run(tc.name, func() {
+			cmd := cli.GetCmdQueryTotalLiquidity()
+			clientCtx := val.ClientCtx
+
+			out, err := clitestutil.ExecTestCLICmd(clientCtx, cmd, tc.args)
+			if tc.expectErr {
+				s.Require().Error(err)
+			} else {
+				resp := types.QueryTotalLiquidityResponse{}
+				s.Require().NoError(err, out.String())
+				s.Require().NoError(clientCtx.JSONMarshaler.UnmarshalJSON(out.Bytes(), &resp), out.String())
+			}
+		})
+	}
+}
+
 func (s *IntegrationTestSuite) TestGetCmdSpotPrice() {
 	val := s.network.Validators[0]
 

diff --git a/x/gamm/client/cli/query.go b/x/gamm/client/cli/query.go
@@ -32,6 +32,7 @@ func GetQueryCmd() *cobra.Command {
 		GetCmdTotalShare(),
 		GetCmdPoolAssets(),
 		GetCmdSpotPrice(),
+		GetCmdQueryTotalLiquidity(),
 		GetCmdEstimateSwapExactAmountIn(),
 		GetCmdEstimateSwapExactAmountOut(),
 	)
@@ -253,6 +254,41 @@ $ %s query gamm total-share 1
 	return cmd
 }
 
+// GetCmdQueryTotalLiquidity return total liquidity
+func GetCmdQueryTotalLiquidity() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "total-liquidity",
+		Short: "Query total-liquidity",
+		Long: strings.TrimSpace(
+			fmt.Sprintf(`Query total-liquidity.
+Example:
+$ %s query gamm total-liquidity
+`,
+				version.AppName,
+			),
+		),
+		Args: cobra.ExactArgs(0),
+		RunE: func(cmd *cobra.Command, args []string) error {
+			clientCtx, err := client.GetClientQueryContext(cmd)
+			if err != nil {
+				return err
+			}
+			queryClient := types.NewQueryClient(clientCtx)
+
+			res, err := queryClient.TotalLiquidity(cmd.Context(), &types.QueryTotalLiquidityRequest{})
+			if err != nil {
+				return err
+			}
+
+			return clientCtx.PrintProto(res)
+		},
+	}
+
+	flags.AddQueryFlagsToCmd(cmd)
+
+	return cmd
+}
+
 // GetCmdPoolAssets return pool-assets for a pool
 func GetCmdPoolAssets() *cobra.Command {
 	cmd := &cobra.Command{

diff --git a/x/gamm/keeper/grpc_query.go b/x/gamm/keeper/grpc_query.go
@@ -204,6 +204,18 @@ func (k Keeper) SpotPrice(ctx context.Context, req *types.QuerySpotPriceRequest)
 	}, nil
 }
 
+func (k Keeper) TotalLiquidity(ctx context.Context, req *types.QueryTotalLiquidityRequest) (*types.QueryTotalLiquidityResponse, error) {
+	if req == nil {
+		return nil, status.Error(codes.InvalidArgument, "empty request")
+	}
+
+	sdkCtx := sdk.UnwrapSDKContext(ctx)
+
+	return &types.QueryTotalLiquidityResponse{
+		Liquidity: k.GetTotalLiquidity(sdkCtx),
+	}, nil
+}
+
 func (k Keeper) EstimateSwapExactAmountIn(ctx context.Context, req *types.QuerySwapExactAmountInRequest) (*types.QuerySwapExactAmountInResponse, error) {
 	if req == nil {
 		return nil, status.Error(codes.InvalidArgument, "empty request")

diff --git a/x/gamm/keeper/grpc_query_test.go b/x/gamm/keeper/grpc_query_test.go
@@ -158,6 +158,22 @@ func (suite *KeeperTestSuite) TestQueryTotalShare() {
 	suite.Require().Equal(types.InitPoolSharesSupply.Add(types.OneShare.MulRaw(10)).String(), res.TotalShare.Amount.String())
 }
 
+func (suite *KeeperTestSuite) TestQueryTotalLiquidity() {
+	queryClient := suite.queryClient
+
+	// Pool not exist
+	res, err := queryClient.TotalLiquidity(gocontext.Background(), &types.QueryTotalLiquidityRequest{})
+	suite.Require().NoError(err)
+	suite.Require().Equal("", sdk.Coins(res.Liquidity).String())
+
+	_ = suite.preparePool()
+
+	// create pool
+	res, err = queryClient.TotalLiquidity(gocontext.Background(), &types.QueryTotalLiquidityRequest{})
+	suite.Require().NoError(err)
+	suite.Require().Equal("5000000bar,5000000baz,5000000foo", sdk.Coins(res.Liquidity).String())
+}
+
 func (suite *KeeperTestSuite) TestQueryPoolAssets() {
 	queryClient := suite.queryClient
 

diff --git a/x/gamm/keeper/pool_service.go b/x/gamm/keeper/pool_service.go
@@ -83,6 +83,7 @@ func (k Keeper) CreatePool(
 	}
 
 	k.hooks.AfterPoolCreated(ctx, sender, pool.GetId())
+	k.RecordTotalLiquidityIncrease(ctx, coins)
 
 	return pool.GetId(), nil
 }
@@ -151,6 +152,7 @@ func (k Keeper) JoinPool(
 	}
 
 	k.hooks.AfterJoinPool(ctx, sender, pool.GetId(), coins, shareOutAmount)
+	k.RecordTotalLiquidityIncrease(ctx, coins)
 
 	return nil
 }
@@ -211,6 +213,7 @@ func (k Keeper) JoinSwapExternAmountIn(
 	}
 
 	k.hooks.AfterJoinPool(ctx, sender, pool.GetId(), sdk.Coins{tokenIn}, shareOutAmount)
+	k.RecordTotalLiquidityIncrease(ctx, sdk.Coins{tokenIn})
 
 	return shareOutAmount, nil
 }
@@ -272,6 +275,7 @@ func (k Keeper) JoinSwapShareAmountOut(
 	}
 
 	k.hooks.AfterJoinPool(ctx, sender, pool.GetId(), sdk.Coins{sdk.NewCoin(tokenInDenom, tokenInAmount)}, shareOutAmount)
+	k.RecordTotalLiquidityIncrease(ctx, sdk.Coins{sdk.NewCoin(tokenInDenom, tokenInAmount)})
 
 	return shareOutAmount, nil
 }
@@ -355,6 +359,7 @@ func (k Keeper) ExitPool(
 	}
 
 	k.hooks.AfterExitPool(ctx, sender, pool.GetId(), shareInAmount, coins)
+	k.RecordTotalLiquidityDecrease(ctx, coins)
 
 	return nil
 }
@@ -433,6 +438,7 @@ func (k Keeper) ExitSwapShareAmountIn(
 	}
 
 	k.hooks.AfterExitPool(ctx, sender, pool.GetId(), shareInAmount, sdk.Coins{sdk.NewCoin(tokenOutDenom, tokenOutAmount)})
+	k.RecordTotalLiquidityDecrease(ctx, sdk.Coins{sdk.NewCoin(tokenOutDenom, tokenOutAmount)})
 
 	return tokenOutAmount, nil
 }
@@ -510,6 +516,37 @@ func (k Keeper) ExitSwapExternAmountOut(
 	}
 
 	k.hooks.AfterExitPool(ctx, sender, pool.GetId(), shareInAmount, sdk.Coins{tokenOut})
+	k.RecordTotalLiquidityDecrease(ctx, sdk.Coins{tokenOut})
 
 	return shareInAmount, nil
 }
+
+func (k Keeper) GetTotalLiquidity(ctx sdk.Context) sdk.Coins {
+	store := ctx.KVStore(k.storeKey)
+	if !store.Has(types.KeyTotalLiquidity) {
+		return sdk.Coins{}
+	}
+
+	bz := store.Get(types.KeyTotalLiquidity)
+	coins, err := sdk.ParseCoinsNormalized(string(bz))
+	if err != nil {
+		panic("invalid total liquidity value set")
+	}
+
+	return coins
+}
+
+func (k Keeper) SetTotalLiquidity(ctx sdk.Context, coins sdk.Coins) {
+	store := ctx.KVStore(k.storeKey)
+	store.Set(types.KeyTotalLiquidity, []byte(coins.String()))
+}
+
+func (k Keeper) RecordTotalLiquidityIncrease(ctx sdk.Context, coins sdk.Coins) {
+	liquidity := k.GetTotalLiquidity(ctx)
+	k.SetTotalLiquidity(ctx, liquidity.Add(coins...))
+}
+
+func (k Keeper) RecordTotalLiquidityDecrease(ctx sdk.Context, coins sdk.Coins) {
+	liquidity := k.GetTotalLiquidity(ctx)
+	k.SetTotalLiquidity(ctx, liquidity.Sub(coins))
+}
diff --git a/x/gamm/keeper/pool_service_test.go b/x/gamm/keeper/pool_service_test.go
@@ -57,6 +57,9 @@ func (suite *KeeperTestSuite) TestCreatePool() {
 			suite.Require().Equal(types.InitPoolSharesSupply.String(), pool.GetTotalShare().Amount.String(),
 				fmt.Sprintf("share token should be minted as %s initially", types.InitPoolSharesSupply.String()),
 			)
+
+			liquidity := suite.app.GAMMKeeper.GetTotalLiquidity(suite.ctx)
+			suite.Require().Equal("10000bar,10000foo", liquidity.String())
 		},
 	}, {
 		fn: func() {
@@ -222,6 +225,9 @@ func (suite *KeeperTestSuite) TestJoinPool() {
 				// Thus, to get the 50*OneShare gamm/pool/1, (10000foo, 10000bar) * (1 / 2) balances should be provided.
 				suite.Require().Equal("5000", deltaBalances.AmountOf("foo").String())
 				suite.Require().Equal("5000", deltaBalances.AmountOf("bar").String())
+
+				liquidity := suite.app.GAMMKeeper.GetTotalLiquidity(suite.ctx)
+				suite.Require().Equal("15000bar,15000foo", liquidity.String())
 			},
 		},
 		{
@@ -258,6 +264,9 @@ func (suite *KeeperTestSuite) TestJoinPool() {
 					sdk.NewCoin("foo", sdk.NewInt(5000)),
 				})
 				suite.Require().NoError(err)
+
+				liquidity := suite.app.GAMMKeeper.GetTotalLiquidity(suite.ctx)
+				suite.Require().Equal("15000bar,15000foo", liquidity.String())
 			},
 		},
 	}

diff --git a/x/gamm/keeper/swap.go b/x/gamm/keeper/swap.go
@@ -144,6 +144,8 @@ func (k Keeper) updatePoolForSwap(
 	}
 
 	k.hooks.AfterSwap(ctx, sender, pool.GetId(), sdk.Coins{tokenIn}, sdk.Coins{tokenOut})
+	k.RecordTotalLiquidityIncrease(ctx, sdk.Coins{tokenIn})
+	k.RecordTotalLiquidityDecrease(ctx, sdk.Coins{tokenOut})
 
 	return err
 }

diff --git a/x/gamm/types/key.go b/x/gamm/types/key.go
@@ -21,6 +21,8 @@ var (
 	KeyNextGlobalPoolNumber = []byte{0x01}
 	// KeyPrefixPools defines prefix to store pools
 	KeyPrefixPools = []byte{0x02}
+	// KeyTotalLiquidity defines key to store total liquidity
+	KeyTotalLiquidity = []byte{0x03}
 )
 
 func GetPoolShareDenom(poolId uint64) string {