use super::*; const WRITER_REPLACEMENT_OPEN_LIMIT_MAX: usize = 128; #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] /// Immutable identity of a writer lifecycle role at replacement reservation time. pub(in crate::transport::middle_proxy) struct WriterRole { /// Telegram DC owning the writer. pub(in crate::transport::middle_proxy) dc: i32, /// Address family of the writer endpoint. pub(in crate::transport::middle_proxy) family: IpFamily, /// Pool generation owning the writer. pub(in crate::transport::middle_proxy) generation: u64, /// Lifecycle contour assigned to the writer. pub(in crate::transport::middle_proxy) contour: WriterContour, } impl WriterRole { /// Captures the current role of an installed writer. pub(in crate::transport::middle_proxy) fn from_writer(writer: &MeWriter) -> Self { Self { dc: writer.writer_dc, family: if writer.addr.is_ipv4() { IpFamily::V4 } else { IpFamily::V6 }, generation: writer.generation, contour: WriterContour::from_u8(writer.contour.load(Ordering::Acquire)), } } /// Revalidates that an installed writer still has this exact role. pub(in crate::transport::middle_proxy) fn matches(self, writer: &MeWriter) -> bool { self == Self::from_writer(writer) } } #[derive(Debug, Clone, Copy, PartialEq, Eq)] /// Capacity policy applied while opening a writer. pub(in crate::transport::middle_proxy) enum WriterOpenIntent { /// Ordinary pool growth constrained by the configured contour cap. Normal, /// Required active coverage allowed to exceed an undersized configured cap temporarily. Coverage, /// Replacement-before-drain capacity owned by an existing victim reservation. Replacement, } /// RAII ownership of one bounded in-flight writer open. pub(in crate::transport::middle_proxy) struct WriterOpenReservation<'a> { counter: Option<&'a AtomicUsize>, } impl Drop for WriterOpenReservation<'_> { fn drop(&mut self) { if let Some(counter) = self.counter { counter.fetch_sub(1, Ordering::AcqRel); } } } impl MePool { /// Computes the authoritative active-writer floor across enabled families and DCs. pub(in crate::transport::middle_proxy) async fn active_coverage_required_total(&self) -> usize { let now_epoch_secs = Self::now_epoch_secs(); let mut required_total = 0usize; let endpoint_snapshot = self.endpoint_snapshot.load_full(); if self.family_enabled_for_drain_coverage(IpFamily::V4, now_epoch_secs) { for addrs in endpoint_snapshot.map_v4.values() { let mut endpoints = HashSet::::new(); for (ip, port) in addrs.iter().copied() { endpoints.insert(SocketAddr::new(ip, port)); } required_total = required_total.saturating_add( self.required_writers_for_dc_with_floor_mode(endpoints.len(), false), ); } } if self.family_enabled_for_drain_coverage(IpFamily::V6, now_epoch_secs) { for addrs in endpoint_snapshot.map_v6.values() { let mut endpoints = HashSet::::new(); for (ip, port) in addrs.iter().copied() { endpoints.insert(SocketAddr::new(ip, port)); } required_total = required_total.saturating_add( self.required_writers_for_dc_with_floor_mode(endpoints.len(), false), ); } } required_total } /// Reports whether one writer may be opened under the selected contour policy. pub(in crate::transport::middle_proxy) async fn can_open_writer_for_contour( &self, contour: WriterContour, intent: WriterOpenIntent, writer_dc: i32, target_addr: SocketAddr, ) -> bool { if intent == WriterOpenIntent::Replacement { return true; } let family = if target_addr.is_ipv4() { IpFamily::V4 } else { IpFamily::V6 }; let endpoint_snapshot = self.endpoint_snapshot.load_full(); let endpoints = endpoint_snapshot.endpoints_for_dc_family(writer_dc, family); if !endpoint_snapshot.contains_dc_endpoint(writer_dc, target_addr) { return false; } let (active_writers, warm_writers, _) = self.non_draining_writer_counts_by_contour().await; let live = match contour { WriterContour::Active => active_writers, WriterContour::Warm => warm_writers, WriterContour::Draining => return true, }; let configured_cap = match contour { WriterContour::Active => self.adaptive_floor_active_cap_configured_total(), WriterContour::Warm => self.adaptive_floor_warm_cap_configured_total(), WriterContour::Draining => usize::MAX, }; if live < configured_cap { return true; } if intent != WriterOpenIntent::Coverage { return false; } let endpoint_count = endpoints.len(); if endpoint_count == 0 { return false; } let required = self.required_writers_for_dc_with_floor_mode(endpoint_count, false); let status = self.reinit.status.load(); let generation = match contour { WriterContour::Active => status.active_generation, WriterContour::Warm => status.pending_hardswap_generation, WriterContour::Draining => 0, }; let family_count = { let writers = self.writers.read().await; writers .iter() .filter(|writer| { !writer.draining.load(Ordering::Relaxed) && writer.writer_dc == writer_dc && writer.generation == generation && WriterContour::from_u8(writer.contour.load(Ordering::Relaxed)) == contour && writer.addr.is_ipv4() == (family == IpFamily::V4) && endpoint_snapshot.contains_dc_endpoint(writer_dc, writer.addr) }) .count() }; if family_count < required { return true; } false } /// Reserves bounded transient capacity for a writer open attempt. pub(in crate::transport::middle_proxy) async fn reserve_writer_open( &self, contour: WriterContour, intent: WriterOpenIntent, writer_dc: i32, target_addr: SocketAddr, ) -> Option> { let counter = match contour { WriterContour::Active => &self.writer_connect_active_reserved, WriterContour::Warm => &self.writer_connect_warm_reserved, WriterContour::Draining => { return Some(WriterOpenReservation { counter: None }); } }; if intent == WriterOpenIntent::Replacement { let configured_cap = match contour { WriterContour::Active => self.adaptive_floor_active_cap_configured_total(), WriterContour::Warm => self.adaptive_floor_warm_cap_configured_total(), WriterContour::Draining => usize::MAX, }; let effective_cap = match contour { WriterContour::Active => self .floor_runtime .me_adaptive_floor_active_cap_effective .load(Ordering::Acquire) as usize, WriterContour::Warm => self .floor_runtime .me_adaptive_floor_warm_cap_effective .load(Ordering::Acquire) as usize, WriterContour::Draining => usize::MAX, }; let replacement_limit = configured_cap .max(effective_cap) .max(1) .min(WRITER_REPLACEMENT_OPEN_LIMIT_MAX); loop { let reserved = self.writer_replacement_open_reserved.load(Ordering::Acquire); if reserved >= replacement_limit { return None; } if self .writer_replacement_open_reserved .compare_exchange_weak( reserved, reserved + 1, Ordering::AcqRel, Ordering::Acquire, ) .is_ok() { return Some(WriterOpenReservation { counter: Some(&self.writer_replacement_open_reserved), }); } } } loop { if !self .can_open_writer_for_contour(contour, intent, writer_dc, target_addr) .await { return None; } let (active_writers, warm_writers, _) = self.non_draining_writer_counts_by_contour().await; let live = match contour { WriterContour::Active => active_writers, WriterContour::Warm => warm_writers, WriterContour::Draining => 0, }; let mut limit = match contour { WriterContour::Active => self.adaptive_floor_active_cap_configured_total(), WriterContour::Warm => self.adaptive_floor_warm_cap_configured_total(), WriterContour::Draining => usize::MAX, }; if intent == WriterOpenIntent::Coverage && matches!(contour, WriterContour::Active | WriterContour::Warm) { limit = limit .max(self.active_coverage_required_total().await) .saturating_add( self.reconnect_runtime .me_reconnect_max_concurrent_per_dc .max(1) as usize, ); } let reserved = counter.load(Ordering::Acquire); if live.saturating_add(reserved) >= limit { return None; } if counter .compare_exchange_weak(reserved, reserved + 1, Ordering::AcqRel, Ordering::Acquire) .is_ok() { return Some(WriterOpenReservation { counter: Some(counter), }); } } } /// Resolves a DC writer floor for static or adaptive-idle operation. pub(in crate::transport::middle_proxy) fn required_writers_for_dc_with_floor_mode( &self, endpoint_count: usize, reduce_for_idle: bool, ) -> usize { let base_required = self.required_writers_for_dc(endpoint_count); if !reduce_for_idle { return base_required; } if self.floor_mode() != MeFloorMode::Adaptive { return base_required; } let min_writers = if endpoint_count == 1 { (self .floor_runtime .me_adaptive_floor_min_writers_single_endpoint .load(Ordering::Relaxed) as usize) .max(1) } else { (self .floor_runtime .me_adaptive_floor_min_writers_multi_endpoint .load(Ordering::Relaxed) as usize) .max(1) }; base_required.min(min_writers) } }