Skip to content

Commit

Permalink
Remove locking around distribution
Browse files Browse the repository at this point in the history
  • Loading branch information
ethanfrey committed Jun 23, 2023
1 parent 33c18d3 commit 426c594
Showing 1 changed file with 17 additions and 54 deletions.
71 changes: 17 additions & 54 deletions contracts/provider/external-staking/src/contract.rs
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ pub struct ExternalStakingContract<'a> {
/// Stakes indexed by `(owner, validator)` pair
pub stakes: Map<'a, (&'a Addr, &'a str), Lockable<Stake>>,
/// Per-validator distribution information
pub distribution: Map<'a, &'a str, Lockable<Distribution>>,
pub distribution: Map<'a, &'a str, Distribution>,
/// Pending txs information
pub tx_count: Item<'a, u64>,
pub pending_txs: Map<'a, u64, Tx>,
Expand Down Expand Up @@ -130,16 +130,17 @@ impl ExternalStakingContract<'_> {
.load(ctx.deps.storage, (&tx_user, &tx_validator))?;

// Load distribution
let mut distribution_lock = self.distribution.load(ctx.deps.storage, &tx_validator)?;
let mut distribution = self
.distribution
.may_load(ctx.deps.storage, &tx_validator)?
.unwrap_or_default();

// Commit amount (need to unlock it first)
stake_lock.unlock_write()?;
let stake = stake_lock.write()?;
stake.stake += tx_amount;

// Commit distribution (need to unlock it first)
distribution_lock.unlock_write()?;
let distribution = distribution_lock.write()?;
// Distribution alignment
stake
.points_alignment
Expand All @@ -152,7 +153,7 @@ impl ExternalStakingContract<'_> {

// Save distribution
self.distribution
.save(ctx.deps.storage, &tx_validator, &distribution_lock)?;
.save(ctx.deps.storage, &tx_validator, &distribution)?;

// Remove tx
self.pending_txs.remove(ctx.deps.storage, tx_id);
Expand Down Expand Up @@ -199,23 +200,13 @@ impl ExternalStakingContract<'_> {
.stakes
.load(ctx.deps.storage, (&tx_user, &tx_validator))?;

// Load distribution
let mut distribution_lock = self.distribution.load(ctx.deps.storage, &tx_validator)?;

// Release stake lock
stake_lock.unlock_write()?;

// Save stake
self.stakes
.save(ctx.deps.storage, (&tx_user, &tx_validator), &stake_lock)?;

// Release distribution lock
distribution_lock.unlock_write()?;

// Save distribution
self.distribution
.save(ctx.deps.storage, &tx_validator, &distribution_lock)?;

// Remove tx
self.pending_txs.remove(ctx.deps.storage, tx_id);

Expand Down Expand Up @@ -265,12 +256,6 @@ impl ExternalStakingContract<'_> {
&stake_lock,
)?;

let mut distribution_lock = self.distribution.load(ctx.deps.storage, &validator)?;

distribution_lock.lock_write()?;
self.distribution
.save(ctx.deps.storage, &validator, &distribution_lock)?;

// Create new tx
let tx_id = self.next_tx_id(ctx.deps.storage)?;

Expand Down Expand Up @@ -329,7 +314,10 @@ impl ExternalStakingContract<'_> {
.load(ctx.deps.storage, (&tx_user, &tx_validator))?;

// Load distribution
let mut distribution_lock = self.distribution.load(ctx.deps.storage, &tx_validator)?;
let mut distribution = self
.distribution
.may_load(ctx.deps.storage, &tx_validator)?
.unwrap_or_default();

// Commit amount (need to unlock it first)
stake_lock.unlock_write()?;
Expand All @@ -344,9 +332,6 @@ impl ExternalStakingContract<'_> {
};
stake.pending_unbonds.push(unbond);

// Commit distribution (need to unlock it first)
distribution_lock.unlock_write()?;
let distribution = distribution_lock.write()?;
// Distribution alignment
stake
.points_alignment
Expand All @@ -359,7 +344,7 @@ impl ExternalStakingContract<'_> {

// Save distribution
self.distribution
.save(ctx.deps.storage, &tx_validator, &distribution_lock)?;
.save(ctx.deps.storage, &tx_validator, &distribution)?;

// Remove tx
self.pending_txs.remove(ctx.deps.storage, tx_id);
Expand Down Expand Up @@ -403,23 +388,13 @@ impl ExternalStakingContract<'_> {
.stakes
.load(ctx.deps.storage, (&tx_user, &tx_validator))?;

// Load distribution
let mut distribution_lock = self.distribution.load(ctx.deps.storage, &tx_validator)?;

// Release stake lock
stake_lock.unlock_write()?;

// Save stake
self.stakes
.save(ctx.deps.storage, (&tx_user, &tx_validator), &stake_lock)?;

// Release distribution lock
distribution_lock.unlock_write()?;

// Save distribution
self.distribution
.save(ctx.deps.storage, &tx_validator, &distribution_lock)?;

// Remove tx
self.pending_txs.remove(ctx.deps.storage, tx_id);
Ok(Response::new())
Expand Down Expand Up @@ -495,11 +470,10 @@ impl ExternalStakingContract<'_> {
let config = self.config.load(ctx.deps.storage)?;
let amount = must_pay(&ctx.info, &config.rewards_denom)?;

let mut distribution_lock = self
let mut distribution = self
.distribution
.may_load(ctx.deps.storage, &validator)?
.unwrap_or_default();
let mut distribution = distribution_lock.write()?;

let total_stake = Uint256::from(distribution.total_stake);
let points_distributed =
Expand All @@ -510,7 +484,7 @@ impl ExternalStakingContract<'_> {
distribution.points_per_stake += points_per_stake;

self.distribution
.save(ctx.deps.storage, &validator, &distribution_lock)?;
.save(ctx.deps.storage, &validator, &distribution)?;

let resp = Response::new()
.add_attribute("action", "distribute_rewards")
Expand All @@ -535,13 +509,12 @@ impl ExternalStakingContract<'_> {

let stake = stake_lock.write()?;

let mut distribution_lock = self
let distribution = self
.distribution
.may_load(ctx.deps.storage, &validator)?
.unwrap_or_default();
let distribution = distribution_lock.write()?;

let amount = Self::calculate_reward(stake, distribution)?;
let amount = Self::calculate_reward(stake, &distribution)?;

let mut resp = Response::new()
.add_attribute("action", "withdraw_rewards")
Expand Down Expand Up @@ -716,13 +689,12 @@ impl ExternalStakingContract<'_> {
.unwrap_or_default();
let stake = stake_lock.read()?;

let distribution_lock = self
let distribution = self
.distribution
.may_load(ctx.deps.storage, &validator)?
.unwrap_or_default();
let distribution = distribution_lock.read()?;

let amount = Self::calculate_reward(stake, distribution)?;
let amount = Self::calculate_reward(stake, &distribution)?;
let config = self.config.load(ctx.deps.storage)?;

let resp = PendingRewards {
Expand Down Expand Up @@ -786,20 +758,11 @@ pub mod cross_staking {
.may_load(ctx.deps.storage, (&owner, &msg.validator))?
.unwrap_or_default();

let mut distribution_lock = self
.distribution
.may_load(ctx.deps.storage, &msg.validator)?
.unwrap_or_default();

// Write lock and save stake and distribution
stake_lock.lock_write()?;
self.stakes
.save(ctx.deps.storage, (&owner, &msg.validator), &stake_lock)?;

distribution_lock.lock_write()?;
self.distribution
.save(ctx.deps.storage, &msg.validator, &distribution_lock)?;

// TODO: Send proper IBC message to remote staking contract

// Save tx
Expand Down

0 comments on commit 426c594

Please sign in to comment.