1  #![allow(clippy::needless_lifetimes)]
 2  
 3  //! A library for cooperative port allocation between multiple processes.
 4  //!
 5  //! Fedimint tests in many places need to allocate ranges of unused ports for
 6  //! Federations and other software under tests, without being able to `bind`
 7  //! them beforehand.
 8  //!
 9  //! We used to mitigate that using a global per-process atomic counter, as
10  //! as simple port allocation mechanism. But this does not prevent conflicts
11  //! between different processes.
12  //!
13  //! Normally this would prevent us from running multiple tests at the same time,
14  //! which also makes it impossible to use `cargo nextest`.
15  //!
16  //! This library keeps track of allocated ports (with an expiration timeout) in
17  //! a shared file, protected by an advisory fs lock, and uses `bind` to make
18  //! sure a given port is actually free
19  
20  mod data;
21  mod envs;
22  mod util;
23  
24  use std::path::PathBuf;
25  
26  use anyhow::bail;
27  
28  use crate::data::DataDir;
29  use crate::envs::FM_PORTALLOC_DATA_DIR_ENV;
30  
31  pub fn port_alloc(range_size: u16) -> anyhow::Result<u16> {
32      if range_size == 0 {
33          bail!("Can't allocate range of 0 ports");
34      }
35  
36      let mut data_dir = DataDir::new(data_dir()?)?;
37  
38      data_dir.with_lock(|data_dir| {
39          let mut data = data_dir.load_data()?;
40          let base_port = data.get_free_port_range(range_size);
41          data_dir.store_data(&data)?;
42          Ok(base_port)
43      })
44  }
45  
46  fn data_dir() -> anyhow::Result<PathBuf> {
47      if let Some(env) = std::env::var_os(FM_PORTALLOC_DATA_DIR_ENV) {
48          Ok(PathBuf::from(env))
49      } else if let Some(dir) = dirs::cache_dir() {
50          Ok(dir.join("fm-portalloc"))
51      } else {
52          bail!("Could not determine port alloc data dir. Try setting FM_PORTALLOC_DATA_DIR");
53      }
54  }