Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Extract retry-read to new crate #1790

Merged
merged 1 commit into from
Nov 2, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion packages/os/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,8 @@ source-groups = [
"webpki-roots-shim",
"logdog",
"models",
"imdsclient"
"imdsclient",
"retry-read",
]

[lib]
Expand Down
8 changes: 8 additions & 0 deletions sources/Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 2 additions & 0 deletions sources/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,8 @@ members = [

"parse-datetime",

"retry-read",

"updater/block-party",
"updater/signpost",
"updater/update_metadata",
Expand Down
1 change: 1 addition & 0 deletions sources/api/early-boot-config/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ flate2 = { version = "1.0", default-features = false, features = ["rust_backend"
http = "0.2"
imdsclient = { path = "../../imdsclient", version = "0.1.0" }
log = "0.4"
retry-read = { path = "../../retry-read", version = "0.1.0" }
serde = { version = "1.0", features = ["derive"] }
serde_json = "1"
serde_plain = "1.0"
Expand Down
39 changes: 2 additions & 37 deletions sources/api/early-boot-config/src/compression.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,9 @@
//! Currently gzip compression is supported.

use flate2::read::GzDecoder;
use retry_read::RetryRead;
use std::fs::File;
use std::io::{BufReader, Chain, Cursor, ErrorKind, Read, Result, Take};
use std::io::{BufReader, Chain, Cursor, Read, Result, Take};
use std::path::Path;

/// "File magic" that indicates file type is stored in a few bytes at the start at the start of the
Expand Down Expand Up @@ -120,42 +121,6 @@ impl<R: Read> Read for OptionalCompressionReader<R> {

// =^..^= =^..^= =^..^= =^..^= =^..^= =^..^= =^..^= =^..^= =^..^= =^..^= =^..^=

/// This trait represents a `Read` operation where we want to retry after standard interruptions
/// (unlike `read()`) but also need to know the number of bytes we read (unlike `read_exact()`).
trait RetryRead<R> {
fn retry_read(&mut self, buf: &mut [u8]) -> Result<usize>;
}

impl<R: Read> RetryRead<R> for R {
// This implementation is based on stdlib Read::read_exact, but hitting EOF isn't a failure, we
// just want to return the number of bytes we could read.
fn retry_read(&mut self, mut buf: &mut [u8]) -> Result<usize> {
let mut count = 0;

// Read until we have no more space in the output buffer
while !buf.is_empty() {
match self.read(buf) {
// No bytes left, done
Ok(0) => break,
// Read n bytes, slide ahead n in the output buffer and read more
Ok(n) => {
count += n;
let tmp = buf;
buf = &mut tmp[n..];
}
// Retry on interrupt
Err(e) if e.kind() == ErrorKind::Interrupted => {}
// Other failures are fatal
Err(e) => return Err(e),
}
}

Ok(count)
}
}

// =^..^= =^..^= =^..^= =^..^= =^..^= =^..^= =^..^= =^..^= =^..^= =^..^= =^..^=

#[cfg(test)]
mod test {
use super::*;
Expand Down
12 changes: 12 additions & 0 deletions sources/retry-read/Cargo.toml
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
[package]
name = "retry-read"
version = "0.1.0"
authors = ["Tom Kirchner <[email protected]>"]
license = "Apache-2.0 OR MIT"
edition = "2018"
publish = false
# Don't rebuild crate just because of changes to README.
exclude = ["README.md"]

[build-dependencies]
cargo-readme = "3.1"
12 changes: 12 additions & 0 deletions sources/retry-read/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
# retry-read

Current version: 0.1.0

This library provides a `RetryRead` trait with a `retry_read` function that's available for any
`Read` type. `retry_read` retries after standard interruptions (unlike `read`) but also
returns the number of bytes read (unlike `read_exact`), and without needing to read to the end
of the input (unlike `read_to_end` and `read_to_string`).

## Colophon

This text was generated from `README.tpl` using [cargo-readme](https://crates.io/crates/cargo-readme), and includes the rustdoc from `src/lib.rs`.
9 changes: 9 additions & 0 deletions sources/retry-read/README.tpl
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
# {{crate}}

Current version: {{version}}

{{readme}}

## Colophon

This text was generated from `README.tpl` using [cargo-readme](https://crates.io/crates/cargo-readme), and includes the rustdoc from `src/lib.rs`.
32 changes: 32 additions & 0 deletions sources/retry-read/build.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
// Automatically generate README.md from rustdoc.

use std::env;
use std::fs::File;
use std::io::Write;
use std::path::PathBuf;

fn main() {
// Check for environment variable "SKIP_README". If it is set,
// skip README generation
if env::var_os("SKIP_README").is_some() {
return;
}

let mut source = File::open("src/lib.rs").unwrap();
let mut template = File::open("README.tpl").unwrap();

let content = cargo_readme::generate_readme(
&PathBuf::from("."), // root
&mut source, // source
Some(&mut template), // template
// The "add x" arguments don't apply when using a template.
true, // add title
false, // add badges
false, // add license
true, // indent headings
)
.unwrap();

let mut readme = File::create("README.md").unwrap();
readme.write_all(content.as_bytes()).unwrap();
}
114 changes: 114 additions & 0 deletions sources/retry-read/src/lib.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,114 @@
//! This library provides a `RetryRead` trait with a `retry_read` function that's available for any
//! `Read` type. `retry_read` retries after standard interruptions (unlike `read`) but also
//! returns the number of bytes read (unlike `read_exact`), and without needing to read to the end
//! of the input (unlike `read_to_end` and `read_to_string`).

use std::io::{ErrorKind, Read, Result};

/// Provides a way to retry standard read operations while also returning the number of bytes read.
pub trait RetryRead<R> {
fn retry_read(&mut self, buf: &mut [u8]) -> Result<usize>;
}

impl<R: Read> RetryRead<R> for R {
// This implementation is based on stdlib Read::read_exact, but hitting EOF isn't a failure, we
// just want to return the number of bytes we could read.
/// Like `Read::read` but retries on ErrorKind::Interrupted, returning the number of bytes read.
webern marked this conversation as resolved.
Show resolved Hide resolved
fn retry_read(&mut self, mut buf: &mut [u8]) -> Result<usize> {
let mut count = 0;

// Read until we have no more space in the output buffer
while !buf.is_empty() {
match self.read(buf) {
// No bytes left, done
Ok(0) => break,
// Read n bytes, slide ahead n in the output buffer and read more
Ok(n) => {
count += n;
let tmp = buf;
buf = &mut tmp[n..];
}
// Retry on interrupt
Err(e) if e.kind() == ErrorKind::Interrupted => {}
// Other failures are fatal
Err(e) => return Err(e),
}
}

Ok(count)
}
}

#[cfg(test)]
mod test {
use super::{ErrorKind, Read, Result, RetryRead};
use std::io::{Error, Write};

// Helper method for simple test cases, confirming we read the full given slice.
fn test(data: &[u8]) {
let mut output = vec![0; data.len()];
let count = (&data[..]).retry_read(&mut output).unwrap();
assert_eq!(count, data.len());
assert_eq!(&data[..], &output);
}

#[test]
fn zero_read() {
test(&[]);
}

#[test]
fn small_read() {
test(&[0, 1, 2, 3, 42]);
}

#[test]
fn large_read() {
test(&[42; 9999]);
}

// Confirm we retry reads when interrupted.
#[test]
fn retried_read() {
let mut reader = InterruptedReader::new(5);
let mut output = vec![0; 5];
let count = reader.retry_read(&mut output).unwrap();
assert_eq!(count, 5);
assert_eq!(output, vec![42, 42, 42, 42, 42]);
}

// Helper that implements Read, eventually returning the requested number of bytes, but returns
// ErrorKind::Interrupted every other call.
struct InterruptedReader {
requested_reads: u64,
finished_reads: u64,
interrupt: bool,
}

impl InterruptedReader {
fn new(requested_reads: u64) -> Self {
Self {
requested_reads,
finished_reads: 0,
interrupt: false,
}
}
}

impl Read for InterruptedReader {
fn read(&mut self, mut buf: &mut [u8]) -> Result<usize> {
if self.finished_reads > self.requested_reads {
return Ok(0);
}

if self.interrupt {
self.interrupt = false;
Err(Error::new(ErrorKind::Interrupted, "you asked for it"))
} else {
self.interrupt = true;
self.finished_reads += 1;
buf.write(&[42])
}
}
}
}