use crate::config::GatewayConfig;
use crate::federation_transport::GatewayFederationTransport;
use crate::ha::{
GatewayHaCoordinator, GatewayHaLifecycle, GatewayHaLifecycleSnapshot,
GatewayHaOwnershipSnapshot, GatewayHaOwnershipSource, GatewayHaSessionSnapshot,
GatewayHaWorkerSnapshot, GatewayRegistryError, GatewayRegistryResult,
};
use crate::metrics::GatewayMetrics;
use crate::tenant_directory::{SharedTenantState, TenantDirectory, TenantDirectorySnapshot};
use crate::GatewayResult;
use appcore_peer_rpc::{BoundedReplayStore, PeerNonceStore, ReplayStoreConfig};
use appcore_security::HashTokenProvider;
use appcore_types::{ClusterId, TenantId};
use parking_lot::Mutex;
use std::sync::Arc;
use tokio::sync::watch;
pub struct GatewayState {
config: GatewayConfig,
tenants: TenantDirectory,
connection_admission: Mutex<()>,
pub metrics: Arc<GatewayMetrics>,
pub token_provider: HashTokenProvider,
connection_replay: Arc<dyn PeerNonceStore>,
ha_lifecycle: Option<Arc<GatewayHaLifecycle>>,
ha_coordinator: Option<Arc<GatewayHaCoordinator>>,
federation_transport: GatewayFederationTransport,
shutdown: watch::Sender<bool>,
}
impl GatewayState {
pub fn new(config: GatewayConfig, token_provider: HashTokenProvider) -> GatewayResult<Self> {
Self::with_replay_store(
config,
token_provider,
Arc::new(BoundedReplayStore::new(ReplayStoreConfig::default())),
)
}
pub fn with_replay_store(
config: GatewayConfig,
token_provider: HashTokenProvider,
connection_replay: Arc<dyn PeerNonceStore>,
) -> GatewayResult<Self> {
Self::build(config, token_provider, connection_replay, None, None)
}
pub fn with_ha_lifecycle(
config: GatewayConfig,
token_provider: HashTokenProvider,
connection_replay: Arc<dyn PeerNonceStore>,
ha_lifecycle: Arc<GatewayHaLifecycle>,
) -> GatewayResult<Self> {
Self::build(
config,
token_provider,
connection_replay,
Some(ha_lifecycle),
None,
)
}
pub fn with_ha_coordinator(
config: GatewayConfig,
token_provider: HashTokenProvider,
connection_replay: Arc<dyn PeerNonceStore>,
coordinator: Arc<GatewayHaCoordinator>,
) -> GatewayResult<Self> {
let lifecycle = coordinator.lifecycle();
Self::build(
config,
token_provider,
connection_replay,
Some(lifecycle),
Some(coordinator),
)
}
fn build(
config: GatewayConfig,
token_provider: HashTokenProvider,
connection_replay: Arc<dyn PeerNonceStore>,
ha_lifecycle: Option<Arc<GatewayHaLifecycle>>,
ha_coordinator: Option<Arc<GatewayHaCoordinator>>,
) -> GatewayResult<Self> {
config.validate()?;
let (shutdown, _) = watch::channel(false);
Ok(Self {
config,
tenants: TenantDirectory::new(),
connection_admission: Mutex::new(()),
metrics: GatewayMetrics::new(),
token_provider,
connection_replay,
ha_lifecycle,
ha_coordinator,
federation_transport: GatewayFederationTransport::default(),
shutdown,
})
}
pub fn config(&self) -> &GatewayConfig {
&self.config
}
pub fn tenant_count(&self) -> usize {
self.tenants.len()
}
pub fn connection_count(&self) -> usize {
self.tenants.connection_count()
}
pub fn ha_lifecycle_snapshot(&self) -> Option<GatewayHaLifecycleSnapshot> {
self.ha_lifecycle
.as_ref()
.map(|lifecycle| lifecycle.snapshot())
}
pub(crate) fn admit_ha_work(&self) -> GatewayRegistryResult<()> {
self.ha_lifecycle
.as_ref()
.map_or(Ok(()), |lifecycle| lifecycle.admit())
}
pub(crate) fn admit_ha_tenant(&self, tenant_id: &TenantId) -> GatewayRegistryResult<()> {
match &self.ha_coordinator {
Some(coordinator) => coordinator.lease_for(tenant_id).map(|_| ()),
None => self.admit_ha_work(),
}
}
pub(crate) fn admit_ha_boundary(
&self,
tenant_id: &TenantId,
cluster_id: &ClusterId,
) -> GatewayRegistryResult<()> {
match &self.ha_coordinator {
Some(coordinator) => coordinator.lease_for(tenant_id).and_then(|lease| {
(lease.cluster_id() == cluster_id)
.then_some(())
.ok_or(GatewayRegistryError::InvalidContract)
}),
None => self.admit_ha_work(),
}
}
pub(crate) fn ha_coordinator(&self) -> Option<Arc<GatewayHaCoordinator>> {
self.ha_coordinator.as_ref().map(Arc::clone)
}
pub(crate) fn federation_transport(&self) -> GatewayFederationTransport {
self.federation_transport.clone()
}
pub fn tenant_partition(&self, tenant_id: &TenantId) -> Option<SharedTenantState> {
self.tenants.get(tenant_id)
}
pub fn tenant_partition_or_insert(
&self,
tenant_id: &TenantId,
) -> GatewayResult<SharedTenantState> {
self.tenants.get_or_insert(tenant_id)
}
pub(crate) fn tenant_snapshot(&self) -> TenantDirectorySnapshot {
self.tenants.snapshot()
}
pub(crate) fn lock_connection_admission(&self) -> parking_lot::MutexGuard<'_, ()> {
self.connection_admission.lock()
}
pub fn request_shutdown(&self) {
if let Some(lifecycle) = &self.ha_lifecycle {
lifecycle.stop();
}
self.shutdown.send_replace(true);
}
pub fn is_shutting_down(&self) -> bool {
*self.shutdown.borrow()
}
pub(crate) fn subscribe_shutdown(&self) -> watch::Receiver<bool> {
self.shutdown.subscribe()
}
pub(crate) fn connection_replay(&self) -> &dyn PeerNonceStore {
self.connection_replay.as_ref()
}
pub(crate) async fn wait_for_shutdown(&self) {
let mut shutdown = self.subscribe_shutdown();
while !*shutdown.borrow() {
if shutdown.changed().await.is_err() {
break;
}
}
}
}
impl GatewayHaOwnershipSource for GatewayState {
fn snapshot(&self, now_ms: u64) -> GatewayRegistryResult<GatewayHaOwnershipSnapshot> {
let mut snapshot = GatewayHaOwnershipSnapshot::default();
let tenants = self.tenant_snapshot();
for (tenant_id, tenant) in tenants.iter() {
let tenant = tenant.read();
for worker in tenant.workers.values() {
let cluster_id = worker
.cluster_id()
.cloned()
.ok_or(GatewayRegistryError::InvalidContract)?;
snapshot.workers.push(GatewayHaWorkerSnapshot {
tenant_id: tenant_id.clone(),
cluster_id,
registration: crate::GatewayWorkerRegistration::new(
worker.key.installation_id.clone(),
worker.key.core_id.clone(),
worker.generation(),
tenant.registry.capabilities_for(&worker.key),
)?,
});
}
snapshot.sessions.extend(
tenant
.sessions
.values()
.filter(|session| !session.is_expired(now_ms))
.map(|session| GatewayHaSessionSnapshot {
tenant_id: tenant_id.clone(),
session_id: session.session_id.clone(),
expires_at_ms: session.expires_at_ms,
}),
);
}
Ok(snapshot)
}
}
#[cfg(test)]
#[path = "state_tests.rs"]
mod tests;