Optimize /etc/hosts writes (#259)

* 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
pull/266/head
Eva Pace 2023-06-02 03:42:54 -03:00 committed by GitHub
parent 33cee129d1
commit de7ec99ad9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 51 additions and 28 deletions

View File

@ -279,8 +279,6 @@ fn update_hosts_file(
hosts_path: PathBuf, hosts_path: PathBuf,
peers: &[Peer], peers: &[Peer],
) -> Result<(), WrappedIoError> { ) -> Result<(), WrappedIoError> {
log::info!("updating {} with the latest peers.", "/etc/hosts".yellow());
let mut hosts_builder = HostsBuilder::new(format!("innernet {interface}")); let mut hosts_builder = HostsBuilder::new(format!("innernet {interface}"));
for peer in peers { for peer in peers {
hosts_builder.add_hostname( hosts_builder.add_hostname(
@ -288,9 +286,16 @@ fn update_hosts_file(
&format!("{}.{}.wg", peer.contents.name, interface), &format!("{}.{}.wg", peer.contents.name, interface),
); );
} }
if let Err(e) = hosts_builder.write_to(&hosts_path).with_path(hosts_path) { match hosts_builder.write_to(&hosts_path).with_path(&hosts_path) {
log::warn!("failed to update hosts ({})", e); 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(()) Ok(())
} }

View File

@ -1,5 +1,5 @@
use std::{ use std::{
collections::HashMap, collections::BTreeMap,
fmt, fmt,
fs::OpenOptions, fs::OpenOptions,
io::{self, BufRead, BufReader, ErrorKind, Write}, io::{self, BufRead, BufReader, ErrorKind, Write},
@ -81,7 +81,7 @@ impl std::error::Error for Error {
/// ``` /// ```
pub struct HostsBuilder { pub struct HostsBuilder {
tag: String, tag: String,
hostname_map: HashMap<IpAddr, Vec<String>>, hostname_map: BTreeMap<IpAddr, Vec<String>>,
} }
impl HostsBuilder { impl HostsBuilder {
@ -90,7 +90,7 @@ impl HostsBuilder {
pub fn new<S: Into<String>>(tag: S) -> Self { pub fn new<S: Into<String>>(tag: S) -> Self {
Self { Self {
tag: tag.into(), tag: tag.into(),
hostname_map: HashMap::new(), hostname_map: BTreeMap::new(),
} }
} }
@ -116,7 +116,8 @@ impl HostsBuilder {
/// Inserts a new section to the system's default hosts file. If there is a section with the /// 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. /// 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()?) self.write_to(Self::default_path()?)
} }
@ -178,7 +179,9 @@ impl HostsBuilder {
/// ///
/// On Windows, the format of one hostname per line will be used, all other systems will use /// 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). /// 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(); let hosts_path = hosts_path.as_ref();
if hosts_path.is_dir() { if hosts_path.is_dir() {
// TODO(jake): use io::ErrorKind::IsADirectory when it's stable. // TODO(jake): use io::ErrorKind::IsADirectory when it's stable.
@ -206,9 +209,31 @@ impl HostsBuilder {
let begin = lines.iter().position(|line| line.trim() == begin_marker); let begin = lines.iter().position(|line| line.trim() == begin_marker);
let end = lines.iter().position(|line| line.trim() == end_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) { let insert = match (begin, end) {
(Some(begin), Some(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 begin
}, },
(None, None) => { (None, None) => {
@ -233,21 +258,12 @@ impl HostsBuilder {
for line in &lines[..insert] { for line in &lines[..insert] {
writeln!(s, "{line}")?; writeln!(s, "{line}")?;
} }
if !self.hostname_map.is_empty() {
writeln!(s, "{begin_marker}")?; // Append hostnames_map section
for (ip, hostnames) in &self.hostname_map { for line in lines_to_insert {
if cfg!(windows) { writeln!(s, "{line}")?;
// 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}")?;
} }
for line in &lines[insert..] { for line in &lines[insert..] {
writeln!(s, "{line}")?; writeln!(s, "{line}")?;
} }
@ -260,8 +276,9 @@ impl HostsBuilder {
_ => { _ => {
log::debug!("wrote hosts file with the write-and-swap strategy"); 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<()> { fn write_and_swap(temp_path: &Path, hosts_path: &Path, contents: &[u8]) -> io::Result<()> {
@ -314,7 +331,8 @@ mod tests {
temp_file.write_all(b"preexisting\ncontent").unwrap(); temp_file.write_all(b"preexisting\ncontent").unwrap();
let mut builder = HostsBuilder::new("foo"); let mut builder = HostsBuilder::new("foo");
builder.add_hostname([1, 1, 1, 1].into(), "whatever"); 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(); let contents = std::fs::read_to_string(&temp_path).unwrap();
println!("contents: {contents}"); println!("contents: {contents}");