use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::Duration;
use tf_tree_arena::AttachMode;
use tf_tree_ipc::{
boot_id, self_start_time, AccessMode, ArenaName, EnvVar, HelloRequest, HelloStatus, IpcError,
OpenOutcome, OwnerServer, Rendezvous, RuntimeDir, SegmentDescriptor, ShutdownHandle,
SocketProbe, SystemEnv, DEFAULT_OPEN_TIMEOUT,
};
use crate::tree::{BuildError, Tree, TreeBuilder, MAX_BACKOFF, MIN_BACKOFF};
pub use tf_tree_ipc::CreatePolicy;
fn open_claim_lock(rv: &Rendezvous) -> Result<std::sync::Arc<tf_tree_ipc::LockFile>, OpenError> {
Ok(std::sync::Arc::new(
tf_tree_ipc::LockFile::open(rv.lock_path()).map_err(OpenError::Rendezvous)?,
))
}
pub(crate) struct LivenessProbe {
lock: tf_tree_ipc::LockFile,
}
impl LivenessProbe {
fn open(rv: &Rendezvous) -> Result<LivenessProbe, IpcError> {
Ok(LivenessProbe {
lock: tf_tree_ipc::LockFile::open(rv.lock_path())?,
})
}
pub(crate) fn is_held(&self, slot: u32) -> Option<bool> {
self.lock.probe_participant(slot).ok().map(|p| p.held)
}
}
pub(crate) type JoinedSession = tf_tree_ipc::Session<tf_tree_ipc::Attached>;
pub(crate) enum Attachment {
Joined {
_session: JoinedSession,
_socket: std::os::fd::OwnedFd,
},
Owner {
_session: JoinedSession,
server: OwnerThread,
},
}
impl Drop for Attachment {
fn drop(&mut self) {
if let Attachment::Owner { server, .. } = self {
server.stop();
}
}
}
const _: () = assert!(
tf_tree_ipc::MAX_PARTICIPANTS == tf_tree_arena::DEFAULT_MAX_PARTICIPANTS,
"the lock file and the arena must agree on the participant slot space"
);
#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)]
#[non_exhaustive]
pub enum OpenError {
#[error("{0}")]
Rendezvous(IpcError),
#[error("{0:?}")]
Map(tf_tree_arena::ShmError),
#[error("{0}")]
Build(BuildError),
#[error("no layout was supplied and the arena had to be created")]
NoLayoutToCreate,
#[error("a read-only attach cannot create an arena: use CreatePolicy::Never, or AttachMode::ReadWrite")]
ReadOnlyCannotCreate,
#[error("an arena is already live at this rendezvous and require_create was set")]
ArenaAlreadyLive,
}
impl From<IpcError> for OpenError {
fn from(e: IpcError) -> OpenError {
OpenError::Rendezvous(e)
}
}
impl From<tf_tree_arena::ShmError> for OpenError {
fn from(e: tf_tree_arena::ShmError) -> OpenError {
OpenError::Map(e)
}
}
impl From<BuildError> for OpenError {
fn from(e: BuildError) -> OpenError {
OpenError::Build(e)
}
}
pub fn open() -> Result<Tree, OpenError> {
Open::new().open()
}
pub struct Open {
domain: Option<u32>,
name: Option<ArenaName>,
mode: AttachMode,
create: CreatePolicy,
timeout: Duration,
layout: Option<TreeBuilder>,
require_create: bool,
}
impl Default for Open {
fn default() -> Open {
Open::new()
}
}
impl Open {
#[must_use]
pub fn new() -> Open {
Open {
domain: None,
name: None,
mode: AttachMode::ReadOnly,
create: CreatePolicy::Never,
timeout: DEFAULT_OPEN_TIMEOUT,
layout: None,
require_create: false,
}
}
#[must_use]
pub fn domain(mut self, domain: u32) -> Open {
self.domain = Some(domain);
self
}
pub fn name(mut self, name: &str) -> Result<Open, OpenError> {
self.name = Some(ArenaName::new(name, EnvVar::Name).map_err(OpenError::Rendezvous)?);
Ok(self)
}
#[must_use]
pub fn mode(mut self, mode: AttachMode) -> Open {
self.mode = mode;
self
}
#[must_use]
pub fn create(mut self, create: CreatePolicy) -> Open {
self.create = create;
self
}
#[must_use]
pub fn require_create(mut self, require: bool) -> Open {
self.require_create = require;
self
}
#[must_use]
pub fn timeout(mut self, timeout: Duration) -> Open {
self.timeout = timeout;
self
}
#[must_use]
pub fn layout_if_creating(mut self, builder: TreeBuilder) -> Open {
self.layout = Some(builder);
self
}
pub fn open(mut self) -> Result<Tree, OpenError> {
let per_attempt = self.timeout;
self.attempt(per_attempt)
}
pub fn await_open(mut self, timeout: Duration) -> Result<Tree, OpenError> {
let start = std::time::Instant::now();
let mut backoff = MIN_BACKOFF;
loop {
let left = timeout.saturating_sub(start.elapsed());
let per_attempt = core::cmp::max(core::cmp::min(self.timeout, left), MIN_BACKOFF);
let err = match self.attempt(per_attempt) {
Ok(tree) => return Ok(tree),
Err(e) if is_retryable(e) => e,
Err(e) => return Err(e),
};
if start.elapsed() >= timeout {
return Err(err);
}
let left = timeout.saturating_sub(start.elapsed());
std::thread::sleep(core::cmp::min(backoff, left));
backoff = core::cmp::min(backoff * 2, MAX_BACKOFF);
}
}
fn attempt(&mut self, per_attempt: Duration) -> Result<Tree, OpenError> {
if self.mode == AttachMode::ReadOnly && self.create != CreatePolicy::Never {
return Err(OpenError::ReadOnlyCannotCreate);
}
let rd = RuntimeDir::resolve().map_err(OpenError::Rendezvous)?;
let domain = match self.domain {
Some(d) => d,
None => tf_tree_ipc::domain_from_env(&SystemEnv).map_err(OpenError::Rendezvous)?,
};
let name = match self.name {
Some(n) => n,
None => tf_tree_ipc::name_from_env(&SystemEnv).map_err(OpenError::Rendezvous)?,
};
let rv = Rendezvous::new(rd, domain, name);
let request = HelloRequest {
format_version: tf_tree_arena::FORMAT_VERSION,
layout_hash: tf_tree_arena::layout_hash(),
mode: match self.mode {
AttachMode::ReadOnly => AccessMode::ReadOnly,
AttachMode::ReadWrite => AccessMode::ReadWrite,
},
client_pid: std::process::id(),
client_start_time: self_start_time().unwrap_or(0),
client_boot_id: boot_id().unwrap_or([0; 16]),
client_name: name_bytes(),
};
let mut probe = SocketProbe::new(request, per_attempt);
let mut session = tf_tree_ipc::Open::new(rv.clone())
.mode(request.mode)
.create(self.create)
.timeout(per_attempt)
.open(&mut probe)
.map_err(OpenError::Rendezvous)?;
match session.outcome() {
OpenOutcome::Joined => {
if self.require_create {
drop(session);
return Err(OpenError::ArenaAlreadyLive);
}
let attached = session
.take_attached()
.ok_or(OpenError::Rendezvous(IpcError::ArenaAbsent))?;
let slot = attached.response.participant_slot;
let mut tree = Tree::attach_shared_at(attached.segment, self.mode, slot)?;
tree.use_ofd_liveness(LivenessProbe::open(&rv)?);
tree.use_claim_leases(open_claim_lock(&rv)?);
tree.hold_attachment(session, attached.socket);
Ok(tree)
}
OpenOutcome::Created | OpenOutcome::TookOver => {
let builder = self.layout.clone().ok_or(OpenError::NoLayoutToCreate)?;
let mut tree = builder.build_shared(rv.name().as_str())?;
tree.use_ofd_liveness(LivenessProbe::open(&rv)?);
tree.use_claim_leases(open_claim_lock(&rv)?);
let server = spawn_owner_server(&rv, &tree)?;
tree.hold_ownership(session, server);
Ok(tree)
}
}
}
}
fn is_retryable(err: OpenError) -> bool {
matches!(
err,
OpenError::Rendezvous(IpcError::ArenaAbsent)
| OpenError::Rendezvous(IpcError::ArenaHeldButUnreachable { .. })
)
}
pub(crate) struct OwnerThread {
shutdown: ShutdownHandle,
join: Option<std::thread::JoinHandle<()>>,
fork_gen: u64,
running: Arc<AtomicBool>,
}
impl OwnerThread {
pub(crate) fn stop(&mut self) {
if self.fork_gen != tf_tree_ipc::fork::generation() {
self.join = None;
return;
}
let _ = self.shutdown.stop();
if let Some(h) = self.join.take() {
let _ = h.join();
}
self.running.store(false, Ordering::Release);
}
}
impl Drop for OwnerThread {
fn drop(&mut self) {
self.stop();
}
}
fn spawn_owner_server(rv: &Rendezvous, tree: &Tree) -> Result<OwnerThread, OpenError> {
let lock_probe = LivenessProbe::open(rv)?;
let view = tree.view();
let header = view.header();
let desc = SegmentDescriptor {
format_version: header.format_version,
layout_hash: header.layout_hash,
arena_size: header.arena_size,
instance_uuid: header.instance_uuid,
boot_id: header.boot_id,
};
let server = OwnerServer::bind_at(rv.sock_path(), desc, std::process::id())
.map_err(OpenError::Rendezvous)?;
let shutdown = server.shutdown_handle().map_err(OpenError::Rendezvous)?;
let segment = tree
.shared_fd()
.ok_or(OpenError::Rendezvous(IpcError::ArenaAbsent))?;
let segment = rustix_dup(segment).map_err(OpenError::Rendezvous)?;
let table_fd = {
use std::os::fd::AsFd;
rustix_dup(segment.as_fd()).map_err(OpenError::Rendezvous)?
};
let table_arena = tf_tree_arena::MappedArena::attach(table_fd, AttachMode::ReadWrite)?;
let running = Arc::new(AtomicBool::new(true));
let flag = Arc::clone(&running);
let join = std::thread::Builder::new()
.name("tf_tree-owner".into())
.spawn(move || {
use std::os::fd::AsFd;
let granted = std::rc::Rc::new(std::cell::Cell::new(0u64));
let granted_assign = std::rc::Rc::clone(&granted);
let granted_hangup = std::rc::Rc::clone(&granted);
let _ = server.serve(
segment.as_fd(),
|_req| {
let view = tf_tree_core::arena_view::ArenaView::new(&table_arena);
let table = view.participants();
let n = table.capacity().min(64) as u32;
for slot in 0..n {
let bit = 1u64 << slot;
if granted_assign.get() & bit != 0 {
continue; }
if table.identity(slot).is_some() {
continue; }
if lock_probe.is_held(slot).unwrap_or(false) {
continue;
}
granted_assign.set(granted_assign.get() | bit);
return Ok(slot);
}
Err(HelloStatus::NoParticipantSlots)
},
|slot| {
let view = tf_tree_core::arena_view::ArenaView::new(&table_arena);
let table = view.participants();
if let Some((_pid, _start, incarnation)) = table.identity(slot) {
table.release(slot, incarnation);
}
granted_hangup.set(granted_hangup.get() & !(1u64 << slot));
},
);
flag.store(false, Ordering::Release);
})
.map_err(|_| OpenError::Rendezvous(IpcError::ArenaAbsent))?;
Ok(OwnerThread {
shutdown,
join: Some(join),
running,
fork_gen: tf_tree_ipc::fork::generation(),
})
}
fn rustix_dup(fd: std::os::fd::BorrowedFd<'_>) -> Result<std::os::fd::OwnedFd, IpcError> {
fd.try_clone_to_owned()
.map_err(|e| IpcError::ClientSocketSetup {
raw_os_error: e.raw_os_error().unwrap_or(0),
})
}
fn name_bytes() -> [u8; 32] {
tf_tree_ipc::self_comm()
}