initial implementation query stakes

Author: Tommytrg

data_structures/src/staking/stakes.rs

@@ -11,6 +11,45 @@ use crate::{chain::PublicKeyHash, get_environment, transaction::StakeTransaction
 
 use super::prelude::*;
 
+/// Message for querying stakes
+#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
+pub enum QueryStakesKey<Address: Default + Ord> {
+    /// Query stakes by validator address
+    Validator(Address),
+    /// Query stakes by withdrawer address
+    Withdrawer(Address),
+    /// Query stakes by validator and withdrawer addresses
+    Key(StakeKey<Address>),
+}
+
+impl<Address: Default + Ord> Default for QueryStakesKey<Address> {
+    fn default() -> Self {
+        QueryStakesKey::Validator(Address::default())
+    }
+}
+
+impl<Address: Default + Ord, T> From<(Option<T>, Option<T>)> for QueryStakesKey<Address>
+where
+    T: Into<Address>,
+{
+    fn from(val: (Option<T>, Option<T>)) -> Self {
+        let validator = val.0;
+        let withdrawer = val.1;
+
+        if validator.is_some() && withdrawer.is_some() {
+            QueryStakesKey::Key(StakeKey {
+                validator: validator.unwrap().into(),
+                withdrawer: withdrawer.unwrap().into(),
+            })
+        } else if validator.is_some() {
+            QueryStakesKey::Validator(validator.unwrap().into())
+        // } else if withdrawer.is_some() {
+        } else {
+            QueryStakesKey::Withdrawer(withdrawer.unwrap().into())
+        }
+    }
+}
+
 /// The main data structure that provides the "stakes tracker" functionality.
 ///
 /// This structure holds indexes of stake entries. Because the entries themselves are reference
@@ -24,6 +63,10 @@ where
 {
     /// A listing of all the stakers, indexed by their address.
     by_key: BTreeMap<StakeKey<Address>, SyncStake<Address, Coins, Epoch, Power>>,
+    /// A listing of all the stakers, indexed by validator.
+    by_validator: BTreeMap<Address, SyncStake<Address, Coins, Epoch, Power>>,
+    /// A listing of all the stakers, indexed by withdrawer.
+    by_withdrawer: BTreeMap<Address, SyncStake<Address, Coins, Epoch, Power>>,
     /// A listing of all the stakers, indexed by their coins and address.
     ///
     /// Because this uses a compound key to prevent duplicates, if we want to know which addresses
@@ -83,7 +126,15 @@ where
             addresses: key,
         };
         self.by_coins.remove(&key);
-        self.by_coins.insert(key, stake.clone());
+        self.by_coins.insert(key.clone(), stake.clone());
+
+        let validator_key = key.clone().addresses.validator;
+        self.by_validator.remove(&validator_key);
+        self.by_validator.insert(validator_key, stake.clone());
+
+        let withdrawer_key = key.addresses.withdrawer;
+        self.by_withdrawer.remove(&withdrawer_key);
+        self.by_withdrawer.insert(withdrawer_key, stake.clone());
 
         Ok(stake.value.read()?.clone())
     }
@@ -228,6 +279,39 @@ where
             ..Default::default()
         }
     }
+
+    /// Query stakes based on different keys.
+    pub fn query_stakes<IQSK>(&mut self, query: IQSK) -> Option<Coins>
+    where
+        IQSK: Into<QueryStakesKey<Address>>,
+    {
+        match query.into() {
+            QueryStakesKey::Key(key) => self.query_by_key(&key),
+            QueryStakesKey::Validator(validator) => self.query_by_validator(&validator),
+            QueryStakesKey::Withdrawer(withdrawer) => self.query_by_withdrawer(&withdrawer),
+        }
+    }
+
+    /// Query stakes by stake key.
+    fn query_by_key(&self, key: &StakeKey<Address>) -> Option<Coins> {
+        self.by_key
+            .get_key_value(key)
+            .and_then(|entry| entry.1.value.read().map(|stake| stake.coins).ok())
+    }
+
+    /// Query stakes by validator address.
+    fn query_by_validator(&self, validator: &Address) -> Option<Coins> {
+        self.by_validator
+            .get_key_value(validator)
+            .and_then(|entry| entry.1.value.read().map(|stake| stake.coins).ok())
+    }
+
+    /// Query stakes by withdrawer address.
+    fn query_by_withdrawer(&self, withdrawer: &Address) -> Option<Coins> {
+        self.by_withdrawer
+            .get_key_value(withdrawer)
+            .and_then(|entry| entry.1.value.read().map(|stake| stake.coins).ok())
+    }
 }
 
 /// Adds stake, based on the data from a stake transaction.
@@ -591,4 +675,27 @@ mod tests {
             ]
         );
     }
+
+    #[test]
+    fn test_query_stakes() {
+        // First, lets create a setup with a few stakers
+        let mut stakes = Stakes::<String, u64, u64, u64>::with_minimum(5);
+        let alice = "Alice";
+        let bob = "Bob";
+        let charlie = "Charlie";
+        let david = "David";
+        let erin = "Erin";
+
+        let alice_charlie = (alice, charlie);
+        let bob_david = (bob, david);
+        let charlie_erin = (charlie, erin);
+
+        stakes.add_stake(alice_charlie, 10, 0).unwrap();
+        stakes.add_stake(bob_david, 20, 20).unwrap();
+        stakes.add_stake(charlie_erin, 30, 30).unwrap();
+
+        let result = stakes.query_stakes(QueryStakesKey::Key(alice_charlie.into()));
+
+        assert_eq!(result, Some(10))
+    }
 }

node/src/actors/chain_manager/handlers.rs

@@ -19,6 +19,7 @@ use witnet_data_structures::{
         Hashable, NodeStats, PublicKeyHash, SuperBlockVote, SupplyInfo,
     },
     error::{ChainInfoError, TransactionError::DataRequestNotFound},
+    staking::{aux::StakeKey, stakes::QueryStakesKey},
     transaction::{DRTransaction, StakeTransaction, Transaction, VTTransaction},
     transaction_factory::{self, NodeBalance},
     types::LastBeacon,
@@ -37,9 +38,9 @@ use crate::{
             GetDataRequestInfo, GetHighestCheckpointBeacon, GetMemoryTransaction, GetMempool,
             GetMempoolResult, GetNodeStats, GetReputation, GetReputationResult, GetSignalingInfo,
             GetState, GetSuperBlockVotes, GetSupplyInfo, GetUtxoInfo, IsConfirmedBlock,
-            PeersBeacons, ReputationStats, Rewind, SendLastBeacon, SessionUnitResult,
-            SetLastBeacon, SetPeersLimits, SignalingInfo, SnapshotExport, SnapshotImport,
-            TryMineBlock,
+            PeersBeacons, QueryStake, QueryStakesParams, ReputationStats, Rewind, SendLastBeacon,
+            SessionUnitResult, SetLastBeacon, SetPeersLimits, SignalingInfo, SnapshotExport,
+            SnapshotImport, TryMineBlock,
         },
         sessions_manager::SessionsManager,
     },
@@ -1356,6 +1357,24 @@ impl Handler<BuildStake> for ChainManager {
     }
 }
 
+impl Handler<QueryStake> for ChainManager {
+    type Result = <QueryStake as Message>::Result;
+
+    fn handle(&mut self, msg: QueryStake, _ctx: &mut Self::Context) -> Self::Result {
+        let query_stake_param = match msg.key {
+            QueryStakesParams::Key(key) => QueryStakesKey::Key(StakeKey {
+                validator: key.0,
+                withdrawer: key.1,
+            }),
+            QueryStakesParams::Validator(v) => QueryStakesKey::Validator(v),
+            QueryStakesParams::Withdrawer(w) => QueryStakesKey::Withdrawer(w),
+        };
+
+        // build address from public key hash
+        Ok(self.chain_state.stakes.query_stakes(query_stake_param))
+    }
+}
+
 impl Handler<BuildDrt> for ChainManager {
     type Result = ResponseActFuture<Self, Result<DRTransaction, failure::Error>>;
 

node/src/actors/json_rpc/api.rs

@@ -45,8 +45,8 @@ use crate::{
             GetConsolidatedPeers, GetDataRequestInfo, GetEpoch, GetHighestCheckpointBeacon,
             GetItemBlock, GetItemSuperblock, GetItemTransaction, GetKnownPeers,
             GetMemoryTransaction, GetMempool, GetNodeStats, GetReputation, GetSignalingInfo,
-            GetState, GetSupplyInfo, GetUtxoInfo, InitializePeers, IsConfirmedBlock, Rewind,
-            SnapshotExport, SnapshotImport, StakeAuthorization,
+            GetState, GetSupplyInfo, GetUtxoInfo, InitializePeers, IsConfirmedBlock, QueryStake,
+            QueryStakesParams, Rewind, SnapshotExport, SnapshotImport, StakeAuthorization,
         },
         peers_manager::PeersManager,
         sessions_manager::SessionsManager,
@@ -139,6 +139,9 @@ pub fn attach_regular_methods<H>(
         Box::pin(signaling_info())
     });
     server.add_actix_method(system, "priority", |_params: Params| Box::pin(priority()));
+    server.add_actix_method(system, "queryStakes", |params: Params| {
+        Box::pin(query_stakes(params.parse()))
+    });
 }
 
 /// Attach the sensitive JSON-RPC methods to a multi-transport server.
@@ -2094,6 +2097,60 @@ pub async fn authorize_stake(params: Result<AuthorizeStake, Error>) -> JsonRpcRe
         .await
 }
 
+/// Param for query_stakes  
+#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
+pub enum QueryStakesArgument {
+    /// To query by stake validator
+    Validator(String),
+    /// To query by stake withdrawer
+    Withdrawer(String),
+    /// To query by stake validator and withdrawer
+    Key((String, String)),
+}
+
+/// Query the amount of nanowits staked by an address.
+pub async fn query_stakes(params: Result<Option<QueryStakesArgument>, Error>) -> JsonRpcResult {
+    // Short-circuit if parameters are wrong
+    let params = params?;
+
+    // If a withdrawer address is not specified, default to local node address
+    let key: QueryStakesParams = if let Some(address) = params {
+        match address {
+            QueryStakesArgument::Validator(k) => QueryStakesParams::Validator(
+                PublicKeyHash::from_bech32(get_environment(), &k).map_err(internal_error)?,
+            ),
+            QueryStakesArgument::Withdrawer(k) => QueryStakesParams::Withdrawer(
+                PublicKeyHash::from_bech32(get_environment(), &k).map_err(internal_error)?,
+            ),
+            QueryStakesArgument::Key((v, w)) => QueryStakesParams::Key((
+                PublicKeyHash::from_bech32(get_environment(), &v).map_err(internal_error)?,
+                PublicKeyHash::from_bech32(get_environment(), &w).map_err(internal_error)?,
+            )),
+        }
+    } else {
+        let pk = signature_mngr::public_key().await.map_err(internal_error)?;
+
+        QueryStakesParams::Validator(PublicKeyHash::from_public_key(&pk))
+    };
+
+    ChainManager::from_registry()
+        .send(QueryStake { key })
+        .map(|res| match res {
+            Ok(Ok(staked_amount)) => {
+                serde_json::to_value(staked_amount.unwrap()).map_err(internal_error)
+            }
+            Ok(Err(e)) => {
+                let err = internal_error_s(e);
+                Err(err)
+            }
+            Err(e) => {
+                let err = internal_error_s(e);
+                Err(err)
+            }
+        })
+        .await
+}
+
 #[cfg(test)]
 mod mock_actix {
     use actix::{MailboxError, Message};

node/src/actors/messages.rs

@@ -34,6 +34,7 @@ use witnet_data_structures::{
     transaction_factory::NodeBalance,
     types::LastBeacon,
     utxo_pool::{UtxoInfo, UtxoSelectionStrategy},
+    wit::Wit,
 };
 use witnet_p2p::{
     error::SessionsError,
@@ -305,6 +306,34 @@ impl Message for StakeAuthorization {
     type Result = Result<String, failure::Error>;
 }
 
+/// Stake key for quering stakes
+#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
+pub enum QueryStakesParams {
+    /// To search by the validator public key hash
+    Validator(PublicKeyHash),
+    /// To search by the withdrawer public key hash
+    Withdrawer(PublicKeyHash),
+    /// To search by validator and withdrawer public key hashes
+    Key((PublicKeyHash, PublicKeyHash)),
+}
+
+impl Default for QueryStakesParams {
+    fn default() -> Self {
+        QueryStakesParams::Validator(PublicKeyHash::default())
+    }
+}
+
+/// Message for querying stakes
+#[derive(Clone, Debug, Default, Eq, PartialEq, Serialize, Deserialize)]
+pub struct QueryStake {
+    /// stake key used to search the stake
+    pub key: QueryStakesParams,
+}
+
+impl Message for QueryStake {
+    type Result = Result<Option<Wit>, failure::Error>;
+}
+
 /// Builds a `DataRequestTransaction` from a `DataRequestOutput`
 #[derive(Clone, Debug, Default, Hash, Eq, PartialEq, Serialize, Deserialize)]
 pub struct BuildDrt {