feat: claim-batch

Author: VladasZ

contract/src/claim/api.rs

@@ -1,9 +1,8 @@
+use std::collections::HashMap;
+
 use near_sdk::{env, ext_contract, json_types::U128, near_bindgen, AccountId, PromiseOrValue};
 use sweat_jar_model::{
-    api::ClaimApi,
-    claimed_amount_view::ClaimedAmountView,
-    jar::{AggregatedTokenAmountView, JarIdView},
-    U32,
+    api::ClaimApi, claimed_amount_view::ClaimedAmountView, jar::AggregatedTokenAmountView, ProductId, TokenAmount,
 };
 
 use crate::{
@@ -12,7 +11,7 @@ use crate::{
     internal::is_promise_success,
     jar::model::Jar,
     score::AccountScore,
-    Contract, ContractExt, JarsStorage,
+    Contract, ContractExt, JarsStorage, Product,
 };
 
 #[allow(dead_code)] // False positive since rust 1.78. It is used from `ext_contract` macro.
@@ -33,58 +32,67 @@ impl ClaimApi for Contract {
     fn claim_total(&mut self, detailed: Option<bool>) -> PromiseOrValue<ClaimedAmountView> {
         let account_id = env::predecessor_account_id();
         self.migrate_account_if_needed(&account_id);
-        let jar_ids = self.account_jars(&account_id).iter().map(|a| U32(a.id)).collect();
-        self.claim_jars_internal(account_id, jar_ids, detailed)
+        self.claim_jars_internal(account_id, detailed)
     }
 }
 
 impl Contract {
     fn claim_jars_internal(
         &mut self,
         account_id: AccountId,
-        jar_ids: Vec<JarIdView>,
         detailed: Option<bool>,
     ) -> PromiseOrValue<ClaimedAmountView> {
         let now = env::block_timestamp_ms();
         let mut accumulator = ClaimedAmountView::new(detailed);
 
-        let unlocked_jars: Vec<Jar> = self
-            .account_jars(&account_id)
-            .iter()
-            .filter(|jar| !jar.is_pending_withdraw && jar_ids.contains(&U32(jar.id)))
-            .cloned()
-            .collect();
+        let account_jars = self.account_jars(&account_id);
 
-        let mut event_data: Vec<ClaimEventItem> = vec![];
+        // UnorderedMap doesn't have cache and deserializes `Product` on each get
+        // This cache significantly reduces gas usage
+        let mut products_cache: HashMap<ProductId, Product> = HashMap::new();
 
         let account_score = self.get_score_mut(&account_id);
 
         let account_score_before_transfer = account_score.as_ref().map(|s| **s);
 
         let score = account_score.map(AccountScore::claim_score).unwrap_or_default();
 
-        for jar in &unlocked_jars {
-            let product = self.get_product(&jar.product_id);
-            let (interest, remainder) = jar.get_interest(&score, &product, now);
+        let mut unlocked_jars: Vec<((TokenAmount, u64), &Jar)> = account_jars
+            .iter()
+            .filter(|jar| !jar.is_pending_withdraw)
+            .map(|jar| {
+                let product = products_cache
+                    .entry(jar.product_id.clone())
+                    .or_insert_with(|| self.get_product(&jar.product_id));
+                (jar.get_interest(&score, product, now), jar)
+            })
+            .collect();
+
+        unlocked_jars.sort_by(|a, b| b.0 .0.cmp(&a.0 .0));
+
+        let jars_to_claim: Vec<_> = unlocked_jars.into_iter().take(100).collect();
+
+        let mut event_data: Vec<ClaimEventItem> = vec![];
 
-            if interest > 0 {
+        for ((available_interest, remainder), jar) in &jars_to_claim {
+            if *available_interest > 0 {
                 let jar = self.get_jar_mut_internal(&jar.account_id, jar.id);
 
-                jar.claim_remainder = remainder;
+                jar.claim_remainder = *remainder;
 
-                jar.claim(interest, now).lock();
+                jar.claim(*available_interest, now).lock();
 
-                accumulator.add(jar.id, interest);
+                accumulator.add(jar.id, *available_interest);
 
-                event_data.push((jar.id, U128(interest)));
+                event_data.push((jar.id, U128(*available_interest)));
             }
         }
 
         if accumulator.get_total().0 > 0 {
             self.claim_interest(
                 &account_id,
                 accumulator,
-                unlocked_jars,
+                jars_to_claim.into_iter().map(|a| a.1).cloned().collect(),
                 account_score_before_transfer,
                 EventKind::Claim(event_data),
                 now,

contract/src/integration_test/integration_test.rs

@@ -1,7 +1,7 @@
 #![cfg(feature = "integration-test")]
 
 use near_sdk::{env, near_bindgen, AccountId, Timestamp};
-use sweat_jar_model::{api::IntegrationTestMethods, jar::JarView, ProductId};
+use sweat_jar_model::{api::IntegrationTestMethods, ProductId};
 
 use crate::{jar::model::Jar, Contract, ContractExt};
 
@@ -12,17 +12,11 @@ impl IntegrationTestMethods for Contract {
         env::block_timestamp_ms()
     }
 
-    fn bulk_create_jars(
-        &mut self,
-        account_id: AccountId,
-        product_id: ProductId,
-        principal: u128,
-        number_of_jars: u16,
-    ) -> Vec<JarView> {
+    fn bulk_create_jars(&mut self, account_id: AccountId, product_id: ProductId, principal: u128, number_of_jars: u16) {
         self.assert_manager();
+        let now = env::block_timestamp_ms();
         (0..number_of_jars)
-            .map(|_| self.create_jar_for_integration_tests(&account_id, &product_id, principal))
-            .collect()
+            .for_each(|_| self.create_jar_for_integration_tests(&account_id, &product_id, principal, now));
     }
 }
 
@@ -33,18 +27,11 @@ impl Contract {
         account_id: &AccountId,
         product_id: &ProductId,
         amount: u128,
-    ) -> JarView {
-        let product = self.get_product(&product_id);
-
-        product.assert_enabled();
-        product.assert_cap(amount);
-
+        now: u64,
+    ) {
         let id = self.increment_and_get_last_jar_id();
-        let now = env::block_timestamp_ms();
         let jar = Jar::create(id, account_id.clone(), product_id.clone(), amount, now);
 
-        self.add_new_jar(account_id, jar.clone());
-
-        jar.into()
+        self.add_new_jar(account_id, jar);
     }
 }

contract/src/jar/api.rs

@@ -4,12 +4,13 @@ use near_sdk::{env, env::panic_str, json_types::U128, near_bindgen, require, Acc
 use sweat_jar_model::{
     api::JarApi,
     jar::{AggregatedInterestView, AggregatedTokenAmountView, JarId, JarIdView, JarView},
-    TokenAmount, U32,
+    ProductId, TokenAmount, U32,
 };
 
 use crate::{
     event::{emit, EventKind, RestakeData},
     jar::model::Jar,
+    product::model::Product,
     score::AccountScore,
     Contract, ContractExt, JarsStorage,
 };
@@ -145,13 +146,21 @@ impl JarApi for Contract {
         let mut detailed_amounts = HashMap::<JarIdView, U128>::new();
         let mut total_amount: TokenAmount = 0;
 
+        // UnorderedMap doesn't have cache and deserializes `Product` on each get
+        // This cache significantly reduces gas usage
+        let mut products_cache: HashMap<ProductId, Product> = HashMap::new();
+
         let score = self
             .get_score(&account_id)
             .map(AccountScore::claimable_score)
             .unwrap_or_default();
 
         for jar in self.account_jars_with_ids(&account_id, &jar_ids) {
-            let interest = jar.get_interest(&score, &self.get_product(&jar.product_id), now).0;
+            let product = products_cache
+                .entry(jar.product_id.clone())
+                .or_insert_with(|| self.get_product(&jar.product_id));
+
+            let interest = jar.get_interest(&score, product, now).0;
 
             detailed_amounts.insert(U32(jar.id), U128(interest));
             total_amount += interest;

integration-tests/src/context.rs

@@ -138,7 +138,7 @@ pub trait ContextHelpers {
         product_id: &ProductId,
         principal: u128,
         number_of_jars: u16,
-    ) -> Result<Vec<JarView>>;
+    ) -> Result<()>;
     async fn account_balance(&self, account: &Account) -> Result<u128>;
 }
 
@@ -159,7 +159,7 @@ impl ContextHelpers for Context {
         product_id: &ProductId,
         principal: u128,
         number_of_jars: u16,
-    ) -> Result<Vec<JarView>> {
+    ) -> Result<()> {
         let total_amount = principal * number_of_jars as u128;
 
         self.ft_contract()
@@ -189,7 +189,9 @@ impl ContextHelpers for Context {
         self.sweat_jar()
             .bulk_create_jars(account.to_near(), product_id.clone(), principal, number_of_jars)
             .with_user(&manager)
-            .await
+            .await?;
+
+        Ok(())
     }
 
     async fn account_balance(&self, account: &Account) -> Result<u128> {

integration-tests/src/many_jars.rs

@@ -1,5 +1,5 @@
 use anyhow::Result;
-use nitka::misc::ToNear;
+use nitka::{misc::ToNear, set_integration_logs_enabled};
 use sweat_jar_model::api::{ClaimApiIntegration, IntegrationTestMethodsIntegration, JarApiIntegration};
 
 use crate::{
@@ -12,32 +12,66 @@ use crate::{
 async fn claim_many_jars() -> Result<()> {
     println!("👷🏽 Claim many jars test");
 
+    set_integration_logs_enabled(false);
+
     let mut context = prepare_contract(None, [Locked5Minutes60000Percents]).await?;
 
     let alice = context.alice().await?;
     let manager = context.manager().await?;
 
     context
         .sweat_jar()
-        .bulk_create_jars(alice.to_near(), Locked5Minutes60000Percents.id(), 10000, 450)
+        .bulk_create_jars(alice.to_near(), Locked5Minutes60000Percents.id(), 1000, 4000)
         .with_user(&manager)
         .await?;
 
+    dbg!(context.sweat_jar().get_jars_for_account(alice.to_near()).await?.len());
+
     context.fast_forward_minutes(5).await?;
 
-    context
-        .sweat_jar()
-        .claim_total(true.into())
-        .with_user(&alice)
-        .result()
-        .await?;
+    let claimed = context.sweat_jar().claim_total(true.into()).with_user(&alice).await?;
 
-    assert!(context
-        .sweat_jar()
-        .get_jars_for_account(alice.to_near())
-        .await?
-        .iter()
-        .all(|j| j.is_pending_withdraw == false));
+    let batch_claim_summ = claimed.get_total().0;
+
+    dbg!(&batch_claim_summ);
+
+    assert_eq!(
+        batch_claim_summ * 39,
+        context
+            .sweat_jar()
+            .get_total_interest(alice.to_near())
+            .await?
+            .amount
+            .total
+            .0
+    );
+
+    for i in 1..40 {
+        let claimed = context.sweat_jar().claim_total(true.into()).with_user(&alice).await?;
+        assert_eq!(claimed.get_total().0, batch_claim_summ);
+
+        assert_eq!(
+            batch_claim_summ * (39 - i),
+            context
+                .sweat_jar()
+                .get_total_interest(alice.to_near())
+                .await?
+                .amount
+                .total
+                .0
+        );
+    }
+
+    assert_eq!(
+        context
+            .sweat_jar()
+            .get_total_interest(alice.to_near())
+            .await?
+            .amount
+            .total
+            .0,
+        0
+    );
 
     Ok(())
 }

integration-tests/src/withdraw_all.rs

@@ -59,6 +59,9 @@ async fn withdraw_all() -> Result<()> {
 
     context.fast_forward_minutes(6).await?;
 
+    // 3 calls to claim all 210 jars
+    context.sweat_jar().claim_total(None).with_user(&alice).await?;
+    context.sweat_jar().claim_total(None).with_user(&alice).await?;
     context.sweat_jar().claim_total(None).with_user(&alice).await?;
 
     let alice_balance = context.ft_contract().ft_balance_of(alice.to_near()).await?;

model/src/api.rs

@@ -295,11 +295,5 @@ pub trait ScoreApi {
 #[make_integration_version]
 pub trait IntegrationTestMethods {
     fn block_timestamp_ms(&self) -> near_sdk::Timestamp;
-    fn bulk_create_jars(
-        &mut self,
-        account_id: AccountId,
-        product_id: ProductId,
-        principal: u128,
-        number_of_jars: u16,
-    ) -> Vec<JarView>;
+    fn bulk_create_jars(&mut self, account_id: AccountId, product_id: ProductId, principal: u128, number_of_jars: u16);
 }