staking,delegation: make undelegate to work like unstake

Author: abarmat

contracts/staking/Staking.sol

@@ -503,6 +503,24 @@ contract Staking is StakingV1Storage, GraphUpgradeable, IStaking {
         return indexerStake.tokensAvailableWithDelegation(tokensDelegated);
     }
 
+    /**
+     * @dev Returns amount of delegated tokens ready to be withdrawn after unbonding period.
+     * @param _delegation Delegation of tokens from delegator to indexer
+     * @return Amount of tokens to withdraw
+     */
+    function getWithdraweableDelegatedTokens(Delegation memory _delegation)
+        public
+        view
+        returns (uint256)
+    {
+        // There must be locked tokens and period passed
+        uint256 currentEpoch = epochManager().currentEpoch();
+        if (_delegation.tokensLockedUntil > 0 && currentEpoch >= _delegation.tokensLockedUntil) {
+            return _delegation.tokensLocked;
+        }
+        return 0;
+    }
+
     /**
      * @dev Authorize an address to be an operator.
      * @param _operator Address to authorize
@@ -665,38 +683,14 @@ contract Staking is StakingV1Storage, GraphUpgradeable, IStaking {
     /**
      * @dev Withdraw delegated tokens once the unbonding period has passed.
      * @param _indexer Withdraw available tokens delegated to indexer
-     * @param _newIndexer Re-delegate to indexer address if non-zero, withdraw if zero address
+     * @param _delegateToIndexer Re-delegate to indexer address if non-zero, withdraw if zero address
      */
-    function withdrawDelegated(address _indexer, address _newIndexer) external override notPaused {
-        address delegator = msg.sender;
-
-        // Get the delegation pool of the indexer
-        DelegationPool storage pool = delegationPools[_indexer];
-        Delegation storage delegation = pool.delegators[delegator];
-
-        // There must be locked tokens and period passed
-        uint256 currentEpoch = epochManager().currentEpoch();
-        require(
-            delegation.tokensLockedUntil > 0 && currentEpoch >= delegation.tokensLockedUntil,
-            "No tokens available to withdraw"
-        );
-
-        // Get tokens available for withdrawal
-        uint256 tokensToWithdraw = delegation.tokensLocked;
-
-        // Reset lock
-        delegation.tokensLocked = 0;
-        delegation.tokensLockedUntil = 0;
-
-        emit StakeDelegatedWithdrawn(_indexer, delegator, tokensToWithdraw);
-
-        if (_newIndexer != address(0)) {
-            // Re-delegate tokens to a new indexer
-            _delegate(delegator, _newIndexer, tokensToWithdraw);
-        } else {
-            // Return tokens to the delegator
-            require(graphToken().transfer(delegator, tokensToWithdraw), "!transfer");
-        }
+    function withdrawDelegated(address _indexer, address _delegateToIndexer)
+        external
+        override
+        notPaused
+    {
+        _withdrawDelegated(msg.sender, _indexer, _delegateToIndexer);
     }
 
     /**
@@ -1162,6 +1156,11 @@ contract Staking is StakingV1Storage, GraphUpgradeable, IStaking {
         // Delegator need to have enough shares in the pool to undelegate
         require(delegation.shares >= _shares, "Delegator does not have enough shares");
 
+        // Withdraw tokens if available
+        if (getWithdraweableDelegatedTokens(delegation) > 0) {
+            _withdrawDelegated(_delegator, _indexer, address(0));
+        }
+
         // Calculate tokens to get in exchange for the shares
         uint256 tokens = _shares.mul(pool.tokens).div(pool.shares);
 
@@ -1185,6 +1184,44 @@ contract Staking is StakingV1Storage, GraphUpgradeable, IStaking {
         return tokens;
     }
 
+    /**
+     * @dev Withdraw delegated tokens once the unbonding period has passed.
+     * @param _delegator Delegator that is withdrawing tokens
+     * @param _indexer Withdraw available tokens delegated to indexer
+     * @param _delegateToIndexer Re-delegate to indexer address if non-zero, withdraw if zero address
+     */
+    function _withdrawDelegated(
+        address _delegator,
+        address _indexer,
+        address _delegateToIndexer
+    ) internal returns (uint256) {
+        // Get the delegation pool of the indexer
+        DelegationPool storage pool = delegationPools[_indexer];
+        Delegation storage delegation = pool.delegators[_delegator];
+
+        // Validation
+        uint256 tokensToWithdraw = getWithdraweableDelegatedTokens(delegation);
+        require(tokensToWithdraw > 0, "No tokens available to withdraw");
+
+        // Reset lock
+        delegation.tokensLocked = 0;
+        delegation.tokensLockedUntil = 0;
+
+        emit StakeDelegatedWithdrawn(_indexer, _delegator, tokensToWithdraw);
+
+        // -- Effects --
+
+        if (_delegateToIndexer != address(0)) {
+            // Re-delegate tokens to a new indexer
+            _delegate(_delegator, _delegateToIndexer, tokensToWithdraw);
+        } else {
+            // Return tokens to the delegator
+            require(graphToken().transfer(_delegator, tokensToWithdraw), "!transfer");
+        }
+
+        return tokensToWithdraw;
+    }
+
     /**
      * @dev Collect the delegation rewards for query fees.
      * This function will assign the collected fees to the delegation pool.

graph.config.yml

@@ -73,8 +73,10 @@ contracts:
         maxAllocationEpochs: 24 # Based on epoch length this is ~1 day (in epochs)
       - fn: "setThawingPeriod"
         thawingPeriod: 275 # ~1 hour (in blocks)
+      - fn: "setDelegationUnbondingPeriod"
+        delegationUnbondingPeriod: 2 # epochs
       - fn: "setDelegationRatio"
-        delegationRatio: 1 # 100% (multiplier)
+        delegationRatio: 100 # 100x (multiplier)
       - fn: "setProtocolPercentage"
         protocolPercentage: 10000 # 1% (in basis points)
       - fn: "setRebateRatio"

test/lib/deployment.ts

@@ -40,6 +40,7 @@ export const defaults = {
     channelDisputeEpochs: 1,
     maxAllocationEpochs: 5,
     thawingPeriod: 20, // in blocks
+    delegationUnbondingPeriod: 1, // in epochs
     alphaNumerator: 85,
     alphaDenominator: 100,
   },
@@ -211,6 +212,9 @@ export async function deployStaking(deployer: Signer, controller: string): Promi
   await staking.connect(deployer).setChannelDisputeEpochs(defaults.staking.channelDisputeEpochs)
   await staking.connect(deployer).setMaxAllocationEpochs(defaults.staking.maxAllocationEpochs)
   await staking.connect(deployer).setThawingPeriod(defaults.staking.thawingPeriod)
+  await staking
+    .connect(deployer)
+    .setDelegationUnbondingPeriod(defaults.staking.delegationUnbondingPeriod)
   await staking
     .connect(deployer)
     .setRebateRatio(defaults.staking.alphaNumerator, defaults.staking.alphaDenominator)

test/staking/delegation.test.ts

@@ -381,6 +381,11 @@ describe('Staking::Delegation', () => {
         await advanceToNextEpoch(epochManager)
         await shouldUndelegate(delegator, toGRT('25'))
       })
+
+      it('should undelegate and withdraw freed tokens from unbonding period', async function () {
+        await shouldDelegate(delegator, toGRT('100'))
+        await shouldUndelegate(delegator2, toGRT('50'))
+      })
     })
 
     describe('withdraw', function () {