From bcff6ae82a14c0f70c0a1957b511095fe891e486 Mon Sep 17 00:00:00 2001
From: Ulrich Hornung <hornunguli@gmx.de>
Date: Tue, 14 May 2024 22:22:38 +0200
Subject: [PATCH] prevent PANIC when leader is None

happend in tests with tifs:
thread 'tokio-runtime-worker' panicked at /home/uli/.cargo/registry/src/index.crates.io-6f17d22bba15001f/tikv-client-0.3.0/src/pd/cluster.rs:270:47:
called `Option::unwrap()` on a `None` value
---
 src/common/errors.rs |  4 +++-
 src/pd/cluster.rs    | 10 ++++++----
 src/region.rs        |  2 +-
 src/store/request.rs |  2 +-
 4 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/src/common/errors.rs b/src/common/errors.rs
index 5e7f6303..29189e33 100644
--- a/src/common/errors.rs
+++ b/src/common/errors.rs
@@ -91,7 +91,9 @@ pub enum Error {
     RegionNotFoundInResponse { region_id: u64 },
     /// No leader is found for the given id.
     #[error("Leader of region {} is not found", region_id)]
-    LeaderNotFound { region_id: u64 },
+    LeaderOfRegionNotFound { region_id: u64 },
+    #[error("Leader of cluster {} is not found", cluster_id)]
+    LeaderOfClusterNotFound { cluster_id: u64 },
     /// Scan limit exceeds the maximum
     #[error("Limit {} exceeds max scan limit {}", limit, max_limit)]
     MaxScanLimitExceeded { limit: u32, max_limit: u32 },
diff --git a/src/pd/cluster.rs b/src/pd/cluster.rs
index e74a78fb..58d7ca24 100644
--- a/src/pd/cluster.rs
+++ b/src/pd/cluster.rs
@@ -279,7 +279,7 @@ impl Connection {
         keyspacepb::keyspace_client::KeyspaceClient<Channel>,
         pdpb::GetMembersResponse,
     )> {
-        let previous_leader = previous.leader.as_ref().unwrap();
+        let previous_leader = previous.leader.as_ref();
         let members = &previous.members;
         let cluster_id = previous.header.as_ref().unwrap().cluster_id;
 
@@ -287,8 +287,8 @@ impl Connection {
         // Try to connect to other members, then the previous leader.
         'outer: for m in members
             .iter()
-            .filter(|m| *m != previous_leader)
-            .chain(Some(previous_leader))
+            .filter(|m| !previous_leader.is_some_and(|pl| *m == pl))
+            .chain(previous_leader)
         {
             for ep in &m.client_urls {
                 match self.try_connect(ep.as_str(), cluster_id, timeout).await {
@@ -306,7 +306,9 @@ impl Connection {
 
         // Then try to connect the PD cluster leader.
         if let Some(resp) = resp {
-            let leader = resp.leader.as_ref().unwrap();
+            let leader = resp.leader.as_ref().ok_or(Error::LeaderOfClusterNotFound{
+                cluster_id,
+            })?;
             for ep in &leader.client_urls {
                 if let Ok((client, keyspace_client, members)) =
                     self.try_connect(ep.as_str(), cluster_id, timeout).await
diff --git a/src/region.rs b/src/region.rs
index 6fb20321..a6fe5514 100644
--- a/src/region.rs
+++ b/src/region.rs
@@ -72,7 +72,7 @@ impl RegionWithLeader {
         self.leader
             .as_ref()
             .cloned()
-            .ok_or_else(|| Error::LeaderNotFound {
+            .ok_or_else(|| Error::LeaderOfRegionNotFound {
                 region_id: self.id(),
             })
             .map(|s| s.store_id)
diff --git a/src/store/request.rs b/src/store/request.rs
index 9eaabe5f..419c71b3 100644
--- a/src/store/request.rs
+++ b/src/store/request.rs
@@ -55,7 +55,7 @@ macro_rules! impl_request {
 
             fn set_leader(&mut self, leader: &RegionWithLeader) -> Result<()> {
                 let ctx = self.context.get_or_insert(kvrpcpb::Context::default());
-                let leader_peer = leader.leader.as_ref().ok_or(Error::LeaderNotFound {
+                let leader_peer = leader.leader.as_ref().ok_or(Error::LeaderOfRegionNotFound {
                     region_id: leader.region.id,
                 })?;
                 ctx.region_id = leader.region.id;