#![cfg(not(target_arch = "wasm32"))]
use std::collections::HashMap;
use std::fmt;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::{watch, Mutex, OnceCell, RwLock};
use crate::client::auth_readiness::{AuthLeg, AuthReadinessStore, AuthReadinessWaitError};
use crate::client::correlation::CorrelationContext;
use crate::clog;
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct DriveGrantConnectionKey {
pub scope: String,
pub peer_node_id: String,
}
impl DriveGrantConnectionKey {
pub fn new(scope: impl Into<String>, peer_node_id: impl Into<String>) -> Self {
Self {
scope: scope.into(),
peer_node_id: peer_node_id.into(),
}
}
}
impl fmt::Display for DriveGrantConnectionKey {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}@{}", self.scope, self.peer_node_id)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum WebRtcState {
Unknown,
Disabled,
Connecting,
Connected,
Failed(String),
}
#[derive(Debug, Clone)]
pub enum ActorError {
Shutdown,
DialFailed(String),
NotReady,
LogicalChannelNotImplemented(String),
AuthNotReady(String),
}
impl fmt::Display for ActorError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
ActorError::Shutdown => write!(f, "drive-grant-actor: already shut down"),
ActorError::DialFailed(reason) => {
write!(f, "drive-grant-actor: dial failed: {reason}")
}
ActorError::NotReady => write!(
f,
"drive-grant-actor: ensure_ready() must resolve before open_logical_channel"
),
ActorError::LogicalChannelNotImplemented(label) => write!(
f,
"drive-grant-actor: logical channel '{label}' wire format pending Phase 6"
),
ActorError::AuthNotReady(reason) => write!(
f,
"drive-grant-actor: auth-readiness gate not ready before first dial: {reason}"
),
}
}
}
impl std::error::Error for ActorError {}
#[async_trait::async_trait]
pub trait DriveGrantTransport: Send + Sync + 'static {
async fn dial(&self, key: &DriveGrantConnectionKey) -> Result<(), String>;
async fn close(&self, _key: &DriveGrantConnectionKey) {}
}
#[derive(Debug, Clone)]
pub struct DriveGrantActorConfig {
pub initial_retry: Duration,
pub max_retry: Duration,
pub max_consecutive_failures: u32,
pub auth_readiness_wait_timeout: Option<Duration>,
}
impl Default for DriveGrantActorConfig {
fn default() -> Self {
Self {
initial_retry: Duration::from_millis(250),
max_retry: Duration::from_secs(4),
max_consecutive_failures: 8,
auth_readiness_wait_timeout: Some(Duration::from_secs(30)),
}
}
}
#[derive(Debug, Clone)]
pub struct LogicalChannelHandle {
label: String,
}
impl LogicalChannelHandle {
pub fn label(&self) -> &str {
&self.label
}
}
#[derive(Debug)]
struct DialState {
consecutive_failures: u32,
ever_succeeded: bool,
shutdown_reason: Option<String>,
}
impl DialState {
fn new() -> Self {
Self {
consecutive_failures: 0,
ever_succeeded: false,
shutdown_reason: None,
}
}
}
pub struct DriveGrantConnectionActor {
key: DriveGrantConnectionKey,
transport: Arc<dyn DriveGrantTransport>,
config: DriveGrantActorConfig,
auth_readiness: Option<Arc<AuthReadinessStore>>,
auth_readiness_legs: Vec<AuthLeg>,
select_loop: Mutex<()>,
state: RwLock<DialState>,
ready_once: OnceCell<()>,
webrtc_tx: watch::Sender<WebRtcState>,
}
impl DriveGrantConnectionActor {
fn new(
key: DriveGrantConnectionKey,
transport: Arc<dyn DriveGrantTransport>,
config: DriveGrantActorConfig,
auth_readiness: Option<Arc<AuthReadinessStore>>,
auth_readiness_legs: Vec<AuthLeg>,
) -> Arc<Self> {
let (webrtc_tx, _rx) = watch::channel(WebRtcState::Unknown);
Arc::new(Self {
key,
transport,
config,
auth_readiness,
auth_readiness_legs,
select_loop: Mutex::new(()),
state: RwLock::new(DialState::new()),
ready_once: OnceCell::new(),
webrtc_tx,
})
}
pub fn key(&self) -> &DriveGrantConnectionKey {
&self.key
}
fn correlation(&self) -> CorrelationContext {
let mut ctx = CorrelationContext::from_drive_grant_scope(&self.key.scope)
.session_kind("drive-grant-guest")
.peer_node_id(&self.key.peer_node_id);
if ctx.scope.is_none() {
ctx = ctx.scope(&self.key.scope);
}
ctx
}
pub fn webrtc_state(&self) -> WebRtcState {
self.webrtc_tx.borrow().clone()
}
pub fn observe_webrtc(&self) -> watch::Receiver<WebRtcState> {
self.webrtc_tx.subscribe()
}
pub fn set_webrtc_state(&self, state: WebRtcState) {
let label = match &state {
WebRtcState::Unknown => "unknown",
WebRtcState::Disabled => "disabled",
WebRtcState::Connecting => "connecting",
WebRtcState::Connected => "connected",
WebRtcState::Failed(_) => "failed",
};
clog!(
"[drive-grant-actor][webrtc]",
&self.correlation(),
"state={}",
label
);
let _ = self.webrtc_tx.send(state);
}
pub async fn ensure_ready(self: &Arc<Self>) -> Result<(), ActorError> {
if self.state.read().await.shutdown_reason.is_some() {
return Err(ActorError::Shutdown);
}
if self.ready_once.initialized() {
return Ok(());
}
let attempt_dial = || async {
if let Some(store) = &self.auth_readiness {
match store
.wait_until_ready(
&self.auth_readiness_legs,
self.config.auth_readiness_wait_timeout,
)
.await
{
Ok(()) => {}
Err(AuthReadinessWaitError::Timeout { remaining, elapsed }) => {
return Err(ActorError::AuthNotReady(format!(
"remaining={:?} elapsed={:?}",
remaining, elapsed
)));
}
Err(AuthReadinessWaitError::Closed) => {
return Err(ActorError::AuthNotReady("store-closed".into()));
}
}
}
let _guard = self.select_loop.lock().await;
self.dial_with_retries().await
};
match self.ready_once.get_or_try_init(attempt_dial).await {
Ok(()) => Ok(()),
Err(error) => Err(error),
}
}
async fn dial_with_retries(&self) -> Result<(), ActorError> {
let mut wait = self.config.initial_retry;
loop {
if self.state.read().await.shutdown_reason.is_some() {
return Err(ActorError::Shutdown);
}
match self.transport.dial(&self.key).await {
Ok(()) => {
let mut state = self.state.write().await;
let was_first = !state.ever_succeeded;
state.consecutive_failures = 0;
state.ever_succeeded = true;
drop(state);
if was_first {
clog!(
"[drive-grant-actor][dial]",
&self.correlation(),
"ready_first"
);
} else {
clog!(
"[drive-grant-actor][dial]",
&self.correlation(),
"ready_recovered"
);
}
return Ok(());
}
Err(error) => {
let consecutive_failures = {
let mut state = self.state.write().await;
state.consecutive_failures = state.consecutive_failures.saturating_add(1);
state.consecutive_failures
};
if consecutive_failures >= self.config.max_consecutive_failures {
clog!(
"[drive-grant-actor][dial]",
&self.correlation(),
"failed_terminal attempts={} error={}",
consecutive_failures,
error
);
return Err(ActorError::DialFailed(error));
}
clog!(
"[drive-grant-actor][dial]",
&self.correlation(),
"failed_retry attempt={} error={}",
consecutive_failures,
error
);
tokio::time::sleep(wait).await;
wait = (wait * 2).min(self.config.max_retry);
}
}
}
}
pub async fn open_logical_channel(
self: &Arc<Self>,
label: &str,
) -> Result<LogicalChannelHandle, ActorError> {
if self.state.read().await.shutdown_reason.is_some() {
return Err(ActorError::Shutdown);
}
if !self.ready_once.initialized() {
return Err(ActorError::NotReady);
}
match label {
"drive-view.list" | "drive-view.fetch" | "drive-view.watch" | "drive-view.buckets"
| "drive-view" => {
let ctx = self.correlation().channel(label);
clog!("[drive-grant-actor][channel]", &ctx, "open");
Ok(LogicalChannelHandle {
label: label.to_string(),
})
}
_ => {
let ctx = self.correlation().channel(label);
clog!(
"[drive-grant-actor][channel]",
&ctx,
"rejected_unknown_label"
);
Err(ActorError::LogicalChannelNotImplemented(label.to_string()))
}
}
}
#[doc(hidden)]
pub async fn open_logical_channel_test_stub(
self: &Arc<Self>,
label: &str,
) -> Result<LogicalChannelHandle, ActorError> {
if self.state.read().await.shutdown_reason.is_some() {
return Err(ActorError::Shutdown);
}
if !self.ready_once.initialized() {
return Err(ActorError::NotReady);
}
Ok(LogicalChannelHandle {
label: label.to_string(),
})
}
pub async fn shutdown(self: &Arc<Self>) {
{
let mut state = self.state.write().await;
if state.shutdown_reason.is_some() {
return;
}
state.shutdown_reason = Some("explicit-shutdown".to_string());
}
clog!(
"[drive-grant-actor][shutdown]",
&self.correlation(),
"begin"
);
let _ = self.webrtc_tx.send(WebRtcState::Failed(
"drive-grant-actor: shutdown".to_string(),
));
self.transport.close(&self.key).await;
clog!(
"[drive-grant-actor][shutdown]",
&self.correlation(),
"complete"
);
}
#[doc(hidden)]
pub async fn debug_state(&self) -> DriveGrantActorDebugState {
let state = self.state.read().await;
DriveGrantActorDebugState {
consecutive_failures: state.consecutive_failures,
ever_succeeded: state.ever_succeeded,
ready: self.ready_once.initialized(),
shutdown: state.shutdown_reason.is_some(),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[doc(hidden)]
pub struct DriveGrantActorDebugState {
pub consecutive_failures: u32,
pub ever_succeeded: bool,
pub ready: bool,
pub shutdown: bool,
}
pub struct DriveGrantActorRegistry {
actors: RwLock<HashMap<DriveGrantConnectionKey, Arc<DriveGrantConnectionActor>>>,
transport: Arc<dyn DriveGrantTransport>,
config: DriveGrantActorConfig,
auth_readiness: Option<Arc<AuthReadinessStore>>,
auth_readiness_legs: Vec<AuthLeg>,
}
impl DriveGrantActorRegistry {
pub fn new(transport: Arc<dyn DriveGrantTransport>) -> Arc<Self> {
Arc::new(Self {
actors: RwLock::new(HashMap::new()),
transport,
config: DriveGrantActorConfig::default(),
auth_readiness: None,
auth_readiness_legs: vec![
AuthLeg::FilesAuth,
AuthLeg::PlutoRtcAuth,
AuthLeg::RuntimeAuth,
AuthLeg::Firestore,
],
})
}
pub fn with_config(
transport: Arc<dyn DriveGrantTransport>,
config: DriveGrantActorConfig,
) -> Arc<Self> {
Arc::new(Self {
actors: RwLock::new(HashMap::new()),
transport,
config,
auth_readiness: None,
auth_readiness_legs: vec![
AuthLeg::FilesAuth,
AuthLeg::PlutoRtcAuth,
AuthLeg::RuntimeAuth,
AuthLeg::Firestore,
],
})
}
pub fn with_auth_readiness(
transport: Arc<dyn DriveGrantTransport>,
config: DriveGrantActorConfig,
auth_readiness: Arc<AuthReadinessStore>,
legs: Vec<AuthLeg>,
) -> Arc<Self> {
Arc::new(Self {
actors: RwLock::new(HashMap::new()),
transport,
config,
auth_readiness: Some(auth_readiness),
auth_readiness_legs: if legs.is_empty() {
vec![
AuthLeg::FilesAuth,
AuthLeg::PlutoRtcAuth,
AuthLeg::RuntimeAuth,
AuthLeg::Firestore,
]
} else {
legs
},
})
}
pub async fn get_or_spawn(
&self,
key: DriveGrantConnectionKey,
) -> Arc<DriveGrantConnectionActor> {
if let Some(actor) = self.actors.read().await.get(&key).cloned() {
return actor;
}
let mut actors = self.actors.write().await;
if let Some(actor) = actors.get(&key).cloned() {
return actor;
}
let actor = DriveGrantConnectionActor::new(
key.clone(),
self.transport.clone(),
self.config.clone(),
self.auth_readiness.clone(),
self.auth_readiness_legs.clone(),
);
actors.insert(key, actor.clone());
actor
}
pub async fn get(
&self,
key: &DriveGrantConnectionKey,
) -> Option<Arc<DriveGrantConnectionActor>> {
self.actors.read().await.get(key).cloned()
}
pub async fn shutdown_all(&self) {
let actors: Vec<_> = {
let mut actors = self.actors.write().await;
actors.drain().map(|(_, actor)| actor).collect()
};
for actor in actors {
actor.shutdown().await;
}
}
#[doc(hidden)]
pub async fn len(&self) -> usize {
self.actors.read().await.len()
}
}