use std::{
sync::{Arc, Mutex},
time::{Duration, Instant},
};
use anyhow::{Context, Result, ensure};
use async_trait::async_trait;
use google_cloud_auth::credentials::{CacheableResource, CredentialsProvider, EntityTag};
use tokio_util::sync::CancellationToken;
use super::{
HostId,
actor_runtime::{ActorActivation, ActorStorage},
};
use crate::{
actor::ActorKey,
bucket::{
Bucket, GcsBucket, RuntimeStorage, WarmGcs,
access::{HostStorageConfig, StorageToken},
},
clock::{Clock, SystemClock},
control_plane::{ControlPlaneClient, LeaseFence},
host_leases::{HostLease, HostLeaseRegistry, HostLeaseRequest},
storage::WritePlan,
};
pub(crate) struct HostStorage {
pub runtime: Arc<RuntimeStorage>,
objects: Option<google_cloud_storage::client::Storage>,
pub(super) stop: CancellationToken,
observer: Arc<ControlPlaneClient>,
host: HostId,
session: String,
region: String,
actor: Option<ActorKey>,
new_actor: bool,
owner_hint: Option<crate::bucket::OwnershipHint>,
activation: Mutex<Option<ActorActivation>>,
fence: Mutex<LeaseFence>,
lease: Mutex<Option<HostLease>>,
renewal: tokio::sync::Mutex<()>,
}
impl HostStorage {
pub async fn new(
config: HostStorageConfig,
host: HostId,
session: String,
client: Arc<ControlPlaneClient>,
stop: CancellationToken,
warm: Option<WarmGcs>,
actor: Option<ActorKey>,
) -> Result<Self> {
let credentials = HostCredentials::new(config.token, client.clone(), stop.clone());
let (authority, clients): (Arc<dyn Bucket>, _) = match config.bucket {
crate::bucket::access::BucketLocation::Gcs { bucket, .. } => {
let bucket = match warm {
Some(warm) => warm.bind(&bucket, credentials.into())?,
None => GcsBucket::with_credentials(&bucket, credentials.into()).await?,
};
let clients = bucket.clients();
(Arc::new(bucket), Some(clients))
}
crate::bucket::access::BucketLocation::File { directory } => {
(Arc::new(crate::bucket::FileBucket::new(directory)?), None)
}
};
let mut runtime = RuntimeStorage::new(authority, Arc::new(SystemClock))?;
if config.persistence.rapid_settings().is_some() {
let clients = clients.clone().context("Rapid persistence requires GCS")?;
let snapshots = Arc::new(crate::bucket::RapidSnapshots::gcs(
&config.persistence,
clients,
stop.clone(),
actor.as_ref(),
)?);
runtime = runtime.with_persistence(config.persistence, snapshots)?;
}
Ok(Self {
objects: clients.map(|clients| clients.storage),
observer: client,
stop,
runtime: Arc::new(runtime),
host,
session,
region: config.region,
actor,
new_actor: false,
owner_hint: None,
activation: Mutex::new(None),
fence: Mutex::new(LeaseFence::default()),
lease: Mutex::new(None),
renewal: tokio::sync::Mutex::new(()),
})
}
pub(crate) async fn install_code(
&self,
artifact: &crate::artifacts::ArtifactManifest,
) -> Result<()> {
artifact
.install(
std::path::Path::new("/customer"),
self.objects
.as_ref()
.context("GCS artifact client missing")?,
)
.await
}
pub(crate) fn with_activation(
mut self,
new_actor: bool,
owner_hint: Option<crate::bucket::OwnershipHint>,
) -> Self {
self.new_actor = new_actor;
self.owner_hint = owner_hint;
self
}
pub(super) fn current_lease(&self) -> Result<HostLease> {
self.ensure_authority()?;
self.lease
.lock()
.unwrap()
.clone()
.context("activation lease missing")
}
pub(super) async fn finish(&self, deadline: tokio::time::Instant) -> Result<()> {
let started = Instant::now();
let actor = self.actor.as_ref().context("host actor identity missing")?;
let drain_deadline = deadline - Duration::from_secs(3);
let completed = tokio::time::timeout_at(
drain_deadline,
self.runtime
.drain_activation(actor, &self.host, &self.session, drain_deadline),
)
.await
.context("snapshot drain timed out")
.and_then(|result| result);
let checkpoint = match completed {
Ok(checkpoint) => Some(checkpoint),
Err(error) => {
tracing::warn!(%error, "clean session shutdown deferred to recovery");
None
}
};
let sealed = checkpoint.is_some();
let storage_ms = started.elapsed().as_secs_f64() * 1000.0;
let release_started = Instant::now();
self.fence.lock().unwrap().fenced = true;
let released = tokio::time::timeout_at(
deadline,
self.runtime
.release_with_checkpoint(actor, &self.host, &self.session, checkpoint),
)
.await
.context("ownership release timed out")
.and_then(|result| result);
if released.is_ok() {
self.notify_observer();
}
tracing::info!(
event = "host_storage_shutdown",
sealed = sealed && released.is_ok(),
released = released.is_ok(),
error = released.as_ref().err().map(|error| format!("{error:#}")),
storage_ms,
release_ms = release_started.elapsed().as_secs_f64() * 1000.0,
duration_ms = started.elapsed().as_secs_f64() * 1000.0
);
released
}
fn notify_observer(&self) {
let client = self.observer.clone();
tokio::spawn(async move {
let _ = tokio::time::timeout(Duration::from_secs(2), client.notify_inventory_changed())
.await;
});
}
fn authorize(&self, actor: &ActorKey, host: &HostId) -> Result<()> {
actor.validate()?;
ensure!(
self.actor.as_ref().is_none_or(|bound| bound == actor),
"actor identity mismatch"
);
ensure!(*host == self.host, "host storage scope mismatch");
self.ensure_authority()
}
}
#[async_trait]
impl ActorStorage for HostStorage {
fn ensure_authority(&self) -> Result<()> {
ensure!(
!self.stop.is_cancelled(),
"host session is permanently fenced"
);
self.fence
.lock()
.map_err(|_| anyhow::anyhow!("lease fence poisoned"))?
.check(Instant::now())
}
async fn acquire_actor(&self, actor: &ActorKey, host: &HostId) -> Result<ActorActivation> {
self.authorize(actor, host)?;
if let Some(activation) = self.activation.lock().unwrap().take() {
return Ok(activation);
}
let lease = self
.lease
.lock()
.unwrap()
.clone()
.context("host lease missing")?;
let loaded = self
.runtime
.activate_actor(actor, &lease, &self.region)
.await?;
self.ensure_authority()?;
Ok(ActorActivation {
owner_epoch: loaded.placement.owner_epoch,
state_version: loaded.placement.state_version,
state: loaded.state,
})
}
async fn verify_actor_ownership(
&self,
actor: &ActorKey,
host: &HostId,
epoch: u64,
) -> Result<()> {
self.authorize(actor, host)?;
let lease = self.current_lease()?;
self.runtime
.verify_actor_ownership(actor, &lease, epoch)
.await?;
self.ensure_authority()
}
async fn load_actor_state(
&self,
actor: &ActorKey,
host: &HostId,
epoch: u64,
) -> Result<(u64, bytes::Bytes)> {
self.authorize(actor, host)?;
let lease = self
.lease
.lock()
.unwrap()
.clone()
.context("host lease missing")?;
let loaded = self.runtime.load_owned_actor(actor, &lease, epoch).await?;
self.ensure_authority()?;
Ok((
loaded.placement.state_version,
loaded.state.unwrap_or_default(),
))
}
async fn read_sqlite(
&self,
actor: &ActorKey,
bytes: &[u8],
) -> Result<crate::state_log::SqliteSnapshot> {
self.authorize(actor, &self.host)?;
let snapshot = crate::state_log::StateSnapshot::decode(bytes)?;
let object = crate::storage::snapshot_object_name(
actor,
snapshot.state_version,
&format!("{:032x}", snapshot.owner_epoch),
)?;
self.runtime
.read_sqlite(&object, bytes::Bytes::copy_from_slice(bytes))
.await
}
async fn prepare_state_write(
&self,
actor: &ActorKey,
host: &HostId,
epoch: u64,
version: u64,
) -> Result<WritePlan> {
self.authorize(actor, host)?;
let lease = self
.lease
.lock()
.unwrap()
.clone()
.context("host lease missing")?;
let ticket = self
.runtime
.prepare_actor_write(
actor,
&lease,
epoch,
version.checked_add(1).context("state version overflow")?,
)
.await?;
self.ensure_authority()?;
Ok(ticket)
}
}
#[async_trait]
impl HostLeaseRegistry for HostStorage {
async fn register(&self, request: &HostLeaseRequest) -> Result<HostLease> {
self.register_with_residents(request, None).await
}
async fn register_with_residents(
&self,
request: &HostLeaseRequest,
residents: Option<&[crate::actor::ActorKey]>,
) -> Result<HostLease> {
self.register_with_inventory(request, residents, None).await
}
async fn register_with_inventory(
&self,
request: &HostLeaseRequest,
residents: Option<&[crate::actor::ActorKey]>,
queues: Option<&[crate::host_leases::ActorQueueInventory]>,
) -> Result<HostLease> {
ensure!(
request.id == self.host && request.session_id == self.session,
"host lease scope mismatch"
);
let _renewal = self.renewal.lock().await;
let started = Instant::now();
self.fence.lock().unwrap().begin(started)?;
let actor = self.actor.as_ref().context("host actor identity missing")?;
ensure!(
residents.is_none_or(
|actors| actors.len() <= 1 && actors.iter().all(|candidate| candidate == actor)
),
"resident actor scope mismatch"
);
ensure!(
queues.is_none_or(|entries| entries.iter().all(|entry| entry.actor == *actor)),
"inventory actor scope mismatch"
);
let inventory = crate::host_leases::ActivationInventory {
resident: residents.map(|actors| actors.contains(actor)),
waiting: queues.map(|entries| {
entries
.iter()
.flat_map(|entry| entry.waiting.clone())
.collect()
}),
};
let first = self.lease.lock().unwrap().is_none();
let lease = if first {
let loaded = self
.runtime
.register_activation(
actor,
request,
&self.region,
self.new_actor,
self.owner_hint.as_ref(),
)
.await?;
let lease = loaded.placement.lease;
*self.activation.lock().unwrap() = Some(ActorActivation {
owner_epoch: loaded.placement.owner_epoch,
state_version: loaded.placement.state_version,
state: loaded.state,
});
lease
} else {
self.runtime
.renew_activation(actor, request, inventory)
.await?
};
self.fence.lock().unwrap().confirm(
started,
Duration::from_millis(request.duration_ms),
Instant::now(),
)?;
self.lease.lock().unwrap().replace(lease.clone());
self.notify_observer();
Ok(lease)
}
async fn unregister(&self, host: &HostId, session: &str) -> Result<()> {
ensure!(
*host == self.host && session == self.session,
"host lease scope mismatch"
);
self.finish(tokio::time::Instant::now() + Duration::from_secs(25))
.await
}
}
#[derive(Clone)]
struct HostCredentials(Arc<std::sync::RwLock<Option<StorageToken>>>);
impl std::fmt::Debug for HostCredentials {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str("HostCredentials")
}
}
impl HostCredentials {
fn new(
token: Option<StorageToken>,
client: Arc<ControlPlaneClient>,
stop: CancellationToken,
) -> Self {
let credentials = Self(Arc::new(std::sync::RwLock::new(token)));
let refreshed = credentials.clone();
tokio::spawn(async move {
loop {
let now = SystemClock.now_ms().unwrap_or(u64::MAX);
let host_expiry = client.token_expires_at_ms().unwrap_or(now);
let expires = refreshed
.0
.read()
.unwrap()
.as_ref()
.map_or(host_expiry, |token| token.expires_at_ms.min(host_expiry));
let delay = Duration::from_millis(
expires
.saturating_sub(now.saturating_add(30_000))
.clamp(1_000, 600_000),
);
tokio::select! { _ = stop.cancelled() => return, _ = tokio::time::sleep(delay) => {} }
match client.refresh_storage_access().await {
Ok(token) => {
*refreshed.0.write().unwrap() = token;
}
Err(error) => {
tracing::warn!(%error, "could not refresh host storage credentials")
}
}
}
});
credentials
}
}
impl CredentialsProvider for HostCredentials {
async fn headers(
&self,
_: axum::http::Extensions,
) -> std::result::Result<
CacheableResource<axum::http::HeaderMap>,
google_cloud_auth::errors::CredentialsError,
> {
let stored = self.0.read().unwrap();
let token = stored.as_ref().ok_or_else(|| {
google_cloud_auth::errors::CredentialsError::from_msg(false, "GCS credentials missing")
})?;
let mut headers = axum::http::HeaderMap::new();
let mut value: axum::http::HeaderValue = format!("Bearer {}", token.access_token)
.parse()
.map_err(|_| {
google_cloud_auth::errors::CredentialsError::from_msg(
false,
"invalid storage credential",
)
})?;
value.set_sensitive(true);
headers.insert(axum::http::header::AUTHORIZATION, value);
Ok(CacheableResource::New {
entity_tag: EntityTag::new(),
data: headers,
})
}
async fn universe_domain(&self) -> Option<String> {
Some("googleapis.com".into())
}
}
#[cfg(test)]
#[path = "../../tests/unit/host/storage.rs"]
mod tests;