Skip to content

Commit

Permalink
Optimize /etc/hosts writes (#259)
Browse files Browse the repository at this point in the history
* hostsfile: change internal map from hash to btree

This change makes the innernet section of /etc/hosts always ordered and
deterministic. We can take advantage of that to avoid writes, that will
be done in another commit.

* hostsfile: reduce number of writes if content hasn't changed

* hostsfile: return bool to inform if file has been written

This commit also makes the logs print accordingly to the new behavior.

* hostsfile: remove has_content_changed in favor of comparing old and new sections

* hostsfile: print the correct hosts path in log message

* hostsfile: remove unnecessary intermediate variable
  • Loading branch information
evaporei authored Jun 2, 2023
1 parent 33cee12 commit de7ec99
Show file tree
Hide file tree
Showing 2 changed files with 51 additions and 28 deletions.
15 changes: 10 additions & 5 deletions client/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -279,18 +279,23 @@ fn update_hosts_file(
hosts_path: PathBuf,
peers: &[Peer],
) -> Result<(), WrappedIoError> {
log::info!("updating {} with the latest peers.", "/etc/hosts".yellow());

let mut hosts_builder = HostsBuilder::new(format!("innernet {interface}"));
for peer in peers {
hosts_builder.add_hostname(
peer.contents.ip,
&format!("{}.{}.wg", peer.contents.name, interface),
);
}
if let Err(e) = hosts_builder.write_to(&hosts_path).with_path(hosts_path) {
log::warn!("failed to update hosts ({})", e);
}
match hosts_builder.write_to(&hosts_path).with_path(&hosts_path) {
Ok(has_written) if has_written => {
log::info!(
"updated {} with the latest peers.",
hosts_path.to_string_lossy().yellow()
)
},
Ok(_) => {},
Err(e) => log::warn!("failed to update hosts ({})", e),
};

Ok(())
}
Expand Down
64 changes: 41 additions & 23 deletions hostsfile/src/lib.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
use std::{
collections::HashMap,
collections::BTreeMap,
fmt,
fs::OpenOptions,
io::{self, BufRead, BufReader, ErrorKind, Write},
Expand Down Expand Up @@ -81,7 +81,7 @@ impl std::error::Error for Error {
/// ```
pub struct HostsBuilder {
tag: String,
hostname_map: HashMap<IpAddr, Vec<String>>,
hostname_map: BTreeMap<IpAddr, Vec<String>>,
}

impl HostsBuilder {
Expand All @@ -90,7 +90,7 @@ impl HostsBuilder {
pub fn new<S: Into<String>>(tag: S) -> Self {
Self {
tag: tag.into(),
hostname_map: HashMap::new(),
hostname_map: BTreeMap::new(),
}
}

Expand All @@ -116,7 +116,8 @@ impl HostsBuilder {

/// Inserts a new section to the system's default hosts file. If there is a section with the
/// same tag name already, it will be replaced with the new list instead.
pub fn write(&self) -> io::Result<()> {
/// Returns true if the hosts file has changed.
pub fn write(&self) -> io::Result<bool> {
self.write_to(Self::default_path()?)
}

Expand Down Expand Up @@ -178,7 +179,9 @@ impl HostsBuilder {
///
/// On Windows, the format of one hostname per line will be used, all other systems will use
/// the same format as Unix and Unix-like systems (i.e. allow multiple hostnames per line).
pub fn write_to<P: AsRef<Path>>(&self, hosts_path: P) -> io::Result<()> {
///
/// Returns true if the hosts file has changed.
pub fn write_to<P: AsRef<Path>>(&self, hosts_path: P) -> io::Result<bool> {
let hosts_path = hosts_path.as_ref();
if hosts_path.is_dir() {
// TODO(jake): use io::ErrorKind::IsADirectory when it's stable.
Expand Down Expand Up @@ -206,9 +209,31 @@ impl HostsBuilder {
let begin = lines.iter().position(|line| line.trim() == begin_marker);
let end = lines.iter().position(|line| line.trim() == end_marker);

let mut lines_to_insert = vec![];
if !self.hostname_map.is_empty() {
lines_to_insert.push(begin_marker);
for (ip, hostnames) in &self.hostname_map {
if cfg!(windows) {
// windows only allows one hostname per line
for hostname in hostnames {
lines_to_insert.push(format!("{ip} {hostname}"));
}
} else {
// assume the same format as Unix
lines_to_insert.push(format!("{} {}", ip, hostnames.join(" ")));
}
}
lines_to_insert.push(end_marker);
}

let insert = match (begin, end) {
(Some(begin), Some(end)) => {
lines.drain(begin..end + 1);
let old_section: Vec<String> = lines.drain(begin..end + 1).collect();

if old_section == lines_to_insert {
return Ok(false);
}

begin
},
(None, None) => {
Expand All @@ -233,21 +258,12 @@ impl HostsBuilder {
for line in &lines[..insert] {
writeln!(s, "{line}")?;
}
if !self.hostname_map.is_empty() {
writeln!(s, "{begin_marker}")?;
for (ip, hostnames) in &self.hostname_map {
if cfg!(windows) {
// windows only allows one hostname per line
for hostname in hostnames {
writeln!(s, "{ip} {hostname}")?;
}
} else {
// assume the same format as Unix
writeln!(s, "{} {}", ip, hostnames.join(" "))?;
}
}
writeln!(s, "{end_marker}")?;

// Append hostnames_map section
for line in lines_to_insert {
writeln!(s, "{line}")?;
}

for line in &lines[insert..] {
writeln!(s, "{line}")?;
}
Expand All @@ -260,8 +276,9 @@ impl HostsBuilder {
_ => {
log::debug!("wrote hosts file with the write-and-swap strategy");
},
}
Ok(())
};

Ok(true)
}

fn write_and_swap(temp_path: &Path, hosts_path: &Path, contents: &[u8]) -> io::Result<()> {
Expand Down Expand Up @@ -314,7 +331,8 @@ mod tests {
temp_file.write_all(b"preexisting\ncontent").unwrap();
let mut builder = HostsBuilder::new("foo");
builder.add_hostname([1, 1, 1, 1].into(), "whatever");
builder.write_to(&temp_path).unwrap();
assert!(builder.write_to(&temp_path).unwrap());
assert!(!builder.write_to(&temp_path).unwrap());

let contents = std::fs::read_to_string(&temp_path).unwrap();
println!("contents: {contents}");
Expand Down

0 comments on commit de7ec99

Please sign in to comment.