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,
#[cfg(feature = "test-hooks")]
probes: std::sync::atomic::AtomicU32,
}
impl LivenessProbe {
fn from_lock(lock: tf_tree_ipc::LockFile) -> LivenessProbe {
LivenessProbe {
lock,
#[cfg(feature = "test-hooks")]
probes: std::sync::atomic::AtomicU32::new(0),
}
}
fn open(rv: &Rendezvous) -> Result<LivenessProbe, IpcError> {
Ok(LivenessProbe::from_lock(tf_tree_ipc::LockFile::open(
rv.lock_path(),
)?))
}
pub(crate) fn lock(&self) -> &tf_tree_ipc::LockFile {
&self.lock
}
pub(crate) fn is_held(&self, slot: u32) -> Option<bool> {
#[cfg(feature = "test-hooks")]
self.probes.fetch_add(1, Ordering::Relaxed);
self.lock.probe_participant(slot).ok().map(|p| p.held)
}
#[cfg(feature = "test-hooks")]
fn probe_count(&self) -> u32 {
self.probes.load(Ordering::Relaxed)
}
}
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub(crate) enum Reclamation {
Reclaimable { observed: u32 },
Live,
Unknown,
}
pub(crate) fn reclamation_verdict(
probe: &LivenessProbe,
own_slot: u32,
slot: u32,
rec: &tf_tree_core::ParticipantRecord,
) -> Reclamation {
if slot == own_slot {
return Reclamation::Live;
}
let observed = rec.state.load(Ordering::Acquire);
if tf_tree_core::participant::state_of(observed) == tf_tree_core::participant::FREE {
return Reclamation::Unknown;
}
match probe.is_held(slot) {
Some(true) => Reclamation::Live,
Some(false) => Reclamation::Reclaimable { observed },
None => Reclamation::Unknown,
}
}
#[cfg(feature = "test-hooks")]
#[doc(hidden)]
#[must_use]
pub fn reclamation_verdict_for_test(
tree: &Tree,
lock_path: &std::path::Path,
own_slot: u32,
slot: u32,
) -> String {
let Ok(lock) = tf_tree_ipc::LockFile::open(lock_path) else {
return "no-lock-file".to_string();
};
let probe = LivenessProbe::from_lock(lock);
let view = tree.view();
let Some(rec) = view.participants().get(slot) else {
return "no-such-slot".to_string();
};
let verdict = match reclamation_verdict(&probe, own_slot, slot, rec) {
Reclamation::Reclaimable { observed } => format!("reclaimable word {observed:#x}"),
Reclamation::Live => "live".to_string(),
Reclamation::Unknown => "unknown".to_string(),
};
format!("{verdict} probes={}", probe.probe_count())
}
pub(crate) type JoinedSession = tf_tree_ipc::Session<tf_tree_ipc::Attached>;
pub(crate) enum Attachment {
Joined {
session: JoinedSession,
socket: std::os::fd::OwnedFd,
rendezvous: Rendezvous,
},
Owner {
_server: OwnerThread,
_session: JoinedSession,
},
}
#[cfg(feature = "crash-points")]
pub const CRASH_SITES: &[&str] = &[
"takeover.after_ownership_lock_before_bind",
"topo.holding_lock",
"open.after_ownership_lock_before_bind",
"open.after_create_before_bind",
"reclaim.after_probe_before_cas",
"hangup.after_probe_before_cas",
];
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub enum Inheritance {
Inherited,
OwnerAlive,
Contended,
ReadOnly,
NotApplicable,
}
impl Tree {
#[cfg(all(feature = "shm", target_os = "linux"))]
pub fn inherit_ownership(&self) -> Result<Inheritance, OpenError> {
if !self.is_joined() {
return Ok(Inheritance::NotApplicable);
}
if !self.owner_lost() {
return Ok(Inheritance::OwnerAlive);
}
if !self.is_writable() {
return Ok(Inheritance::ReadOnly);
}
let (mut session, socket, rendezvous) = match self.take_attachment() {
Some(Attachment::Joined {
session,
socket,
rendezvous,
}) => (session, socket, rendezvous),
other => {
self.put_attachment(other);
return Ok(Inheritance::NotApplicable);
}
};
let acquired = match session.take_over_ownership() {
Ok(v) => v,
Err(e) => {
self.put_attachment(Some(Attachment::Joined {
session,
socket,
rendezvous,
}));
return Err(OpenError::Rendezvous(e));
}
};
if !acquired {
self.put_attachment(Some(Attachment::Joined {
session,
socket,
rendezvous,
}));
return Ok(Inheritance::Contended);
}
#[cfg(feature = "crash-points")]
tf_tree_core::crash::maybe_abort(CRASH_SITES[0]);
match spawn_owner_server(&rendezvous, self) {
Ok(server) => {
drop(socket);
self.put_attachment(Some(Attachment::Owner {
_server: server,
_session: session,
}));
Ok(Inheritance::Inherited)
}
Err(e) => {
let _ = session.release_ownership();
self.put_attachment(Some(Attachment::Joined {
session,
socket,
rendezvous,
}));
Err(e)
}
}
}
}
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,
#[error("this process's participant lock byte and its arena participant record are different slots; the arena was not published")]
ParticipantSlotDiverged,
}
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 per_attempt =
Duration::new(per_attempt.as_secs(), per_attempt.subsec_micros() * 1_000);
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 ipc_open = tf_tree_ipc::Open::new(rv.clone())
.mode(request.mode)
.create(self.create)
.timeout(per_attempt);
let mut session = ipc_open.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_joined_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, rv);
Ok(tree)
}
OpenOutcome::Created => {
#[cfg(feature = "crash-points")]
tf_tree_core::crash::maybe_abort(CRASH_SITES[2]);
let builder = self.layout.clone().ok_or(OpenError::NoLayoutToCreate)?;
let mut tree = builder.build_shared(rv.name().as_str())?;
#[cfg(feature = "crash-points")]
tf_tree_core::crash::maybe_abort(CRASH_SITES[3]);
tree.use_ofd_liveness(LivenessProbe::open(&rv)?);
tree.use_claim_leases(open_claim_lock(&rv)?);
if session.slot() != tree.participant_slot() {
drop(tree);
drop(session);
return Err(OpenError::ParticipantSlotDiverged);
}
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 own_slot = tree.participant_slot();
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; }
let Some(rec) = table.get(slot) else {
continue;
};
let word = rec.state.load(Ordering::Acquire);
if tf_tree_core::participant::state_of(word)
!= tf_tree_core::participant::FREE
{
match reclamation_verdict(&lock_probe, own_slot, slot, rec) {
Reclamation::Reclaimable { observed } => {
if !table.reclaim(slot, observed) {
continue;
}
}
Reclamation::Live => continue,
Reclamation::Unknown => 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 revoked =
crate::tree::reap_claims(&view, lock_probe.lock(), Some(slot), own_slot);
let _ = revoked;
let table = view.participants();
if let Some(rec) = table.get(slot) {
let observed = rec.state.load(Ordering::Acquire);
if tf_tree_core::participant::state_of(observed)
!= tf_tree_core::participant::FREE
{
#[cfg(feature = "crash-points")]
tf_tree_core::crash::maybe_abort(CRASH_SITES[5]);
let _ = table.reclaim(slot, observed);
}
}
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] {
let mut out = [0u8; 32];
let comm = tf_tree_ipc::self_comm();
out[..comm.len()].copy_from_slice(&comm);
out
}