From fadf1efa417541549005557c82971b35fab1aedd Mon Sep 17 00:00:00 2001
From: "mergify[bot]" <37929162+mergify[bot]@users.noreply.github.com>
Date: Thu, 22 Apr 2021 21:02:46 +0000
Subject: [PATCH] Update getLeaderSchedule options (#16749) (#16752)

(cherry picked from commit 636b5987af7d0739fdb28c59b21316180de9d118)

Co-authored-by: Tyera Eulberg <teulberg@gmail.com>
---
 client/src/rpc_config.rs | 16 ++++++++++++++++
 core/src/rpc.rs          | 13 +++++++++----
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/client/src/rpc_config.rs b/client/src/rpc_config.rs
index eb67da9cc54539..def38d053106e1 100644
--- a/client/src/rpc_config.rs
+++ b/client/src/rpc_config.rs
@@ -49,6 +49,22 @@ pub struct RpcLeaderScheduleConfig {
     pub commitment: Option<CommitmentConfig>,
 }
 
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+#[serde(untagged)]
+pub enum RpcLeaderScheduleConfigWrapper {
+    SlotOnly(Option<Slot>),
+    ConfigOnly(Option<RpcLeaderScheduleConfig>),
+}
+
+impl RpcLeaderScheduleConfigWrapper {
+    pub fn unzip(&self) -> (Option<Slot>, Option<RpcLeaderScheduleConfig>) {
+        match &self {
+            RpcLeaderScheduleConfigWrapper::SlotOnly(slot) => (*slot, None),
+            RpcLeaderScheduleConfigWrapper::ConfigOnly(config) => (None, config.clone()),
+        }
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub enum RpcLargestAccountsFilter {
diff --git a/core/src/rpc.rs b/core/src/rpc.rs
index c19e163e0e9caf..593315bd4b4814 100644
--- a/core/src/rpc.rs
+++ b/core/src/rpc.rs
@@ -2108,7 +2108,7 @@ pub mod rpc_minimal {
         fn get_leader_schedule(
             &self,
             meta: Self::Metadata,
-            slot: Option<Slot>,
+            options: Option<RpcLeaderScheduleConfigWrapper>,
             config: Option<RpcLeaderScheduleConfig>,
         ) -> Result<Option<RpcLeaderSchedule>>;
     }
@@ -2213,10 +2213,11 @@ pub mod rpc_minimal {
         fn get_leader_schedule(
             &self,
             meta: Self::Metadata,
-            slot: Option<Slot>,
+            options: Option<RpcLeaderScheduleConfigWrapper>,
             config: Option<RpcLeaderScheduleConfig>,
         ) -> Result<Option<RpcLeaderSchedule>> {
-            let config = config.unwrap_or_default();
+            let (slot, maybe_config) = options.map(|options| options.unzip()).unwrap_or_default();
+            let config = maybe_config.or(config).unwrap_or_default();
 
             if let Some(ref identity) = config.identity {
                 let _ = verify_pubkey(identity)?;
@@ -4122,6 +4123,10 @@ pub mod tests {
                 r#"{{"jsonrpc":"2.0","id":1,"method":"getLeaderSchedule", "params": [null, {{ "identity": "{}" }}]}}"#,
                 bank.collector_id().to_string()
             ),
+            &format!(
+                r#"{{"jsonrpc":"2.0","id":1,"method":"getLeaderSchedule", "params": [{{ "identity": "{}" }}]}}"#,
+                bank.collector_id().to_string()
+            ),
         ]
         .iter()
         {
@@ -4171,7 +4176,7 @@ pub mod tests {
 
         // `bob` is not in the leader schedule, look for an empty response
         let req = format!(
-            r#"{{"jsonrpc":"2.0","id":1,"method":"getLeaderSchedule", "params": [null, {{ "identity": "{}"}}]}}"#,
+            r#"{{"jsonrpc":"2.0","id":1,"method":"getLeaderSchedule", "params": [{{ "identity": "{}"}}]}}"#,
             bob_pubkey
         );