Merge pull request #113 from milkyway-labs/fabo/use-index

fabo/use index and reduce query load

contracts/staking/src/contract.rs

@@ -5,8 +5,8 @@ use crate::execute::{
 use crate::helpers::validate_addresses;
 use crate::ibc::{receive_ack, receive_timeout};
 use crate::query::{
-    query_batch, query_batches, query_claimable, query_config, query_ibc_queue,
-    query_pending_batch, query_reply_queue, query_state,
+    query_batch, query_batches, query_config, query_ibc_queue, query_pending_batch,
+    query_reply_queue, query_state, query_unstake_requests,
 };
 use crate::state::{
     new_unstake_request, Config, MultisigAddressConfig, ProtocolFeeConfig, State, ADMIN, BATCHES,
@@ -30,6 +30,7 @@ use cw2::set_contract_version;
 use cw_utils::must_pay;
 use milky_way::staking::Batch;
 use osmosis_std::types::osmosis::tokenfactory::v1beta1::MsgCreateDenom;
+use schemars::Map;
 use semver::Version;
 
 // Version information for migration
@@ -265,11 +266,16 @@ pub fn query(deps: Deps, _env: Env, msg: QueryMsg) -> StdResult<Binary> {
             status,
         } => to_binary(&query_batches(deps, start_after, limit, status)?),
         QueryMsg::PendingBatch {} => to_binary(&query_pending_batch(deps)?),
-        QueryMsg::ClaimableBatches {
+        QueryMsg::UnstakeRequests {
             user,
             start_after,
             limit,
-        } => to_binary(&query_claimable(deps, user, start_after, limit)?),
+        } => to_binary(&query_unstake_requests(
+            deps,
+            user.to_string(),
+            start_after,
+            limit,
+        )?),
 
         // dev only, depr
         QueryMsg::IbcQueue { start_after, limit } => {
@@ -318,6 +324,7 @@ pub fn migrate(mut deps: DepsMut, _env: Env, _msg: MigrateMsg) -> Result<Respons
         .into());
     }
     let mut batch_ids = Vec::<u64>::new();
+    let mut request_count = Map::<u64, u64>::new();
     let requests = BATCHES
         .range(deps.storage, None, None, cosmwasm_std::Order::Ascending)
         .map(|v| {
@@ -326,6 +333,7 @@ pub fn migrate(mut deps: DepsMut, _env: Env, _msg: MigrateMsg) -> Result<Respons
             let _requests = _v.liquid_unstake_requests;
             if _requests.is_some() {
                 let requests = _requests.unwrap().into_iter();
+                request_count.insert(k, requests.len() as u64);
                 return requests
                     .filter(|r| !r.1.redeemed)
                     .map(|r| return (r.0, k.clone(), r.1.shares))
@@ -342,6 +350,7 @@ pub fn migrate(mut deps: DepsMut, _env: Env, _msg: MigrateMsg) -> Result<Respons
         BATCHES.update(deps.storage, batch_id, |b| -> StdResult<Batch> {
             let mut batch = b.unwrap();
             batch.liquid_unstake_requests = None;
+            batch.unstake_requests_count = request_count.get(&batch_id).cloned();
             Ok(batch)
         })?;
     }

contracts/staking/src/execute.rs

@@ -10,9 +10,7 @@ use crate::state::{
     Config, IbcWaitingForReply, MultisigAddressConfig, ProtocolFeeConfig, State, ADMIN, BATCHES,
     CONFIG, IBC_WAITING_FOR_REPLY, INFLIGHT_PACKETS, PENDING_BATCH_ID, STATE,
 };
-use crate::state::{
-    new_unstake_request, remove_unstake_request, UNSTAKE_REQUESTS, UNSTAKE_REQUEST_COUNTERS,
-};
+use crate::state::{new_unstake_request, remove_unstake_request, unstake_requests, UnstakeRequest};
 use cosmwasm_std::{
     ensure, Addr, Binary, CosmosMsg, Deps, DepsMut, Env, IbcTimeout, MessageInfo, Order, ReplyOn,
     Response, SubMsg, SubMsgResponse, SubMsgResult, Timestamp, Uint128,
@@ -243,16 +241,26 @@ pub fn execute_liquid_unstake(
 
     // Add unstake request to pending batch
     let pending_unstake_request =
-        UNSTAKE_REQUESTS.load(deps.storage, (pending_batch_id, info.sender.to_string()));
+        unstake_requests().may_load(deps.storage, (pending_batch_id, info.sender.to_string()))?;
+    let is_new_request = pending_unstake_request.is_none();
     match pending_unstake_request {
-        Ok(current_amount) => {
-            UNSTAKE_REQUESTS.save(
+        Some(_) => {
+            unstake_requests().update(
                 deps.storage,
                 (pending_batch_id, info.sender.to_string()),
-                &(current_amount + amount),
+                |or| -> Result<UnstakeRequest, ContractError> {
+                    match or {
+                        Some(r) => Ok(UnstakeRequest {
+                            batch_id: r.batch_id,
+                            user: r.user.clone(),
+                            amount: r.amount + amount,
+                        }),
+                        None => Err(ContractError::NoRequestInBatch {}),
+                    }
+                },
             )?;
         }
-        Err(_) => {
+        None => {
             new_unstake_request(&mut deps, info.sender.to_string(), pending_batch_id, amount)?;
         }
     }
@@ -264,28 +272,13 @@ pub fn execute_liquid_unstake(
         |_batch| -> Result<Batch, ContractError> {
             let mut batch = _batch.unwrap();
             batch.batch_total_liquid_stake += amount;
+            if is_new_request {
+                batch.unstake_requests_count = Some(batch.unstake_requests_count.unwrap_or(0) + 1);
+            }
             Ok(batch)
         },
     )?;
 
-    // let mut msgs: Vec<CosmosMsg> = vec![];
-    // if batch period has elapsed, submit batch
-    // for simplicity not doing this for now
-    // if let Some(est_next_batch_action) = pending_batch.next_batch_action_time {
-    //     if est_next_batch_action >= env.block.time.seconds() {
-    //         msgs.push(CosmosMsg::Wasm(WasmMsg::Execute {
-    //             contract_addr: env.contract.address.to_string(),
-    //             msg: to_binary(&ExecuteMsg::SubmitBatch {
-    //                 batch_id: pending_batch_id,
-    //             })?,
-    //             funds: vec![],
-    //         }))
-    //     }
-
-    //     // Save updated pending batch
-    //     PENDING_BATCH.save(deps.storage, &pending_batch)?;
-    // }
-
     Ok(Response::new()
         .add_attribute("action", "liquid_unstake")
         .add_attribute("sender", info.sender.to_string())
@@ -323,7 +316,12 @@ pub fn execute_submit_batch(
             expected: 0u64,
         });
     }
-    let unstake_requests = UNSTAKE_REQUEST_COUNTERS.load(deps.storage, batch.id)?;
+
+    let unstake_requests = unstake_requests()
+        .prefix(pending_batch_id)
+        .range(deps.storage, None, None, Order::Ascending)
+        .take(1)
+        .count();
 
     if unstake_requests == 0 {
         return Err(ContractError::BatchEmpty {});
@@ -428,16 +426,15 @@ pub fn execute_withdraw(
     let received_native_unstaked = batch.received_native_unstaked.as_ref().unwrap();
 
     let _liquid_unstake_request =
-        UNSTAKE_REQUESTS.load(deps.storage, (batch.id, info.sender.to_string()));
-
-    if _liquid_unstake_request.is_err() {
+        unstake_requests().may_load(deps.storage, (batch.id, info.sender.to_string()))?;
+    if _liquid_unstake_request.is_none() {
         return Err(ContractError::NoRequestInBatch {});
     }
 
-    let liquid_unstake_request = _liquid_unstake_request.unwrap();
+    let unstake_request_amount = _liquid_unstake_request.unwrap().amount;
 
     let amount = received_native_unstaked
-        .multiply_ratio(liquid_unstake_request, batch.batch_total_liquid_stake);
+        .multiply_ratio(unstake_request_amount, batch.batch_total_liquid_stake);
 
     // TODO: if all liquid unstake requests have been withdrawn, delete the batch?
     remove_unstake_request(&mut deps, info.sender.to_string(), batch.id)?;

contracts/staking/src/msg.rs

@@ -126,9 +126,9 @@ pub struct BatchResponse {
     pub batch_total_liquid_stake: Uint128,
     pub expected_native_unstaked: Uint128,
     pub received_native_unstaked: Uint128,
+    pub unstake_request_count: u64,
     pub next_batch_action_time: Timestamp,
     pub status: String,
-    pub requests: Vec<LiquidUnstakeRequestResponse>,
 }
 #[derive(Serialize, Deserialize, Clone, PartialEq, JsonSchema, Debug, Default)]
 pub struct BatchesResponse {
@@ -161,7 +161,7 @@ pub enum QueryMsg {
     #[returns(BatchResponse)]
     PendingBatch {},
     #[returns(BatchesResponse)]
-    ClaimableBatches {
+    UnstakeRequests {
         user: Addr,
         start_after: Option<u64>,
         limit: Option<u32>,

contracts/staking/src/query.rs

@@ -1,14 +1,14 @@
 use crate::helpers::{get_redemption_rate, paginate_map};
 use crate::msg::{
     BatchResponse, BatchesResponse, ConfigResponse, IBCQueueResponse, IBCReplyQueueResponse,
-    LiquidUnstakeRequestResponse, StateResponse,
+    StateResponse,
 };
 use crate::state::ibc::IBCTransfer;
 use crate::state::{
-    BATCHES, CONFIG, IBC_WAITING_FOR_REPLY, INFLIGHT_PACKETS, PENDING_BATCH_ID, STATE,
-    UNSTAKE_REQUESTS, UNSTAKE_REQUESTS_BY_USER,
+    unstake_requests, BATCHES, CONFIG, IBC_WAITING_FOR_REPLY, INFLIGHT_PACKETS, PENDING_BATCH_ID,
+    STATE,
 };
-use cosmwasm_std::{Addr, Deps, StdResult, Timestamp, Uint128};
+use cosmwasm_std::{Deps, StdResult, Timestamp, Uint128};
 use cw_storage_plus::Bound;
 use milky_way::staking::{Batch, BatchStatus};
 
@@ -66,12 +66,7 @@ pub fn query_state(deps: Deps) -> StdResult<StateResponse> {
     Ok(res)
 }
 
-fn batch_to_response(deps: Deps, batch: Batch) -> BatchResponse {
-    let unstake_requests = UNSTAKE_REQUESTS
-        .prefix(batch.id)
-        .range(deps.storage, None, None, cosmwasm_std::Order::Ascending)
-        .map(|v| v.unwrap())
-        .collect::<Vec<_>>();
+fn batch_to_response(batch: Batch) -> BatchResponse {
     BatchResponse {
         id: batch.id,
         batch_total_liquid_stake: batch.batch_total_liquid_stake,
@@ -81,19 +76,13 @@ fn batch_to_response(deps: Deps, batch: Batch) -> BatchResponse {
             batch.next_batch_action_time.unwrap_or(0u64),
         ),
         status: batch.status.as_str().to_string(),
-        requests: unstake_requests
-            .into_iter()
-            .map(|v| LiquidUnstakeRequestResponse {
-                user: v.0,
-                amount: v.1,
-            })
-            .collect(),
+        unstake_request_count: batch.unstake_requests_count.unwrap_or(0), // Fallback. Only is none if migration failed. Would be set in updates for new batches though
     }
 }
 
 pub fn query_batch(deps: Deps, id: u64) -> StdResult<BatchResponse> {
     let batch: Batch = BATCHES.load(deps.storage, id)?;
-    Ok(batch_to_response(deps, batch))
+    Ok(batch_to_response(batch))
 }
 
 pub fn query_batches(
@@ -115,10 +104,7 @@ pub fn query_batches(
     )?;
 
     let res = BatchesResponse {
-        batches: batches
-            .into_iter()
-            .map(|v| batch_to_response(deps, v))
-            .collect(),
+        batches: batches.into_iter().map(|v| batch_to_response(v)).collect(),
     };
     Ok(res)
 }
@@ -127,7 +113,7 @@ pub fn query_pending_batch(deps: Deps) -> StdResult<BatchResponse> {
     let pending_batch_id = PENDING_BATCH_ID.load(deps.storage)?;
     let pending_batch = BATCHES.load(deps.storage, pending_batch_id)?;
 
-    Ok(batch_to_response(deps, pending_batch))
+    Ok(batch_to_response(pending_batch))
 }
 
 pub fn query_ibc_queue(
@@ -169,15 +155,15 @@ pub fn query_reply_queue(
     Ok(res)
 }
 
-pub fn query_claimable(
+pub fn query_unstake_requests(
     deps: Deps,
-    user: Addr,
+    user: String,
     start_after: Option<u64>,
     limit: Option<u32>,
 ) -> StdResult<BatchesResponse> {
-    deps.api.addr_validate(&user.to_string())?;
-
-    let unstaking_requests = UNSTAKE_REQUESTS_BY_USER
+    let unstaking_requests = unstake_requests()
+        .idx
+        .by_user
         .prefix(user.to_string())
         .range(
             deps.storage,
@@ -192,15 +178,11 @@ pub fn query_claimable(
     let batches = unstaking_requests
         .into_iter()
         .filter_map(|v| {
-            let batch_id = v.0;
+            let batch_id = v.1.batch_id;
             let batch = BATCHES.load(deps.storage, batch_id).ok()?;
-            if batch.status == BatchStatus::Received {
-                Some(batch)
-            } else {
-                None
-            }
+            Some(batch)
         })
-        .map(|v| batch_to_response(deps, v))
+        .map(|v| batch_to_response(v))
         .collect();
 
     let res = BatchesResponse { batches };

contracts/staking/src/state.rs

@@ -1,7 +1,7 @@
 use cosmwasm_schema::cw_serde;
 use cosmwasm_std::{Addr, DepsMut, StdError, Timestamp, Uint128};
 use cw_controllers::Admin;
-use cw_storage_plus::{Item, Map};
+use cw_storage_plus::{Index, IndexList, IndexedMap, Item, Map, UniqueIndex};
 use milky_way::staking::Batch;
 
 #[cw_serde]
@@ -51,24 +51,51 @@ pub const CONFIG: Item<Config> = Item::new("config");
 pub const ADMIN: Admin = Admin::new("admin");
 pub const STATE: Item<State> = Item::new("state");
 pub const BATCHES: Map<u64, Batch> = Map::new("batches");
-// <batch id, user address, request>
-pub const UNSTAKE_REQUESTS: Map<(u64, String), Uint128> = Map::new("unstake_requests");
-pub const UNSTAKE_REQUESTS_BY_USER: Map<(String, u64), bool> = Map::new("unstake_requests_by_user");
-pub const UNSTAKE_REQUEST_COUNTERS: Map<u64, u64> = Map::new("unstake_request_counters");
 pub const PENDING_BATCH_ID: Item<u64> = Item::new("pending_batch_id");
 
+#[cw_serde]
+pub struct UnstakeRequest {
+    pub batch_id: u64,
+    pub user: String,
+    pub amount: Uint128,
+}
+
+pub struct UnstakeRequestIndexes<'a> {
+    pub by_user: UniqueIndex<'a, (String, u64), UnstakeRequest>,
+}
+
+impl<'a> IndexList<UnstakeRequest> for UnstakeRequestIndexes<'a> {
+    fn get_indexes(&'_ self) -> Box<dyn Iterator<Item = &'_ dyn Index<UnstakeRequest>> + '_> {
+        let v: Vec<&dyn Index<UnstakeRequest>> = vec![&self.by_user];
+        Box::new(v.into_iter())
+    }
+}
+
+pub fn unstake_requests<'a>(
+) -> IndexedMap<'a, (u64, String), UnstakeRequest, UnstakeRequestIndexes<'a>> {
+    let indexes = UnstakeRequestIndexes {
+        by_user: UniqueIndex::new(|r| (r.user.clone(), r.batch_id), "unstake_requests_by_user"),
+    };
+
+    // depr version
+    IndexedMap::new("unstake_requests", indexes)
+}
+
 pub fn new_unstake_request(
     deps: &mut DepsMut,
     user: String,
     batch_id: u64,
     amount: Uint128,
 ) -> Result<(), StdError> {
-    UNSTAKE_REQUESTS.save(deps.storage, (batch_id, user.clone()), &amount)?;
-    UNSTAKE_REQUEST_COUNTERS.update(deps.storage, batch_id, |c| match c {
-        Some(c) => Ok::<u64, StdError>(c + 1),
-        None => Ok(1),
-    })?;
-    UNSTAKE_REQUESTS_BY_USER.save(deps.storage, (user, batch_id), &true)?;
+    unstake_requests().save(
+        deps.storage,
+        (batch_id, user.clone()),
+        &UnstakeRequest {
+            batch_id,
+            user,
+            amount,
+        },
+    )?;
     Ok(())
 }
 
@@ -77,12 +104,9 @@ pub fn remove_unstake_request(
     user: String,
     batch_id: u64,
 ) -> Result<(), StdError> {
-    UNSTAKE_REQUESTS.remove(deps.storage, (batch_id, user.clone()));
-    UNSTAKE_REQUEST_COUNTERS.update(deps.storage, batch_id, |c| match c {
-        Some(c) => Ok::<u64, StdError>(c - 1),
-        None => Ok(0),
-    })?;
-    UNSTAKE_REQUESTS_BY_USER.remove(deps.storage, (user, batch_id));
+    unstake_requests()
+        .remove(deps.storage, (batch_id, user.clone()))
+        .unwrap();
     Ok(())
 }