use core::cell::{Cell, RefCell};
use std::path::PathBuf;
#[cfg(not(target_os = "linux"))]
use std::process::Command;
use lean_rs::LeanRuntime;
use lean_rs::Obj;
use lean_rs::ResourceExhaustedFacts;
use lean_rs::error::LeanError;
use lean_rs::error::LeanResult;
use crate::host::cancellation::{LeanCancellationToken, check_cancellation};
use crate::host::capabilities::LeanCapabilities;
use crate::host::progress::LeanProgressSink;
use crate::host::session::{LeanImportStats, LeanSession, LeanSessionImportProfile};
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
pub struct PoolStats {
pub imports_performed: u64,
pub reused: u64,
pub acquired: u64,
pub released_to_pool: u64,
pub released_dropped: u64,
pub drains: u64,
pub drained: u64,
pub fresh_import_refusals: u64,
pub rss_samples: u64,
pub rss_samples_unavailable: u64,
pub key_hits: u64,
pub key_misses: u64,
pub distinct_keys_seen: u64,
pub fresh_imports_avoided: u64,
pub miss_empty_pool: u64,
pub miss_reuse_disabled: u64,
pub miss_no_matching_key: u64,
pub last_miss_reason: Option<SessionPoolKeyMissReason>,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum SessionPoolKeyMissReason {
EmptyPool,
ReuseDisabled,
NoMatchingKey,
}
impl SessionPoolKeyMissReason {
pub const fn label(self) -> &'static str {
match self {
Self::EmptyPool => "empty_pool",
Self::ReuseDisabled => "reuse_disabled",
Self::NoMatchingKey => "no_matching_key",
}
}
}
#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct SessionPoolMemoryPolicy {
max_fresh_imports: Option<u64>,
max_rss_kib: Option<u64>,
}
impl SessionPoolMemoryPolicy {
#[must_use]
pub fn disabled() -> Self {
Self::default()
}
#[must_use]
pub fn max_fresh_imports(mut self, limit: u64) -> Self {
self.max_fresh_imports = Some(limit.max(1));
self
}
#[must_use]
pub fn max_rss_kib(mut self, limit_kib: u64) -> Self {
self.max_rss_kib = Some(limit_kib.max(1));
self
}
#[must_use]
pub fn max_fresh_imports_limit(&self) -> Option<u64> {
self.max_fresh_imports
}
#[must_use]
pub fn max_rss_kib_limit(&self) -> Option<u64> {
self.max_rss_kib
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct SessionPoolConfig {
capacity: usize,
memory_policy: SessionPoolMemoryPolicy,
}
impl SessionPoolConfig {
#[must_use]
pub fn new(capacity: usize) -> Self {
Self {
capacity,
memory_policy: SessionPoolMemoryPolicy::disabled(),
}
}
#[must_use]
pub fn memory_policy(mut self, policy: SessionPoolMemoryPolicy) -> Self {
self.memory_policy = policy;
self
}
#[must_use]
pub fn capacity(&self) -> usize {
self.capacity
}
#[must_use]
pub fn memory_policy_ref(&self) -> &SessionPoolMemoryPolicy {
&self.memory_policy
}
}
#[derive(Clone, Eq, PartialEq)]
struct SessionPoolKey {
project_root: PathBuf,
imports: Vec<String>,
import_profile: LeanSessionImportProfile,
}
impl SessionPoolKey {
fn from_capabilities(
caps: &LeanCapabilities<'_, '_>,
imports: &[&str],
import_profile: LeanSessionImportProfile,
) -> Self {
Self {
project_root: caps.host().project().root().to_path_buf(),
imports: imports.iter().map(|&s| s.to_owned()).collect(),
import_profile,
}
}
}
struct PooledEntry<'lean> {
key: SessionPoolKey,
environment: Obj<'lean>,
import_stats: LeanImportStats,
}
struct PoolInner<'lean> {
free: Vec<PooledEntry<'lean>>,
seen_keys: Vec<SessionPoolKey>,
}
impl<'lean> PoolInner<'lean> {
fn take_matching(&mut self, key: &SessionPoolKey) -> Option<PooledEntry<'lean>> {
let idx = self.free.iter().rposition(|entry| &entry.key == key)?;
Some(self.free.remove(idx))
}
}
pub struct SessionPool<'lean> {
runtime: &'lean LeanRuntime,
capacity: usize,
memory_policy: SessionPoolMemoryPolicy,
inner: RefCell<PoolInner<'lean>>,
last_import_stats: RefCell<Option<LeanImportStats>>,
stats: Cell<PoolStats>,
}
impl<'lean> SessionPool<'lean> {
#[must_use]
pub fn with_capacity(runtime: &'lean LeanRuntime, capacity: usize) -> Self {
Self::with_config(runtime, SessionPoolConfig::new(capacity))
}
#[must_use]
pub fn with_config(runtime: &'lean LeanRuntime, config: SessionPoolConfig) -> Self {
let capacity = config.capacity;
Self {
runtime,
capacity,
memory_policy: config.memory_policy,
inner: RefCell::new(PoolInner {
free: Vec::with_capacity(capacity),
seen_keys: Vec::new(),
}),
last_import_stats: RefCell::new(None),
stats: Cell::new(PoolStats::default()),
}
}
#[must_use]
pub fn with_memory_policy(runtime: &'lean LeanRuntime, capacity: usize, policy: SessionPoolMemoryPolicy) -> Self {
Self::with_config(runtime, SessionPoolConfig::new(capacity).memory_policy(policy))
}
pub fn acquire<'p, 'c>(
&'p self,
caps: &'c LeanCapabilities<'lean, 'c>,
imports: &[&str],
cancellation: Option<&LeanCancellationToken>,
progress: Option<&dyn LeanProgressSink>,
) -> LeanResult<PooledSession<'lean, 'p, 'c>> {
self.acquire_with_profile(
caps,
imports,
LeanSessionImportProfile::default(),
cancellation,
progress,
)
}
pub fn acquire_with_profile<'p, 'c>(
&'p self,
caps: &'c LeanCapabilities<'lean, 'c>,
imports: &[&str],
import_profile: LeanSessionImportProfile,
cancellation: Option<&LeanCancellationToken>,
progress: Option<&dyn LeanProgressSink>,
) -> LeanResult<PooledSession<'lean, 'p, 'c>> {
let _span = tracing::debug_span!(
target: "lean_rs",
"lean_rs.host.pool.acquire",
profile = import_profile.label(),
imports_len = imports.len(),
imports_first = imports.first().copied().unwrap_or("<empty>"),
)
.entered();
check_cancellation(cancellation)?;
debug_assert!(
core::ptr::eq(self.runtime, caps.host().runtime()),
"pool runtime and capability runtime must agree; the shared 'lean parameter normally enforces this",
);
let key = SessionPoolKey::from_capabilities(caps, imports, import_profile);
self.remember_seen_key(&key);
let (session, hit) = {
let mut inner = self.inner.borrow_mut();
if let Some(entry) = inner.take_matching(&key) {
self.bump_reused();
(
LeanSession::from_environment_with_import_stats(caps, entry.environment, entry.import_stats)?,
true,
)
} else {
let reason = self.miss_reason(&inner);
drop(inner);
self.bump_key_miss(reason);
self.enforce_before_fresh_import(imports)?;
let session = caps.session_with_profile(imports, import_profile, cancellation, progress)?;
self.remember_import_stats(session.import_stats().clone());
self.bump_imported();
(session, false)
}
};
tracing::debug!(target: "lean_rs", hit = hit, "lean_rs.host.pool.acquire.result");
Ok(PooledSession {
pool: self,
key,
session: Some(session),
})
}
#[must_use]
pub fn stats(&self) -> PoolStats {
self.stats.get()
}
#[must_use]
pub fn len(&self) -> usize {
self.inner.borrow().free.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.len() == 0
}
#[must_use]
pub fn capacity(&self) -> usize {
self.capacity
}
#[must_use]
pub fn memory_policy(&self) -> &SessionPoolMemoryPolicy {
&self.memory_policy
}
pub fn drain(&self) -> usize {
let mut inner = self.inner.borrow_mut();
let drained = inner.free.len();
inner.free.clear();
let mut s = self.stats.get();
s.drains = s.drains.saturating_add(1);
s.drained = s.drained.saturating_add(u64::try_from(drained).unwrap_or(u64::MAX));
self.stats.set(s);
tracing::debug!(
target: "lean_rs",
drained = drained,
"lean_rs.host.pool.drain",
);
drained
}
fn bump_reused(&self) {
let mut s = self.stats.get();
s.reused = s.reused.saturating_add(1);
s.acquired = s.acquired.saturating_add(1);
s.key_hits = s.key_hits.saturating_add(1);
s.fresh_imports_avoided = s.fresh_imports_avoided.saturating_add(1);
s.last_miss_reason = None;
self.stats.set(s);
}
fn bump_imported(&self) {
let mut s = self.stats.get();
s.imports_performed = s.imports_performed.saturating_add(1);
s.acquired = s.acquired.saturating_add(1);
self.stats.set(s);
}
fn remember_seen_key(&self, key: &SessionPoolKey) {
let mut inner = self.inner.borrow_mut();
if inner.seen_keys.iter().all(|seen| seen != key) {
inner.seen_keys.push(key.clone());
let mut s = self.stats.get();
s.distinct_keys_seen = u64::try_from(inner.seen_keys.len()).unwrap_or(u64::MAX);
self.stats.set(s);
}
}
fn miss_reason(&self, inner: &PoolInner<'_>) -> SessionPoolKeyMissReason {
if self.capacity == 0 {
SessionPoolKeyMissReason::ReuseDisabled
} else if inner.free.is_empty() {
SessionPoolKeyMissReason::EmptyPool
} else {
SessionPoolKeyMissReason::NoMatchingKey
}
}
fn bump_key_miss(&self, reason: SessionPoolKeyMissReason) {
let mut s = self.stats.get();
s.key_misses = s.key_misses.saturating_add(1);
match reason {
SessionPoolKeyMissReason::EmptyPool => {
s.miss_empty_pool = s.miss_empty_pool.saturating_add(1);
}
SessionPoolKeyMissReason::ReuseDisabled => {
s.miss_reuse_disabled = s.miss_reuse_disabled.saturating_add(1);
}
SessionPoolKeyMissReason::NoMatchingKey => {
s.miss_no_matching_key = s.miss_no_matching_key.saturating_add(1);
}
}
s.last_miss_reason = Some(reason);
self.stats.set(s);
}
fn bump_fresh_import_refusal(&self) {
let mut s = self.stats.get();
s.fresh_import_refusals = s.fresh_import_refusals.saturating_add(1);
self.stats.set(s);
}
fn bump_rss_sample(&self, unavailable: bool) {
let mut s = self.stats.get();
s.rss_samples = s.rss_samples.saturating_add(1);
if unavailable {
s.rss_samples_unavailable = s.rss_samples_unavailable.saturating_add(1);
}
self.stats.set(s);
}
fn remember_import_stats(&self, stats: LeanImportStats) {
*self.last_import_stats.borrow_mut() = Some(stats);
}
fn latest_import_stats_diagnostic(&self) -> String {
self.last_import_stats.borrow().as_ref().map_or_else(
|| String::from("last_import_stats=unavailable"),
|stats| format!("last_import_stats=available {}", stats.memory_diagnostic()),
)
}
fn latest_import_stats_for_resource_facts(&self) -> Option<String> {
self.last_import_stats
.borrow()
.as_ref()
.map(LeanImportStats::memory_diagnostic)
}
fn resource_refusal(
&self,
cause: &str,
message: String,
current_rss_kib: Option<u64>,
limit_kib: Option<u64>,
import_count: Option<u64>,
import_limit: Option<u64>,
requested_imports: u64,
) -> LeanError {
lean_rs::__host_internals::host_resource_exhausted_with_facts(
message,
ResourceExhaustedFacts {
cause: cause.to_owned(),
work_entered_lean: false,
current_rss_kib,
limit_kib,
import_count,
import_limit,
requested_imports: Some(requested_imports),
last_import_stats: self.latest_import_stats_for_resource_facts(),
},
)
}
fn enforce_before_fresh_import(&self, imports: &[&str]) -> LeanResult<()> {
let stats = self.stats.get();
if let Some(limit) = self.memory_policy.max_fresh_imports
&& stats.imports_performed >= limit
{
self.bump_fresh_import_refusal();
return Err(self.resource_refusal(
"same_process_fresh_import_limit",
format!(
"same-process SessionPool refused fresh import #{} for {} import(s): max_fresh_imports={limit}; {}; reuse a pooled environment or cycle the worker process",
stats.imports_performed.saturating_add(1),
imports.len(),
self.latest_import_stats_diagnostic(),
),
None,
None,
Some(stats.imports_performed),
Some(limit),
imports.len() as u64,
));
}
if let Some(limit_kib) = self.memory_policy.max_rss_kib {
match current_process_rss_kib() {
Some(current_kib) if current_kib >= limit_kib => {
self.bump_rss_sample(false);
self.bump_fresh_import_refusal();
return Err(self.resource_refusal(
"same_process_rss_ceiling",
format!(
"same-process SessionPool refused fresh import for {} import(s): current RSS {current_kib} KiB reached max_rss_kib={limit_kib}; {}; cycle the worker process to reset Lean process-global import state",
imports.len(),
self.latest_import_stats_diagnostic(),
),
Some(current_kib),
Some(limit_kib),
Some(stats.imports_performed),
None,
imports.len() as u64,
));
}
Some(_) => self.bump_rss_sample(false),
None => {
self.bump_rss_sample(true);
self.bump_fresh_import_refusal();
return Err(self.resource_refusal(
"same_process_rss_sample_unavailable",
format!(
"same-process SessionPool refused fresh import for {} import(s): current RSS sample unavailable while max_rss_kib={limit_kib} is configured; {}",
imports.len(),
self.latest_import_stats_diagnostic(),
),
None,
Some(limit_kib),
Some(stats.imports_performed),
None,
imports.len() as u64,
));
}
}
}
Ok(())
}
fn release(&self, key: SessionPoolKey, env: Obj<'lean>, import_stats: LeanImportStats) {
let mut inner = self.inner.borrow_mut();
let mut s = self.stats.get();
let kept = inner.free.len() < self.capacity;
if kept {
inner.free.push(PooledEntry {
key,
environment: env,
import_stats,
});
s.released_to_pool = s.released_to_pool.saturating_add(1);
} else {
drop(env);
s.released_dropped = s.released_dropped.saturating_add(1);
}
self.stats.set(s);
tracing::trace!(
target: "lean_rs",
kept = kept,
"lean_rs.host.pool.release",
);
}
}
impl core::fmt::Debug for SessionPool<'_> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("SessionPool")
.field("capacity", &self.capacity)
.field("memory_policy", &self.memory_policy)
.field("len", &self.len())
.field("stats", &self.stats.get())
.finish()
}
}
#[cfg(target_os = "linux")]
fn current_process_rss_kib() -> Option<u64> {
let status = std::fs::read_to_string("/proc/self/status").ok()?;
status.lines().find_map(|line| {
let rest = line.strip_prefix("VmRSS:")?;
rest.split_whitespace().next()?.parse::<u64>().ok()
})
}
#[cfg(not(target_os = "linux"))]
fn current_process_rss_kib() -> Option<u64> {
let output = Command::new("ps")
.args(["-o", "rss=", "-p", &std::process::id().to_string()])
.output()
.ok()?;
if !output.status.success() {
return None;
}
let text = String::from_utf8_lossy(&output.stdout);
text.trim().parse::<u64>().ok().filter(|value| *value > 0)
}
pub struct PooledSession<'lean, 'p, 'c> {
pool: &'p SessionPool<'lean>,
key: SessionPoolKey,
session: Option<LeanSession<'lean, 'c>>,
}
impl<'lean, 'c> core::ops::Deref for PooledSession<'lean, '_, 'c> {
type Target = LeanSession<'lean, 'c>;
#[allow(clippy::expect_used, reason = "see PROOF OBLIGATION above")]
fn deref(&self) -> &Self::Target {
self.session
.as_ref()
.expect("session is Some between PooledSession::acquire and Drop::drop")
}
}
#[allow(
single_use_lifetimes,
clippy::elidable_lifetime_names,
reason = "the named lifetimes line up with `Deref::Target = LeanSession<'lean, 'c>` above; \
elision flips the inferred bound and breaks the trait-signature check"
)]
impl<'lean, 'c> core::ops::DerefMut for PooledSession<'lean, '_, 'c> {
#[allow(clippy::expect_used, reason = "see PROOF OBLIGATION on Deref impl")]
fn deref_mut(&mut self) -> &mut LeanSession<'lean, 'c> {
self.session
.as_mut()
.expect("session is Some between PooledSession::acquire and Drop::drop")
}
}
impl Drop for PooledSession<'_, '_, '_> {
fn drop(&mut self) {
if let Some(session) = self.session.take() {
let import_stats = session.import_stats().clone();
let env = session.into_environment();
self.pool.release(self.key.clone(), env, import_stats);
}
}
}
impl core::fmt::Debug for PooledSession<'_, '_, '_> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("PooledSession").finish_non_exhaustive()
}
}