use parking_lot::Mutex;
use std::{
collections::HashSet,
error::Error,
fmt,
sync::{Arc, Weak},
task::Poll,
time::Duration,
};
use uqa_core::notifications::{NotificationEvent, NotificationFailureKind, NotificationIdentity};
use uqa_sql::catalog::roles::session::SessionAuthorization;
use uqa_storage::{PersistentStorageBackend, PersistentStorageProvider};
use super::{inbox::SubscriptionInbox, registration, NotificationHubOwner};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct NotificationSubscriptionOptions {
pub max_active_subscriptions: usize,
pub max_channels: usize,
pub max_queued_notifications: usize,
pub max_queued_bytes: usize,
pub max_registry_entries_per_poll: usize,
}
impl NotificationSubscriptionOptions {
pub(super) fn validate(self) -> Result<(), NotificationSubscriptionError> {
if self.max_active_subscriptions == 0
|| self.max_channels == 0
|| self.max_queued_notifications == 0
|| self.max_queued_bytes == 0
|| self.max_registry_entries_per_poll == 0
{
return Err(NotificationSubscriptionError::new(
NotificationFailureKind::InvalidRequest,
));
}
Ok(())
}
}
pub struct NotificationSubscriptionPermit {
hub: Weak<NotificationHubOwner>,
admission: Arc<registration::SubscriptionAdmission>,
options: NotificationSubscriptionOptions,
}
impl fmt::Debug for NotificationSubscriptionPermit {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("NotificationSubscriptionPermit")
.finish_non_exhaustive()
}
}
fn validate_channels(
channels: &[&str],
options: NotificationSubscriptionOptions,
cancellation: &uqa_core::CancellationToken,
) -> Result<(), NotificationSubscriptionError> {
options.validate()?;
registration::check(cancellation)?;
if channels.is_empty() || channels.len() > options.max_channels {
return Err(NotificationSubscriptionError::new(
NotificationFailureKind::InvalidRequest,
));
}
let mut seen = HashSet::new();
seen.try_reserve(channels.len()).map_err(|error| {
NotificationSubscriptionError::with_source(NotificationFailureKind::Capacity, error)
})?;
for channel in channels {
registration::check(cancellation)?;
if channel.contains('\0')
|| super::validate_channel(channel).is_err()
|| !seen.insert(*channel)
{
return Err(NotificationSubscriptionError::new(
NotificationFailureKind::InvalidRequest,
));
}
}
registration::check(cancellation)
}
#[derive(Clone)]
pub struct NotificationSubscriptionError {
kind: NotificationFailureKind,
original: Option<Arc<dyn Error + Send + Sync>>,
}
impl NotificationSubscriptionError {
pub(super) fn new(kind: NotificationFailureKind) -> Self {
Self {
kind,
original: None,
}
}
pub(super) fn with_source(
kind: NotificationFailureKind,
error: impl Error + Send + Sync + 'static,
) -> Self {
Self {
kind,
original: Some(Arc::new(error)),
}
}
pub fn kind(&self) -> NotificationFailureKind {
self.kind
}
pub fn code(&self) -> &'static str {
self.kind.code()
}
pub fn original_error(&self) -> Option<&(dyn Error + Send + Sync + 'static)> {
self.original.as_deref()
}
}
impl fmt::Debug for NotificationSubscriptionError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("NotificationSubscriptionError")
.field("kind", &self.kind)
.finish_non_exhaustive()
}
}
impl fmt::Display for NotificationSubscriptionError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(self.code())
}
}
impl Error for NotificationSubscriptionError {}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum NotificationWait {
Event(NotificationEvent),
TimedOut,
Closed,
}
pub(super) struct ListenerResources {
pub(super) hub: Arc<NotificationHubOwner>,
pub(super) session_id: u64,
pub(super) locks: Arc<crate::row_locks::RowLockManager>,
pub(super) _provider: Option<Arc<dyn PersistentStorageProvider>>,
pub(super) _backend: Option<Arc<dyn PersistentStorageBackend>>,
pub(super) _admission: Arc<registration::SubscriptionAdmission>,
}
impl Drop for ListenerResources {
fn drop(&mut self) {
self.hub.retire_subscription(self.session_id);
self.locks.release_session(self.session_id);
}
}
pub struct NotificationSubscription {
pub(super) inbox: Arc<SubscriptionInbox>,
pub(super) resources: Mutex<Option<ListenerResources>>,
pub(super) authorization: SessionAuthorization,
}
impl NotificationSubscription {
pub fn identity(&self) -> &NotificationIdentity {
self.inbox.identity()
}
pub fn poll(&self) -> Result<Poll<Option<NotificationEvent>>, NotificationSubscriptionError> {
self.inbox.poll()
}
pub fn wait(
&self,
timeout: Duration,
) -> Result<NotificationWait, NotificationSubscriptionError> {
self.inbox.wait(timeout)
}
pub async fn next_event(
&self,
) -> Result<Option<NotificationEvent>, NotificationSubscriptionError> {
self.inbox.wait_async().await
}
pub fn is_closed(&self) -> bool {
self.inbox.is_closed()
}
pub fn role_identity(&self) -> uqa_core::catalog_role::RoleIdentity {
self.authorization.current().identity()
}
pub fn stop_delivery(&self) {
self.inbox.close();
}
pub fn close(&self) {
self.stop_delivery();
let mut resources = self.resources.lock();
drop(resources.take());
}
}
impl Drop for NotificationSubscription {
fn drop(&mut self) {
self.close();
}
}
impl crate::Engine {
pub fn subscribe_notifications(
&self,
channels: &[&str],
options: NotificationSubscriptionOptions,
) -> Result<NotificationSubscription, NotificationSubscriptionError> {
self.subscribe_notifications_with_cancellation(
channels,
options,
&uqa_core::CancellationToken::new(),
)
}
pub fn subscribe_notifications_with_cancellation(
&self,
channels: &[&str],
options: NotificationSubscriptionOptions,
cancellation: &uqa_core::CancellationToken,
) -> Result<NotificationSubscription, NotificationSubscriptionError> {
validate_channels(channels, options, cancellation)?;
let permit = self.reserve_notification_subscription(options, cancellation)?;
self.register_notification_subscription(channels, permit, cancellation)
}
pub fn reserve_notification_subscription(
&self,
options: NotificationSubscriptionOptions,
cancellation: &uqa_core::CancellationToken,
) -> Result<NotificationSubscriptionPermit, NotificationSubscriptionError> {
options.validate()?;
let admission = self
.notification_hub
.admissions
.reserve(options.max_active_subscriptions, cancellation)?;
Ok(NotificationSubscriptionPermit {
hub: Arc::downgrade(&self.notification_hub),
admission,
options,
})
}
pub fn subscribe_notifications_with_permit(
&self,
channels: &[&str],
permit: NotificationSubscriptionPermit,
cancellation: &uqa_core::CancellationToken,
) -> Result<NotificationSubscription, NotificationSubscriptionError> {
if !Weak::ptr_eq(&permit.hub, &Arc::downgrade(&self.notification_hub)) {
return Err(NotificationSubscriptionError::new(
NotificationFailureKind::InvalidRequest,
));
}
validate_channels(channels, permit.options, cancellation)?;
self.register_notification_subscription(channels, permit, cancellation)
}
fn register_notification_subscription(
&self,
channels: &[&str],
permit: NotificationSubscriptionPermit,
cancellation: &uqa_core::CancellationToken,
) -> Result<NotificationSubscription, NotificationSubscriptionError> {
let _statement = self
.runtime
.statement_gate
.lock_with_cancellation(cancellation)
.map_err(|error| registration::failure(error.into()))?;
self.prepare_notification_recovery_with_cancellation(Some(cancellation))
.map_err(registration::failure)?;
let authorization = registration::read(&self.session.state, cancellation)
.map_err(registration::failure)?
.authorization
.clone();
let mut bytes = [0u8; 16];
getrandom::fill(&mut bytes).map_err(|error| {
NotificationSubscriptionError::with_source(
NotificationFailureKind::SourceUnavailable,
std::io::Error::other(error.to_string()),
)
})?;
bytes[6] = (bytes[6] & 0x0f) | 0x40;
bytes[8] = (bytes[8] & 0x3f) | 0x80;
let identity = NotificationIdentity {
epoch: uqa_core::notifications::NotificationEpoch::from_bytes(bytes)
.expect("version and variant set above"),
request_id: None,
};
let session_id = self.row_locks.allocate_session();
let inbox = SubscriptionInbox::new(identity, permit.options);
self.notification_hub.register_subscription(
session_id,
self.backend_process_id(),
channels
.iter()
.map(|channel| (*channel).to_owned())
.collect(),
&inbox,
cancellation,
)?;
Ok(NotificationSubscription {
inbox,
resources: Mutex::new(Some(ListenerResources {
hub: self.notification_hub.clone(),
session_id,
locks: self.row_locks.clone(),
_provider: self.storage.provider.clone(),
_backend: self.storage.backend.clone(),
_admission: permit.admission,
})),
authorization,
})
}
}