pd: bound every wait on the PD path — a frozen peer must not hang the client

Cargo.toml

@@ -59,7 +59,7 @@ rstest = "0.18.2"
 serde_json = "1"
 serial_test = "0.5.0"
 tempfile = "=3.14.0"
-tokio = { version = "1", features = ["sync", "rt-multi-thread", "macros"] }
+tokio = { version = "1", features = ["sync", "rt-multi-thread", "macros", "test-util"] }
 
 [[test]]
 name = "failpoint_tests"

src/common/security.rs

@@ -76,10 +76,27 @@ impl SecurityManager {
     }
 
     /// Connect to gRPC server using TLS connection. If TLS is not configured, use normal connection.
-    pub async fn connect<Factory, Client>(
+    ///
+    /// The dial is bounded by the default request timeout — see
+    /// [`Self::connect_with_timeout`] to choose the bound.
+    pub async fn connect<Factory, Client>(&self, addr: &str, factory: Factory) -> Result<Client>
+    where
+        Factory: FnOnce(Channel) -> Client,
+    {
+        self.connect_with_timeout(addr, crate::config::DEFAULT_REQUEST_TIMEOUT, factory)
+            .await
+    }
+
+    /// Like [`Self::connect`], with `connect_timeout` bounding the whole dial,
+    /// HTTP/2 handshake included. A frozen peer can complete the TCP handshake
+    /// from its kernel (which keeps running when the process does not) and then
+    /// never answer the HTTP/2 preface, hanging an unbounded connect until the
+    /// peer thaws (#516).
+    pub async fn connect_with_timeout<Factory, Client>(
         &self,
         // env: Arc<Environment>,
         addr: &str,
+        connect_timeout: Duration,
         factory: Factory,
     ) -> Result<Client>
     where
@@ -91,7 +108,15 @@ impl SecurityManager {
         } else {
             self.default_channel(addr).await?
         };
-        let ch = channel.connect().await?;
+        let ch = tokio::time::timeout(connect_timeout, channel.connect())
+            .await
+            .map_err(|_| {
+                internal_err!(
+                    "connecting to {} timed out after {:?}",
+                    addr,
+                    connect_timeout
+                )
+            })??;
 
         Ok(factory(ch))
     }

src/config.rs

@@ -28,7 +28,7 @@ pub struct Config {
     pub keyspace: Option<String>,
 }
 
-const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(2);
+pub(crate) const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(2);
 const DEFAULT_GRPC_MAX_DECODING_MESSAGE_SIZE: usize = 4 * 1024 * 1024; // 4MB
 
 impl Default for Config {

src/pd/cluster.rs

@@ -128,7 +128,7 @@ impl Connection {
         let members = self.validate_endpoints(endpoints, timeout).await?;
         let (client, keyspace_client, members) = self.try_connect_leader(&members, timeout).await?;
         let id = members.header.as_ref().unwrap().cluster_id;
-        let tso = TimestampOracle::new(id, &client)?;
+        let tso = TimestampOracle::new(id, &client, timeout)?;
         let cluster = Cluster {
             id,
             client,
@@ -145,7 +145,7 @@ impl Connection {
         let start = Instant::now();
         let (client, keyspace_client, members) =
             self.try_connect_leader(&cluster.members, timeout).await?;
-        let tso = TimestampOracle::new(cluster.id, &client)?;
+        let tso = TimestampOracle::new(cluster.id, &client, timeout)?;
         *cluster = Cluster {
             id: cluster.id,
             client,
@@ -213,27 +213,33 @@ impl Connection {
     async fn connect(
         &self,
         addr: &str,
-        _timeout: Duration,
+        timeout: Duration,
     ) -> Result<(
         pdpb::pd_client::PdClient<Channel>,
         keyspacepb::keyspace_client::KeyspaceClient<Channel>,
         pdpb::GetMembersResponse,
     )> {
         let mut client = self
             .security_mgr
-            .connect(addr, pdpb::pd_client::PdClient::<Channel>::new)
+            .connect_with_timeout(addr, timeout, pdpb::pd_client::PdClient::<Channel>::new)
             .await?;
         let keyspace_client = self
             .security_mgr
-            .connect(
+            .connect_with_timeout(
                 addr,
+                timeout,
                 keyspacepb::keyspace_client::KeyspaceClient::<Channel>::new,
             )
             .await?;
-        let resp: pdpb::GetMembersResponse = client
-            .get_members(pdpb::GetMembersRequest::default())
-            .await?
-            .into_inner();
+        // Bounded like every other request in this path: this is the first RPC
+        // on a fresh connection, exactly where a frozen peer would hang it.
+        let resp: pdpb::GetMembersResponse = tokio::time::timeout(
+            timeout,
+            client.get_members(pdpb::GetMembersRequest::default()),
+        )
+        .await
+        .map_err(|_| internal_err!("get_members timed out after {:?}", timeout))??
+        .into_inner();
         if let Some(err) = resp
             .header
             .as_ref()