f - Reconnect and retry HTTP requests

jkczyz · jkczyz · commit 91691a8697de · 2021-01-15T19:18:54.000-08:00
diff --git a/lightning-block-sync/src/http.rs b/lightning-block-sync/src/http.rs
@@ -140,9 +140,8 @@ impl HttpClient {
 			 Host: {}\r\n\
 			 Connection: keep-alive\r\n\
 			 \r\n", uri, host);
-		self.write_request(request).await?;
-		let bytes = self.read_response().await?;
-		F::try_from(bytes)
+		let response_body = self.send_request_with_retry(&request).await?;
+		F::try_from(response_body)
 	}
 
 	/// Sends a `POST` request for a resource identified by `uri` at the `host` using the given HTTP
@@ -162,13 +161,37 @@ impl HttpClient {
 			 Content-Length: {}\r\n\
 			 \r\n\
 			 {}", uri, host, auth, content.len(), content);
+		let response_body = self.send_request_with_retry(&request).await?;
+		F::try_from(response_body)
+	}
+
+	/// Sends an HTTP request message and reads the response, returning its body. Attempts to
+	/// reconnect and retry if the connection has been closed.
+	async fn send_request_with_retry(&mut self, request: &str) -> std::io::Result<Vec<u8>> {
+		let endpoint = self.stream.peer_addr().unwrap();
+		match self.send_request(request).await {
+			Ok(bytes) => Ok(bytes),
+			Err(e) => match e.kind() {
+				std::io::ErrorKind::ConnectionReset |
+				std::io::ErrorKind::ConnectionAborted |
+				std::io::ErrorKind::UnexpectedEof => {
+					// Reconnect if the connection was closed.
+					*self = Self::connect(endpoint)?;
+					self.send_request(request).await
+				},
+				_ => Err(e),
+			},
+		}
+	}
+
+	/// Sends an HTTP request message and reads the response, returning its body.
+	async fn send_request(&mut self, request: &str) -> std::io::Result<Vec<u8>> {
 		self.write_request(request).await?;
-		let bytes = self.read_response().await?;
-		F::try_from(bytes)
+		self.read_response().await
 	}
 
 	/// Writes an HTTP request message.
-	async fn write_request(&mut self, request: String) -> std::io::Result<()> {
+	async fn write_request(&mut self, request: &str) -> std::io::Result<()> {
 		#[cfg(feature = "tokio")]
 		{
 			self.stream.write_all(request.as_bytes()).await?;
@@ -214,14 +237,14 @@ impl HttpClient {
 
 		// Read and parse status line
 		let status_line = read_line!()
-			.ok_or(std::io::Error::new(std::io::ErrorKind::InvalidData, "no status line"))?;
+			.ok_or(std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "no status line"))?;
 		let status = HttpStatus::parse(&status_line)?;
 
 		// Read and parse relevant headers
 		let mut message_length = HttpMessageLength::Empty;
 		loop {
 			let line = read_line!()
-				.ok_or(std::io::Error::new(std::io::ErrorKind::InvalidData, "unexpected eof"))?;
+				.ok_or(std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "no headers"))?;
 			if line.is_empty() { break; }
 
 			let header = HttpHeader::parse(&line)?;
@@ -512,21 +535,23 @@ pub(crate) mod client_tests {
 			let shutdown = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
 			let shutdown_signaled = std::sync::Arc::clone(&shutdown);
 			let handler = std::thread::spawn(move || {
-				let (mut stream, _) = listener.accept().unwrap();
-				stream.set_write_timeout(Some(Duration::from_secs(1))).unwrap();
-
-				let lines_read = std::io::BufReader::new(&stream)
-					.lines()
-					.take_while(|line| !line.as_ref().unwrap().is_empty())
-					.count();
-				if lines_read == 0 { return; }
-
-				for chunk in response.as_bytes().chunks(16) {
-					if shutdown_signaled.load(std::sync::atomic::Ordering::SeqCst) {
-						break;
-					} else {
-						stream.write(chunk).unwrap();
-						stream.flush().unwrap();
+				for stream in listener.incoming() {
+					let mut stream = stream.unwrap();
+					stream.set_write_timeout(Some(Duration::from_secs(1))).unwrap();
+
+					let lines_read = std::io::BufReader::new(&stream)
+						.lines()
+						.take_while(|line| !line.as_ref().unwrap().is_empty())
+						.count();
+					if lines_read == 0 { continue; }
+
+					for chunk in response.as_bytes().chunks(16) {
+						if shutdown_signaled.load(std::sync::atomic::Ordering::SeqCst) {
+							return;
+						} else {
+							stream.write(chunk).unwrap();
+							stream.flush().unwrap();
+						}
 					}
 				}
 			});
@@ -587,7 +612,7 @@ pub(crate) mod client_tests {
 		drop(server);
 		match client.get::<BinaryResponse>("/foo", "foo.com").await {
 			Err(e) => {
-				assert_eq!(e.kind(), std::io::ErrorKind::InvalidData);
+				assert_eq!(e.kind(), std::io::ErrorKind::UnexpectedEof);
 				assert_eq!(e.get_ref().unwrap().to_string(), "no status line");
 			},
 			Ok(_) => panic!("Expected error"),
@@ -602,8 +627,8 @@ pub(crate) mod client_tests {
 		drop(server);
 		match client.get::<BinaryResponse>("/foo", "foo.com").await {
 			Err(e) => {
-				assert_eq!(e.kind(), std::io::ErrorKind::InvalidData);
-				assert_eq!(e.get_ref().unwrap().to_string(), "unexpected eof");
+				assert_eq!(e.kind(), std::io::ErrorKind::UnexpectedEof);
+				assert_eq!(e.get_ref().unwrap().to_string(), "no headers");
 			},
 			Ok(_) => panic!("Expected error"),
 		}
@@ -620,8 +645,8 @@ pub(crate) mod client_tests {
 		let mut client = HttpClient::connect(&server.endpoint()).unwrap();
 		match client.get::<BinaryResponse>("/foo", "foo.com").await {
 			Err(e) => {
-				assert_eq!(e.kind(), std::io::ErrorKind::InvalidData);
-				assert_eq!(e.get_ref().unwrap().to_string(), "unexpected eof");
+				assert_eq!(e.kind(), std::io::ErrorKind::UnexpectedEof);
+				assert_eq!(e.get_ref().unwrap().to_string(), "no headers");
 			},
 			Ok(_) => panic!("Expected error"),
 		}
@@ -699,6 +724,18 @@ pub(crate) mod client_tests {
 		}
 	}
 
+	#[tokio::test]
+	async fn reconnect_closed_connection() {
+		let server = HttpServer::responding_with_ok::<String>(MessageBody::Empty);
+
+		let mut client = HttpClient::connect(&server.endpoint()).unwrap();
+		assert!(client.get::<BinaryResponse>("/foo", "foo.com").await.is_ok());
+		match client.get::<BinaryResponse>("/foo", "foo.com").await {
+			Err(e) => panic!("Unexpected error: {:?}", e),
+			Ok(bytes) => assert_eq!(bytes.0, Vec::<u8>::new()),
+		}
+	}
+
 	#[test]
 	fn from_bytes_into_binary_response() {
 		let bytes = b"foo";
diff --git a/lightning-block-sync/src/rest.rs b/lightning-block-sync/src/rest.rs
@@ -6,21 +6,21 @@ use std::convert::TryInto;
 /// A simple REST client for requesting resources using HTTP `GET`.
 pub struct RESTClient {
 	endpoint: HttpEndpoint,
+	client: HttpClient,
 }
 
 impl RESTClient {
-	pub fn new(endpoint: HttpEndpoint) -> Self {
-		Self { endpoint }
+	pub fn new(endpoint: HttpEndpoint) -> std::io::Result<Self> {
+		let client = HttpClient::connect(&endpoint)?;
+		Ok(Self { endpoint, client })
 	}
 
 	/// Requests a resource encoded in `F` format and interpreted as type `T`.
-	async fn request_resource<F, T>(&self, resource_path: &str) -> std::io::Result<T>
+	async fn request_resource<F, T>(&mut self, resource_path: &str) -> std::io::Result<T>
 	where F: TryFrom<Vec<u8>, Error = std::io::Error> + TryInto<T, Error = std::io::Error> {
 		let host = format!("{}:{}", self.endpoint.host(), self.endpoint.port());
 		let uri = format!("{}/{}", self.endpoint.path().trim_end_matches("/"), resource_path);
-
-		let mut client = HttpClient::connect(&self.endpoint)?;
-		client.get::<F>(&uri, &host).await?.try_into()
+		self.client.get::<F>(&uri, &host).await?.try_into()
 	}
 }
 
@@ -48,7 +48,7 @@ mod tests {
 	#[tokio::test]
 	async fn request_unknown_resource() {
 		let server = HttpServer::responding_with_not_found();
-		let client = RESTClient::new(server.endpoint());
+		let mut client = RESTClient::new(server.endpoint()).unwrap();
 
 		match client.request_resource::<BinaryResponse, u32>("/").await {
 			Err(e) => assert_eq!(e.kind(), std::io::ErrorKind::NotFound),
@@ -59,7 +59,7 @@ mod tests {
 	#[tokio::test]
 	async fn request_malformed_resource() {
 		let server = HttpServer::responding_with_ok(MessageBody::Content("foo"));
-		let client = RESTClient::new(server.endpoint());
+		let mut client = RESTClient::new(server.endpoint()).unwrap();
 
 		match client.request_resource::<BinaryResponse, u32>("/").await {
 			Err(e) => assert_eq!(e.kind(), std::io::ErrorKind::InvalidData),
@@ -70,7 +70,7 @@ mod tests {
 	#[tokio::test]
 	async fn request_valid_resource() {
 		let server = HttpServer::responding_with_ok(MessageBody::Content(42));
-		let client = RESTClient::new(server.endpoint());
+		let mut client = RESTClient::new(server.endpoint()).unwrap();
 
 		match client.request_resource::<BinaryResponse, u32>("/").await {
 			Err(e) => panic!("Unexpected error: {:?}", e),
diff --git a/lightning-block-sync/src/rpc.rs b/lightning-block-sync/src/rpc.rs
@@ -11,20 +11,23 @@ use std::sync::atomic::{AtomicUsize, Ordering};
 pub struct RPCClient {
 	basic_auth: String,
 	endpoint: HttpEndpoint,
+	client: HttpClient,
 	id: AtomicUsize,
 }
 
 impl RPCClient {
-	pub fn new(user_auth: &str, endpoint: HttpEndpoint) -> Self {
-		Self {
+	pub fn new(user_auth: &str, endpoint: HttpEndpoint) -> std::io::Result<Self> {
+		let client = HttpClient::connect(&endpoint)?;
+		Ok(Self {
 			basic_auth: "Basic ".to_string() + &base64::encode(user_auth),
 			endpoint,
+			client,
 			id: AtomicUsize::new(0),
-		}
+		})
 	}
 
 	/// Calls a method with the response encoded in JSON format and interpreted as type `T`.
-	async fn call_method<T>(&self, method: &str, params: &[serde_json::Value]) -> std::io::Result<T>
+	async fn call_method<T>(&mut self, method: &str, params: &[serde_json::Value]) -> std::io::Result<T>
 	where JsonResponse: TryFrom<Vec<u8>, Error = std::io::Error> + TryInto<T, Error = std::io::Error> {
 		let host = format!("{}:{}", self.endpoint.host(), self.endpoint.port());
 		let uri = self.endpoint.path();
@@ -34,8 +37,8 @@ impl RPCClient {
 			"id": &self.id.fetch_add(1, Ordering::AcqRel).to_string()
 		});
 
-		let mut client = HttpClient::connect(&self.endpoint)?;
-		let mut response = client.post::<JsonResponse>(&uri, &host, &self.basic_auth, content).await?.0;
+		let mut response = self.client.post::<JsonResponse>(&uri, &host, &self.basic_auth, content)
+			.await?.0;
 		if !response.is_object() {
 			return Err(std::io::Error::new(std::io::ErrorKind::InvalidData, "expected JSON object"));
 		}
@@ -76,7 +79,7 @@ mod tests {
 	#[tokio::test]
 	async fn call_method_returning_unknown_response() {
 		let server = HttpServer::responding_with_not_found();
-		let client = RPCClient::new("credentials", server.endpoint());
+		let mut client = RPCClient::new("credentials", server.endpoint()).unwrap();
 
 		match client.call_method::<u64>("getblockcount", &[]).await {
 			Err(e) => assert_eq!(e.kind(), std::io::ErrorKind::NotFound),
@@ -88,7 +91,7 @@ mod tests {
 	async fn call_method_returning_malfomred_response() {
 		let response = serde_json::json!("foo");
 		let server = HttpServer::responding_with_ok(MessageBody::Content(response));
-		let client = RPCClient::new("credentials", server.endpoint());
+		let mut client = RPCClient::new("credentials", server.endpoint()).unwrap();
 
 		match client.call_method::<u64>("getblockcount", &[]).await {
 			Err(e) => {
@@ -105,7 +108,7 @@ mod tests {
 			"error": { "code": -8, "message": "invalid parameter" },
 		});
 		let server = HttpServer::responding_with_ok(MessageBody::Content(response));
-		let client = RPCClient::new("credentials", server.endpoint());
+		let mut client = RPCClient::new("credentials", server.endpoint()).unwrap();
 
 		let invalid_block_hash = serde_json::json!("foo");
 		match client.call_method::<u64>("getblock", &[invalid_block_hash]).await {
@@ -121,7 +124,7 @@ mod tests {
 	async fn call_method_returning_missing_result() {
 		let response = serde_json::json!({ "result": null });
 		let server = HttpServer::responding_with_ok(MessageBody::Content(response));
-		let client = RPCClient::new("credentials", server.endpoint());
+		let mut client = RPCClient::new("credentials", server.endpoint()).unwrap();
 
 		match client.call_method::<u64>("getblockcount", &[]).await {
 			Err(e) => {
@@ -136,7 +139,7 @@ mod tests {
 	async fn call_method_returning_valid_result() {
 		let response = serde_json::json!({ "result": 654470 });
 		let server = HttpServer::responding_with_ok(MessageBody::Content(response));
-		let client = RPCClient::new("credentials", server.endpoint());
+		let mut client = RPCClient::new("credentials", server.endpoint()).unwrap();
 
 		match client.call_method::<u64>("getblockcount", &[]).await {
 			Err(e) => panic!("Unexpected error: {:?}", e),
