use crate::{NFSVersion, NfsError, OperationOutcome, RecoveryAction, Result};
use async_trait::async_trait;
use std::collections::VecDeque;
use std::fmt;
use std::future::Future;
use std::sync::atomic::{AtomicU8, AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use tokio::sync::Notify;
const READY: u8 = 0;
const CLOSING: u8 = 1;
const CLOSED: u8 = 2;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ClientLifecycle {
Ready,
Closing,
Closed,
}
#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct ResourceKey(u64);
impl fmt::Display for ResourceKey {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
self.0.fmt(formatter)
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct CoreOperation {
pub name: String,
pub safe_path: Option<String>,
}
#[async_trait]
pub trait ClientDriver: fmt::Debug + Send + Sync + 'static {
async fn execute(&self, operation: CoreOperation) -> Result<()>;
async fn close_resource(&self, key: ResourceKey) -> Result<()>;
async fn umount(&self) -> Result<()>;
}
#[derive(Debug, Default)]
pub struct ClientCloseReport {
errors: Vec<Arc<NfsError>>,
}
impl ClientCloseReport {
pub fn errors(&self) -> &[Arc<NfsError>] {
&self.errors
}
}
#[derive(Debug)]
pub struct ClientCore {
driver: Arc<dyn ClientDriver>,
lifecycle: AtomicU8,
next_resource_key: AtomicU64,
in_flight: DrainCounter,
resources: Mutex<Vec<ResourceKey>>,
owned_tasks: DrainCounter,
recovery_events: Mutex<RecoveryEventQueue>,
close_state: Mutex<CloseState>,
close_notify: Notify,
lifecycle_notify: Notify,
}
#[derive(Debug, Default)]
struct CloseState {
started: bool,
report: Option<Arc<ClientCloseReport>>,
}
#[derive(Debug, Default)]
struct DrainCounter {
count: AtomicU64,
notify: Notify,
}
impl DrainCounter {
fn increment(&self) {
self.count.fetch_add(1, Ordering::AcqRel);
}
fn decrement(&self) {
if self.count.fetch_sub(1, Ordering::AcqRel) == 1 {
self.notify.notify_waiters();
}
}
fn count(&self) -> u64 {
self.count.load(Ordering::Acquire)
}
async fn wait_for_zero(&self) {
while self.count() != 0 {
let notified = self.notify.notified();
if self.count() != 0 {
notified.await;
}
}
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct CoreRecoveryEvent {
pub operation: String,
pub safe_path: Option<String>,
pub protocol: NFSVersion,
pub outcome: OperationOutcome,
pub recovery: RecoveryAction,
pub completed_bytes: Option<u64>,
pub message: String,
}
#[derive(Debug)]
struct RecoveryEventQueue {
capacity: usize,
dropped: u64,
events: VecDeque<CoreRecoveryEvent>,
}
impl ClientCore {
pub fn new(driver: Arc<dyn ClientDriver>) -> Arc<Self> {
Self::build(driver, 256)
}
pub fn with_recovery_event_capacity(
driver: Arc<dyn ClientDriver>,
recovery_event_capacity: usize,
) -> Result<Arc<Self>> {
if recovery_event_capacity == 0 {
return Err(NfsError::InvalidInput(
"recovery-event capacity must be positive".to_string(),
));
}
Ok(Self::build(driver, recovery_event_capacity))
}
fn build(driver: Arc<dyn ClientDriver>, recovery_event_capacity: usize) -> Arc<Self> {
Arc::new(Self {
driver,
lifecycle: AtomicU8::new(READY),
next_resource_key: AtomicU64::new(1),
in_flight: DrainCounter::default(),
resources: Mutex::new(Vec::new()),
owned_tasks: DrainCounter::default(),
recovery_events: Mutex::new(RecoveryEventQueue {
capacity: recovery_event_capacity,
dropped: 0,
events: VecDeque::new(),
}),
close_state: Mutex::new(CloseState::default()),
close_notify: Notify::new(),
lifecycle_notify: Notify::new(),
})
}
fn ensure_ready(&self) -> Result<()> {
if self.lifecycle.load(Ordering::Acquire) == READY {
Ok(())
} else {
Err(NfsError::ClientClosed(
"connected client is closing or closed".to_string(),
))
}
}
pub fn lifecycle(&self) -> ClientLifecycle {
match self.lifecycle.load(Ordering::Acquire) {
READY => ClientLifecycle::Ready,
CLOSING => ClientLifecycle::Closing,
_ => ClientLifecycle::Closed,
}
}
pub async fn wait_for_lifecycle(&self, expected: ClientLifecycle) {
while self.lifecycle() != expected {
let notified = self.lifecycle_notify.notified();
if self.lifecycle() != expected {
notified.await;
}
}
}
pub async fn execute(self: &Arc<Self>, operation: CoreOperation) -> Result<()> {
let _operation = self.begin_operation()?;
self.driver.execute(operation).await
}
pub fn record_recovery_event(&self, event: CoreRecoveryEvent) -> Result<()> {
let mut queue = self
.recovery_events
.lock()
.map_err(|_| NfsError::Rpc("recovery-event queue lock poisoned".to_string()))?;
if queue.events.len() == queue.capacity {
queue.events.pop_front();
queue.dropped = queue.dropped.saturating_add(1);
}
queue.events.push_back(event);
Ok(())
}
pub fn recovery_events(&self) -> Result<Vec<CoreRecoveryEvent>> {
self.recovery_events
.lock()
.map(|queue| queue.events.iter().cloned().collect())
.map_err(|_| NfsError::Rpc("recovery-event queue lock poisoned".to_string()))
}
pub fn drain_recovery_events(&self) -> Result<Vec<CoreRecoveryEvent>> {
self.recovery_events
.lock()
.map(|mut queue| queue.events.drain(..).collect())
.map_err(|_| NfsError::Rpc("recovery-event queue lock poisoned".to_string()))
}
pub fn dropped_recovery_event_count(&self) -> Result<u64> {
self.recovery_events
.lock()
.map(|queue| queue.dropped)
.map_err(|_| NfsError::Rpc("recovery-event queue lock poisoned".to_string()))
}
pub fn register_resource(&self) -> Result<ResourceKey> {
let key = self.allocate_resource_key()?;
self.publish_resource(key)?;
Ok(key)
}
pub fn allocate_resource_key(&self) -> Result<ResourceKey> {
self.ensure_ready()?;
Ok(ResourceKey(
self.next_resource_key.fetch_add(1, Ordering::Relaxed),
))
}
pub fn publish_resource(&self, key: ResourceKey) -> Result<()> {
let mut resources = self
.resources
.lock()
.map_err(|_| NfsError::Rpc("client resource registry lock poisoned".to_string()))?;
self.ensure_ready()?;
resources.push(key);
Ok(())
}
pub fn unregister_resource(&self, key: ResourceKey) -> Result<bool> {
let mut resources = self
.resources
.lock()
.map_err(|_| NfsError::Rpc("client resource registry lock poisoned".to_string()))?;
let Some(position) = resources.iter().position(|candidate| *candidate == key) else {
return Ok(false);
};
resources.remove(position);
Ok(true)
}
pub fn resource_count(&self) -> usize {
self.resources
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.len()
}
pub fn begin_operation(self: &Arc<Self>) -> Result<OperationGuard> {
self.ensure_ready()?;
self.in_flight.increment();
if let Err(error) = self.ensure_ready() {
self.finish_operation();
return Err(error);
}
Ok(OperationGuard {
core: Some(Arc::clone(self)),
})
}
fn finish_operation(&self) {
self.in_flight.decrement();
}
pub fn spawn_owned<F>(self: &Arc<Self>, future: F) -> Result<()>
where
F: Future<Output = ()> + Send + 'static,
{
self.ensure_ready()?;
self.owned_tasks.increment();
if let Err(error) = self.ensure_ready() {
self.finish_owned_task();
return Err(error);
}
let core = Arc::clone(self);
tokio::spawn(async move {
let _guard = OwnedTaskGuard { core };
future.await;
});
Ok(())
}
fn finish_owned_task(&self) {
self.owned_tasks.decrement();
}
pub fn owned_task_count(&self) -> u64 {
self.owned_tasks.count()
}
pub async fn close(self: &Arc<Self>) -> Arc<ClientCloseReport> {
let start_cleanup = {
let mut state = self
.close_state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if state.started {
false
} else {
state.started = true;
true
}
};
if start_cleanup {
self.lifecycle.store(CLOSING, Ordering::Release);
self.lifecycle_notify.notify_waiters();
let core = Arc::clone(self);
tokio::spawn(async move {
core.run_close().await;
});
}
loop {
let notified = self.close_notify.notified();
if let Some(report) = self
.close_state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.report
.clone()
{
return report;
}
notified.await;
}
}
async fn run_close(&self) {
self.in_flight.wait_for_zero().await;
self.owned_tasks.wait_for_zero().await;
let mut errors = Vec::new();
if self.close_state.is_poisoned() {
errors.push(Arc::new(NfsError::Rpc(
"client close-state lock poisoned".to_string(),
)));
}
if self.resources.is_poisoned() {
errors.push(Arc::new(NfsError::Rpc(
"client resource registry lock poisoned".to_string(),
)));
}
let resources = {
let mut resources = self
.resources
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
std::mem::take(&mut *resources)
};
for key in resources {
if let Err(error) = self.driver.close_resource(key).await {
errors.push(Arc::new(error));
}
}
if let Err(error) = self.driver.umount().await {
errors.push(Arc::new(error));
}
self.lifecycle.store(CLOSED, Ordering::Release);
self.lifecycle_notify.notify_waiters();
let report = Arc::new(ClientCloseReport { errors });
self.close_state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.report = Some(report);
self.close_notify.notify_waiters();
}
}
#[derive(Debug)]
pub struct OperationGuard {
core: Option<Arc<ClientCore>>,
}
impl Drop for OperationGuard {
fn drop(&mut self) {
if let Some(core) = self.core.take() {
core.finish_operation();
}
}
}
struct OwnedTaskGuard {
core: Arc<ClientCore>,
}
impl Drop for OwnedTaskGuard {
fn drop(&mut self) {
self.core.finish_owned_task();
}
}
#[cfg(test)]
mod poison_tests {
use super::*;
use std::panic::{AssertUnwindSafe, catch_unwind};
#[derive(Debug, Default)]
struct Driver {
closed: Mutex<Vec<ResourceKey>>,
}
#[async_trait]
impl ClientDriver for Driver {
async fn execute(&self, _operation: CoreOperation) -> Result<()> {
Ok(())
}
async fn close_resource(&self, key: ResourceKey) -> Result<()> {
self.closed.lock().unwrap().push(key);
Ok(())
}
async fn umount(&self) -> Result<()> {
Ok(())
}
}
#[tokio::test]
async fn poisoned_resource_registry_is_reported_without_skipping_cleanup() {
let driver = Arc::new(Driver::default());
let core = ClientCore::new(driver.clone());
let key = core.register_resource().unwrap();
let _ = catch_unwind(AssertUnwindSafe(|| {
let _guard = core.resources.lock().unwrap();
panic!("poison resource registry");
}));
let report = tokio::time::timeout(std::time::Duration::from_secs(1), core.close())
.await
.expect("poisoned close must terminate");
assert!(
report.errors()[0]
.to_string()
.contains("registry lock poisoned")
);
assert_eq!(*driver.closed.lock().unwrap(), vec![key]);
}
#[tokio::test]
async fn poisoned_close_state_is_reported_and_publishes_terminal_report() {
let core = ClientCore::new(Arc::new(Driver::default()));
let _ = catch_unwind(AssertUnwindSafe(|| {
let _guard = core.close_state.lock().unwrap();
panic!("poison close state");
}));
let report = tokio::time::timeout(std::time::Duration::from_secs(1), core.close())
.await
.expect("poisoned close must terminate");
assert!(
report.errors()[0]
.to_string()
.contains("close-state lock poisoned")
);
assert_eq!(core.lifecycle(), ClientLifecycle::Closed);
}
}