Atomic Lifecycle for Writer Refresh + Hardswap

This commit is contained in:
Alexey
2026-09-12 20:14:41 +03:00
parent 54197e848a
commit 1ea5f7a1d6
44 changed files with 3020 additions and 581 deletions
@@ -54,23 +54,45 @@ impl MePool {
contour: WriterContour,
writer_dc: i32,
) -> Result<()> {
self.connect_one_with_generation_contour_for_dc_with_cap_policy(
addr, rng, generation, contour, writer_dc, false,
self.connect_one_with_generation_contour_for_dc_with_intent(
addr,
rng,
generation,
contour,
writer_dc,
WriterOpenIntent::Normal,
)
.await
}
pub(in crate::transport::middle_proxy) async fn connect_one_with_generation_contour_for_dc_with_cap_policy(
/// Connects and publishes one writer under an explicit capacity intent.
pub(in crate::transport::middle_proxy) async fn connect_one_with_generation_contour_for_dc_with_intent(
self: &Arc<Self>,
addr: SocketAddr,
rng: &SecureRandom,
generation: u64,
contour: WriterContour,
writer_dc: i32,
allow_coverage_override: bool,
intent: WriterOpenIntent,
) -> Result<()> {
let Some(_writer_open_reservation) = self
.reserve_writer_open(contour, allow_coverage_override, writer_dc)
let prepared = self
.prepare_writer_with_intent(addr, rng, generation, contour, writer_dc, intent)
.await?;
self.publish_connected_writer(prepared).await
}
/// Completes all cancellable connection work without publishing pool visibility.
pub(super) async fn prepare_writer_with_intent<'a>(
self: &'a Arc<Self>,
addr: SocketAddr,
rng: &SecureRandom,
generation: u64,
contour: WriterContour,
writer_dc: i32,
intent: WriterOpenIntent,
) -> Result<PreparedWriter<'a>> {
let Some(writer_open_reservation) = self
.reserve_writer_open(contour, intent, writer_dc)
.await
else {
return Err(ProxyError::Proxy(format!(
@@ -264,36 +286,15 @@ impl MePool {
})
};
self.publish_prepared_writer(writer, tx, byte_budget, task_registration, writer_task)
.await;
Ok(())
}
/// Commits writer visibility and lifecycle ownership after all cancellation points.
#[allow(clippy::too_many_arguments)]
pub(in crate::transport::middle_proxy) async fn publish_prepared_writer<F>(
self: &Arc<Self>,
writer: MeWriter,
tx: mpsc::Sender<WriterCommand>,
byte_budget: Arc<tokio::sync::Semaphore>,
task_registration: MeTaskRegistration<'_>,
writer_task: F,
) where
F: Future<Output = ()> + Send + 'static,
{
let (mut writers, mut registry_registration) = tokio::join!(
self.writers.write(),
self.registry.prepare_writer_registration()
);
registry_registration.install(writer.id, tx, byte_budget);
writers.push(writer);
self.conn_count.fetch_add(1, Ordering::Relaxed);
self.lifecycle
.spawn_registered_writer(task_registration, writer_task);
drop(writers);
drop(registry_registration);
self.notify_writer_epoch();
Ok(PreparedWriter {
writer,
tx,
byte_budget,
task_registration,
writer_task,
intent,
_open_reservation: writer_open_reservation,
})
}
pub(crate) async fn remove_writer_and_close_clients(self: &Arc<Self>, writer_id: u64) {
@@ -331,7 +332,7 @@ impl MePool {
) -> bool {
let mut close_tx: Option<mpsc::Sender<WriterCommand>> = None;
let mut removed_addr: Option<SocketAddr> = None;
let mut removed_dc: Option<i32> = None;
let mut removed_role: Option<WriterRole> = None;
let mut removed_uptime: Option<Duration> = None;
let mut trigger_refill = false;
let mut removed = false;
@@ -352,7 +353,7 @@ impl MePool {
self.stats.increment_me_writer_removed_total();
w.cancel.cancel();
removed_addr = Some(w.addr);
removed_dc = Some(w.writer_dc);
removed_role = Some(WriterRole::from_writer(&w));
removed_uptime = Some(w.created_at.elapsed());
trigger_refill = !was_draining;
if trigger_refill {
@@ -397,8 +398,8 @@ impl MePool {
);
}
}
if trigger_refill && let Some(writer_dc) = removed_dc {
self.trigger_immediate_refill_for_dc(addr, writer_dc);
if trigger_refill && let Some(role) = removed_role {
self.trigger_immediate_refill_for_role(addr, role);
}
}
if removed {
@@ -415,26 +416,13 @@ impl MePool {
) {
let timeout = timeout.filter(|d| !d.is_zero());
let found = {
let mut ws = self.writers.write().await;
if let Some(w) = ws.iter_mut().find(|w| w.id == writer_id) {
let already_draining = w.draining.swap(true, Ordering::Relaxed);
w.allow_drain_fallback
.store(allow_drain_fallback, Ordering::Relaxed);
let now_epoch_secs = Self::now_epoch_secs();
w.draining_started_at_epoch_secs
.store(now_epoch_secs, Ordering::Relaxed);
let drain_deadline_epoch_secs = timeout
.map(|duration| now_epoch_secs.saturating_add(duration.as_secs()))
.unwrap_or(0);
w.drain_deadline_epoch_secs
.store(drain_deadline_epoch_secs, Ordering::Relaxed);
if !already_draining {
self.stats.increment_pool_drain_active();
self.increment_draining_active_runtime();
let ws = self.writers.write().await;
let mut registry_registration = self.registry.prepare_writer_registration().await;
if let Some(w) = ws.iter().find(|w| w.id == writer_id) {
if !allow_drain_fallback {
registry_registration.retire(writer_id);
}
w.contour
.store(WriterContour::Draining.as_u8(), Ordering::Relaxed);
w.draining.store(true, Ordering::Relaxed);
self.apply_writer_draining_state(w, timeout, allow_drain_fallback);
true
} else {
false
@@ -452,6 +440,39 @@ impl MePool {
);
}
/// Publishes drain metadata before the release-store that makes draining observable.
pub(in crate::transport::middle_proxy) fn apply_writer_draining_state(
&self,
writer: &MeWriter,
timeout: Option<Duration>,
allow_drain_fallback: bool,
) -> bool {
let timeout = timeout.filter(|duration| !duration.is_zero());
let already_draining = writer.draining.load(Ordering::Acquire);
writer
.allow_drain_fallback
.store(allow_drain_fallback, Ordering::Release);
let now_epoch_secs = Self::now_epoch_secs();
writer
.draining_started_at_epoch_secs
.store(now_epoch_secs, Ordering::Release);
let drain_deadline_epoch_secs = timeout
.map(|duration| now_epoch_secs.saturating_add(duration.as_secs()))
.unwrap_or(0);
writer
.drain_deadline_epoch_secs
.store(drain_deadline_epoch_secs, Ordering::Release);
writer
.contour
.store(WriterContour::Draining.as_u8(), Ordering::Release);
writer.draining.store(true, Ordering::Release);
if !already_draining {
self.stats.increment_pool_drain_active();
self.increment_draining_active_runtime();
}
!already_draining
}
pub(crate) async fn mark_writer_draining(self: &Arc<Self>, writer_id: u64) {
self.mark_writer_draining_with_timeout(writer_id, Some(Duration::from_secs(300)), false)
.await;
@@ -461,7 +482,7 @@ impl MePool {
&self,
writer: &MeWriter,
) -> bool {
if !writer.draining.load(Ordering::Relaxed) {
if !writer.draining.load(Ordering::Acquire) {
return true;
}
if !writer.allow_drain_fallback.load(Ordering::Relaxed) {