Optimise query stake-snapshot command

Author: newhoggy

cabal.project

@@ -263,8 +263,8 @@ source-repository-package
 source-repository-package
   type: git
   location: https://github.com/input-output-hk/ouroboros-network
-  tag: c764553561bed8978d2c6753d1608dc65449617a
-  --sha256: 0hdh7xdrvxw943r6qr0xr4kwszindh5mnsn1lww6qdnxnmn7wcsc
+  tag: a06a2472827df073493f64307ef0d854679aaa77
+  --sha256: 0fchrb4cba2j2jxcc8dk9sn3hnrhc1x56pjy37q7rq9hv1lcx5y2
   subdir:
     monoidal-synchronisation
     network-mux

cardano-api/src/Cardano/Api/KeysShelley.hs

@@ -1184,7 +1184,9 @@ instance SerialiseAsBech32 (SigningKey StakePoolKey) where
     bech32PrefixesPermitted _ = ["pool_sk"]
 
 newtype instance Hash StakePoolKey =
-    StakePoolKeyHash (Shelley.KeyHash Shelley.StakePool StandardCrypto)
+    StakePoolKeyHash
+    { unStakePoolKeyHash :: Shelley.KeyHash Shelley.StakePool StandardCrypto
+    }
   deriving stock (Eq, Ord)
   deriving (Show, IsString) via UsingRawBytesHex (Hash StakePoolKey)
   deriving (ToCBOR, FromCBOR) via UsingRawBytes (Hash StakePoolKey)

cardano-api/src/Cardano/Api/Orphans.hs

@@ -6,6 +6,7 @@
 {-# LANGUAGE GADTs #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE NamedFieldPuns #-}
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE StandaloneDeriving #-}
@@ -18,7 +19,7 @@ module Cardano.Api.Orphans () where
 
 import           Prelude
 
-import           Data.Aeson (FromJSON (..), ToJSON (..), object, (.=))
+import           Data.Aeson (FromJSON (..), ToJSON (..), object, pairs, (.=))
 import qualified Data.Aeson as Aeson
 import           Data.Aeson.Types (ToJSONKey (..), toJSONKeyText)
 import           Data.BiMap (BiMap (..), Bimap)
@@ -63,6 +64,7 @@ import           Cardano.Ledger.Shelley.PParams (PParamsUpdate)
 import qualified Cardano.Ledger.Shelley.Rewards as Shelley
 import qualified Cardano.Ledger.Shelley.RewardUpdate as Shelley
 import qualified Ouroboros.Consensus.Shelley.Eras as Consensus
+import qualified Ouroboros.Consensus.Shelley.Ledger.Query as Consensus
 
 import           Cardano.Api.Script
 
@@ -681,3 +683,37 @@ instance Crypto.Crypto crypto => ToJSON (VMap VB VB (Shelley.KeyHash    'Shelley
 instance Crypto.Crypto crypto => ToJSON (VMap VB VP (Shelley.Credential 'Shelley.Staking   crypto) (Shelley.CompactForm Shelley.Coin)) where
   toJSON = toJSON . fmap fromCompact . VMap.toMap
   toEncoding = toEncoding . fmap fromCompact . VMap.toMap
+
+-----
+
+instance ToJSON (Consensus.StakeSnapshots crypto) where
+  toJSON = object . stakeSnapshotsToPair
+  toEncoding = pairs . mconcat . stakeSnapshotsToPair
+
+stakeSnapshotsToPair :: Aeson.KeyValue a => Consensus.StakeSnapshots crypto -> [a]
+stakeSnapshotsToPair Consensus.StakeSnapshots
+    { Consensus.ssStakeSnapshots
+    , Consensus.ssMarkTotal
+    , Consensus.ssSetTotal
+    , Consensus.ssGoTotal
+    } = mconcat
+    -- Only output the first pool in order to preserve backwards compatibility of the output
+    -- format.  The output format will have to change to support multiple pools when that
+    -- functionality is added.
+    [ take 1 (Map.elems ssStakeSnapshots) >>= stakeSnapshotToPair
+    , [ "activeStakeMark" .= ssMarkTotal
+      , "activeStakeSet" .= ssSetTotal
+      , "activeStakeGo" .= ssGoTotal
+      ]
+    ]
+
+stakeSnapshotToPair :: Aeson.KeyValue a => Consensus.StakeSnapshot crypto -> [a]
+stakeSnapshotToPair Consensus.StakeSnapshot
+    { Consensus.ssMarkPool
+    , Consensus.ssSetPool
+    , Consensus.ssGoPool
+    } =
+    [ "poolStakeMark" .= ssMarkPool
+    , "poolStakeSet" .= ssSetPool
+    , "poolStakeGo" .= ssGoPool
+    ]

cardano-api/src/Cardano/Api/Query.hs

@@ -48,6 +48,10 @@ module Cardano.Api.Query (
     PoolState(..),
     decodePoolState,
 
+    SerialisedStakeSnapshots(..),
+    StakeSnapshot(..),
+    decodeStakeSnapshot,
+
     EraHistory(..),
     SystemStart(..),
 
@@ -245,6 +249,10 @@ data QueryInShelleyBasedEra era result where
     :: Maybe (Set PoolId)
     -> QueryInShelleyBasedEra era (SerialisedPoolState era)
 
+  QueryStakeSnapshot
+    :: PoolId
+    -> QueryInShelleyBasedEra era (SerialisedStakeSnapshots era)
+
 deriving instance Show (QueryInShelleyBasedEra era result)
 
 
@@ -403,6 +411,18 @@ decodePoolState
   -> Either DecoderError (PoolState era)
 decodePoolState (SerialisedPoolState (Serialised ls)) = PoolState <$> decodeFull ls
 
+newtype SerialisedStakeSnapshots era
+  = SerialisedStakeSnapshots (Serialised (Consensus.StakeSnapshots (Ledger.Crypto (ShelleyLedgerEra era))))
+
+newtype StakeSnapshot era = StakeSnapshot (Consensus.StakeSnapshots (Ledger.Crypto (ShelleyLedgerEra era)))
+
+decodeStakeSnapshot
+  :: forall era. ()
+  => FromCBOR (Consensus.StakeSnapshots (Ledger.Crypto (ShelleyLedgerEra era)))
+  => SerialisedStakeSnapshots era
+  -> Either DecoderError (StakeSnapshot era)
+decodeStakeSnapshot (SerialisedStakeSnapshots (Serialised ls)) = StakeSnapshot <$> decodeFull ls
+
 toShelleyAddrSet :: CardanoEra era
                  -> Set AddressAny
                  -> Set (Shelley.Addr Consensus.StandardCrypto)
@@ -571,7 +591,7 @@ toConsensusQueryShelleyBased erainmode (QueryStakePoolParameters poolids) =
     Some (consensusQueryInEraInMode erainmode (Consensus.GetStakePoolParams poolids'))
   where
     poolids' :: Set (Shelley.KeyHash Shelley.StakePool Consensus.StandardCrypto)
-    poolids' = Set.map (\(StakePoolKeyHash kh) -> kh) poolids
+    poolids' = Set.map unStakePoolKeyHash poolids
 
 toConsensusQueryShelleyBased erainmode QueryDebugLedgerState =
     Some (consensusQueryInEraInMode erainmode (Consensus.GetCBOR Consensus.DebugNewEpochState))
@@ -583,10 +603,10 @@ toConsensusQueryShelleyBased erainmode QueryCurrentEpochState =
     Some (consensusQueryInEraInMode erainmode (Consensus.GetCBOR Consensus.DebugEpochState))
 
 toConsensusQueryShelleyBased erainmode (QueryPoolState poolIds) =
-    Some (consensusQueryInEraInMode erainmode (Consensus.GetCBOR (Consensus.GetPoolState (getPoolIds <$> poolIds))))
-  where
-    getPoolIds :: Set PoolId -> Set (Shelley.KeyHash Shelley.StakePool Consensus.StandardCrypto)
-    getPoolIds = Set.map (\(StakePoolKeyHash kh) -> kh)
+    Some (consensusQueryInEraInMode erainmode (Consensus.GetCBOR (Consensus.GetPoolState (Set.map unStakePoolKeyHash <$> poolIds))))
+
+toConsensusQueryShelleyBased erainmode (QueryStakeSnapshot poolId) =
+    Some (consensusQueryInEraInMode erainmode (Consensus.GetCBOR (Consensus.GetStakeSnapshots (Just (Set.singleton (unStakePoolKeyHash poolId))))))
 
 consensusQueryInEraInMode
   :: forall era mode erablock modeblock result result' xs.
@@ -823,6 +843,11 @@ fromConsensusQueryResultShelleyBased _ QueryPoolState{} q' r' =
     Consensus.GetCBOR Consensus.GetPoolState {} -> SerialisedPoolState r'
     _                                           -> fromConsensusQueryResultMismatch
 
+fromConsensusQueryResultShelleyBased _ QueryStakeSnapshot{} q' r' =
+  case q' of
+    Consensus.GetCBOR Consensus.GetStakeSnapshots {} -> SerialisedStakeSnapshots r'
+    _                                                -> fromConsensusQueryResultMismatch
+
 -- | This should /only/ happen if we messed up the mapping in 'toConsensusQuery'
 -- and 'fromConsensusQueryResult' so they are inconsistent with each other.
 --

cardano-api/src/Cardano/Api/Shelley.hs

@@ -207,9 +207,15 @@ module Cardano.Api.Shelley
     CurrentEpochState(..),
     SerialisedCurrentEpochState(..),
     decodeCurrentEpochState,
+
     PoolState(..),
     SerialisedPoolState(..),
     decodePoolState,
+
+    StakeSnapshot(..),
+    SerialisedStakeSnapshots(..),
+    decodeStakeSnapshot,
+
     UTxO(..),
     AcquireFailure(..),
     SystemStart(..),
@@ -235,8 +241,8 @@ import           Cardano.Api.Address
 import           Cardano.Api.Block
 import           Cardano.Api.Certificate
 import           Cardano.Api.Eras
-import           Cardano.Api.IPC
 import           Cardano.Api.InMode
+import           Cardano.Api.IPC
 import           Cardano.Api.KeysByron
 import           Cardano.Api.KeysPraos
 import           Cardano.Api.KeysShelley

cardano-cli/src/Cardano/CLI/Shelley/Run/Query.hs

@@ -53,7 +53,6 @@ import qualified Data.Text.IO as Text
 import           Data.Text.Lazy.Builder (toLazyText)
 import           Data.Time.Clock
 import qualified Data.Vector as Vector
-import qualified Data.VMap as VMap
 import           Formatting.Buildable (build)
 import           Numeric (showEFloat)
 import qualified System.IO as IO
@@ -75,18 +74,14 @@ import qualified Cardano.Crypto.Hash.Blake2b as Blake2b
 import qualified Cardano.Crypto.VRF as Crypto
 import qualified Cardano.Ledger.Alonzo.PParams as Alonzo
 import           Cardano.Ledger.BaseTypes (Seed, UnitInterval)
-import           Cardano.Ledger.Coin
-import           Cardano.Ledger.Compactible
 import qualified Cardano.Ledger.Core as Core
 import qualified Cardano.Ledger.Credential as Ledger
 import qualified Cardano.Ledger.Crypto as Crypto
 import qualified Cardano.Ledger.Era as Era
 import qualified Cardano.Ledger.Era as Ledger
 import           Cardano.Ledger.Keys (KeyHash (..), KeyRole (..))
 import           Cardano.Ledger.Shelley.Constraints
-import           Cardano.Ledger.Shelley.EpochBoundary
-import           Cardano.Ledger.Shelley.LedgerState (EpochState (esSnapshots),
-                   NewEpochState (nesEs), PState (_fPParams, _pParams, _retiring))
+import           Cardano.Ledger.Shelley.LedgerState (PState (_fPParams, _pParams, _retiring))
 import qualified Cardano.Ledger.Shelley.LedgerState as SL
 import qualified Cardano.Ledger.Shelley.PParams as Shelley
 import           Cardano.Ledger.Shelley.Scripts ()
@@ -134,6 +129,7 @@ data ShelleyQueryCmdError
       FilePath
       -- ^ Operational certificate of the unknown stake pool.
   | ShelleyQueryCmdPoolStateDecodeError DecoderError
+  | ShelleyQueryCmdStakeSnapshotDecodeError DecoderError
 
   deriving Show
 
@@ -171,6 +167,8 @@ renderShelleyQueryCmdError err =
                   "in the current epoch, you must wait until the following epoch for the registration to take place."
     ShelleyQueryCmdPoolStateDecodeError decoderError ->
       "Failed to decode PoolState.  Error: " <> Text.pack (show decoderError)
+    ShelleyQueryCmdStakeSnapshotDecodeError decoderError ->
+      "Failed to decode StakeSnapshot.  Error: " <> Text.pack (show decoderError)
 
 runQueryCmd :: QueryCmd -> ExceptT ShelleyQueryCmdError IO ()
 runQueryCmd cmd =
@@ -629,7 +627,7 @@ runQueryStakeSnapshot
   -> NetworkId
   -> Hash StakePoolKey
   -> ExceptT ShelleyQueryCmdError IO ()
-runQueryStakeSnapshot (AnyConsensusModeParams cModeParams) network poolid = do
+runQueryStakeSnapshot (AnyConsensusModeParams cModeParams) network poolId = do
   SocketPath sockPath <- firstExceptT ShelleyQueryCmdEnvVarSocketErr readEnvSocketPath
   let localNodeConnInfo = LocalNodeConnectInfo cModeParams network sockPath
 
@@ -640,9 +638,9 @@ runQueryStakeSnapshot (AnyConsensusModeParams cModeParams) network poolid = do
   eInMode <- toEraInMode era cMode
     & hoistMaybe (ShelleyQueryCmdEraConsensusModeMismatch (AnyConsensusMode cMode) anyE)
 
-  let qInMode = QueryInEra eInMode . QueryInShelleyBasedEra sbe $ QueryDebugLedgerState
+  let qInMode = QueryInEra eInMode . QueryInShelleyBasedEra sbe $ QueryStakeSnapshot poolId
   result <- executeQuery era cModeParams localNodeConnInfo qInMode
-  obtainLedgerEraClassConstraints sbe (writeStakeSnapshot poolid) result
+  obtainLedgerEraClassConstraints sbe writeStakeSnapshot result
 
 
 runQueryLedgerState
@@ -794,44 +792,15 @@ writeLedgerState mOutFile qState@(SerialisedDebugLedgerState serLedgerState) =
 writeStakeSnapshot :: forall era ledgerera. ()
   => ShelleyLedgerEra era ~ ledgerera
   => Era.Crypto ledgerera ~ StandardCrypto
-  => FromCBOR (DebugLedgerState era)
-  => PoolId
-  -> SerialisedDebugLedgerState era
+  => SerialisedStakeSnapshots era
   -> ExceptT ShelleyQueryCmdError IO ()
-writeStakeSnapshot (StakePoolKeyHash hk) qState =
-  case decodeDebugLedgerState qState of
-    -- In the event of decode failure print the CBOR instead
-    Left bs -> firstExceptT ShelleyQueryCmdHelpersError $ pPrintCBOR bs
-
-    Right ledgerState -> do
-      -- Ledger State
-      let (DebugLedgerState snapshot) = ledgerState
-
-      -- The three stake snapshots, obtained from the ledger state
-      let (SnapShots markS setS goS _) = esSnapshots $ nesEs snapshot
+writeStakeSnapshot qState =
+  case decodeStakeSnapshot qState of
+    Left err -> left (ShelleyQueryCmdStakeSnapshotDecodeError err)
 
+    Right (StakeSnapshot snapshot) -> do
       -- Calculate the three pool and active stake values for the given pool
-      liftIO . LBS.putStrLn $ encodePretty $ Stakes
-        { markPool = getPoolStake hk markS
-        , setPool = getPoolStake hk setS
-        , goPool = getPoolStake hk goS
-        , markTotal = getAllStake markS
-        , setTotal = getAllStake setS
-        , goTotal = getAllStake goS
-        }
-
--- | Sum all the stake that is held by the pool
-getPoolStake :: KeyHash Cardano.Ledger.Keys.StakePool crypto -> SnapShot crypto -> Integer
-getPoolStake hash ss = pStake
-  where
-    Coin pStake = fold (Map.map fromCompact $ VMap.toMap s)
-    Stake s = poolStake hash (_delegations ss) (_stake ss)
-
--- | Sum the active stake from a snapshot
-getAllStake :: SnapShot crypto -> Integer
-getAllStake (SnapShot stake _ _) = activeStake
-  where
-    Coin activeStake = fold (fmap fromCompact (VMap.toMap (unStake stake)))
+      liftIO . LBS.putStrLn $ encodePretty snapshot
 
 -- | This function obtains the pool parameters, equivalent to the following jq query on the output of query ledger-state
 --   .nesEs.esLState._delegationState._pstate._pParams.<pool_id>