Skip to content

Commit

Permalink
Extract retry-read to new crate
Browse files Browse the repository at this point in the history
  • Loading branch information
tjkirch committed Nov 1, 2021
1 parent 97400c3 commit 0b6baa5
Show file tree
Hide file tree
Showing 10 changed files with 194 additions and 38 deletions.
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.
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])
}
}
}
}

0 comments on commit 0b6baa5

Please sign in to comment.