mod builder;
mod lifecycle;
pub use builder::NntpProxyBuilder;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, AtomicUsize};
use std::time::Instant;
use anyhow::Result;
use tracing::info;
use crate::auth::AuthHandler;
use crate::cache::UnifiedCache;
use crate::config::{Memory, RoutingMode, Server};
use crate::metrics::{ConnectionStatsAggregator, MetricsCollector};
use crate::pool::{BufferPool, DeadpoolConnectionProvider};
use crate::router;
#[derive(Debug, Clone)]
pub struct NntpProxy {
pub(super) servers: Arc<[Server]>,
pub(super) router: Arc<router::BackendSelector>,
pub(super) connection_providers: Vec<DeadpoolConnectionProvider>,
pub(super) buffer_pool: BufferPool,
pub(super) routing_mode: RoutingMode,
pub(super) auth_handler: Arc<AuthHandler>,
pub(super) metrics: MetricsCollector,
pub(super) connection_stats: ConnectionStatsAggregator,
pub(super) cache: Arc<UnifiedCache>,
pub(super) memory: Memory,
pub(super) store_article_bodies: bool,
pub(super) adaptive_precheck: bool,
pub(super) last_activity_nanos: Arc<AtomicU64>,
pub(super) active_clients: Arc<AtomicUsize>,
pub(super) start_instant: Instant,
}
impl NntpProxy {
pub async fn new(config: crate::config::Config, routing_mode: RoutingMode) -> Result<Self> {
NntpProxyBuilder::new(config)
.with_routing_mode(routing_mode)
.build()
.await
}
pub fn new_sync(config: crate::config::Config, routing_mode: RoutingMode) -> Result<Self> {
NntpProxyBuilder::new(config)
.with_routing_mode(routing_mode)
.build_sync()
}
#[must_use]
pub const fn builder(config: crate::config::Config) -> NntpProxyBuilder {
NntpProxyBuilder::new(config)
}
pub async fn graceful_shutdown(&self) {
info!("Initiating graceful shutdown...");
for provider in &self.connection_providers {
provider.shutdown();
}
info!("Flushing disk cache writes...");
match tokio::time::timeout(crate::constants::timeout::CACHE_CLOSE, self.cache.close()).await
{
Ok(Ok(())) => {}
Ok(Err(e)) => tracing::warn!("Error closing cache: {}", e),
Err(_) => tracing::warn!(
"Cache close timed out after {}s, forcing shutdown",
crate::constants::timeout::CACHE_CLOSE.as_secs()
),
}
info!("Shutting down connection pools...");
for provider in &self.connection_providers {
provider.graceful_shutdown().await;
}
info!("All connection pools have been shut down gracefully");
}
pub async fn prewarm_connections(&self) -> Result<()> {
crate::pool::prewarm_pools(&self.connection_providers, &self.servers).await
}
#[must_use]
#[inline]
pub fn servers(&self) -> &[Server] {
&self.servers
}
#[must_use]
#[inline]
pub const fn router(&self) -> &Arc<router::BackendSelector> {
&self.router
}
#[must_use]
#[inline]
pub fn connection_providers(&self) -> &[DeadpoolConnectionProvider] {
&self.connection_providers
}
#[must_use]
#[inline]
pub const fn buffer_pool(&self) -> &BufferPool {
&self.buffer_pool
}
#[must_use]
#[inline]
pub const fn cache(&self) -> &Arc<UnifiedCache> {
&self.cache
}
#[must_use]
#[inline]
pub const fn metrics(&self) -> &MetricsCollector {
&self.metrics
}
#[must_use]
#[inline]
pub const fn connection_stats(&self) -> &ConnectionStatsAggregator {
&self.connection_stats
}
}
#[cfg(test)]
pub mod tests {
use super::*;
use crate::config::{Config, Server};
use crate::types::{MaxConnections, Port};
pub(in crate::proxy) fn create_test_config() -> Config {
Config {
servers: vec![
Server::builder("server1.example.com", Port::try_new(119).unwrap())
.name("Test Server 1")
.max_connections(MaxConnections::try_new(5).unwrap())
.build()
.unwrap(),
Server::builder("server2.example.com", Port::try_new(119).unwrap())
.name("Test Server 2")
.max_connections(MaxConnections::try_new(8).unwrap())
.build()
.unwrap(),
Server::builder("server3.example.com", Port::try_new(119).unwrap())
.name("Test Server 3")
.max_connections(MaxConnections::try_new(12).unwrap())
.build()
.unwrap(),
],
..Default::default()
}
}
#[test]
fn test_proxy_creation_with_servers() {
let config = create_test_config();
let proxy = Arc::new(
NntpProxy::new_sync(config, RoutingMode::Stateful).expect("Failed to create proxy"),
);
assert_eq!(proxy.servers().len(), 3);
assert_eq!(proxy.servers()[0].name.as_str(), "Test Server 1");
}
#[test]
fn test_proxy_creation_with_empty_servers() {
let config = Config {
servers: vec![],
..Default::default()
};
let result = NntpProxy::new_sync(config, RoutingMode::Stateful);
assert!(result.is_err());
assert!(
result
.unwrap_err()
.to_string()
.contains("No servers configured")
);
}
#[test]
fn test_proxy_has_router() {
let config = create_test_config();
let proxy = Arc::new(
NntpProxy::new_sync(config, RoutingMode::Stateful).expect("Failed to create proxy"),
);
assert_eq!(proxy.router.backend_count(), 3);
}
#[test]
fn test_new_sync_constructs_proxy() {
let config = create_test_config();
let proxy = NntpProxy::new_sync(config, RoutingMode::Stateful)
.expect("Failed to create proxy with new_sync()");
assert_eq!(proxy.servers().len(), 3);
assert_eq!(proxy.router.backend_count(), 3);
}
}