From a8ff0a8405c7d7f25f401214fadf3fc751270f1f Mon Sep 17 00:00:00 2001 From: Artur Skubei <41114720+gdlbo@users.noreply.github.com> Date: Thu, 8 Oct 2026 19:04:14 +0300 Subject: [PATCH 1/6] feat(mtproto): bind temp keys on live sockets --- native/mtproto-rs/src/auth/auth_key.rs | 62 ++- native/mtproto-rs/src/auth/bind_temp.rs | 470 ++++++++++++++++++ native/mtproto-rs/src/auth/mod.rs | 1 + native/mtproto-rs/src/client/auth.rs | 4 +- native/mtproto-rs/src/client/dispatch.rs | 6 +- native/mtproto-rs/src/client/lanes.rs | 24 +- .../mtproto-rs/src/client/media_download.rs | 82 ++- native/mtproto-rs/src/client/mod.rs | 83 +++- native/mtproto-rs/src/client/persist.rs | 22 +- native/mtproto-rs/src/client/push_rpc.rs | 10 +- native/mtproto-rs/src/client/session.rs | 4 +- native/mtproto-rs/src/client/updates.rs | 16 +- native/mtproto-rs/src/media/download.rs | 4 +- native/mtproto-rs/src/media/index.rs | 18 +- native/mtproto-rs/src/media/page_plain.rs | 20 +- native/mtproto-rs/src/rpc.rs | 3 +- native/mtproto-rs/src/scheduler.rs | 26 +- native/mtproto-rs/src/updates/updates_rpc.rs | 5 +- native/mtproto-rs/src/upload/upload_rpc.rs | 2 +- native/mtproto-rs/src/vpx.rs | 25 +- .../tests/unit/client_mgr_channel_tests.rs | 3 + .../tests/unit/client_mgr_lane_tests.rs | 61 ++- .../tests/unit/client_mgr_logout_tests.rs | 4 + .../mtproto-rs/tests/unit/media_page_tests.rs | 103 ++-- .../tests/unit/media_thumbs_tests.rs | 4 +- .../mtproto-rs/tests/unit/scheduler_tests.rs | 10 + 26 files changed, 897 insertions(+), 175 deletions(-) create mode 100644 native/mtproto-rs/src/auth/bind_temp.rs diff --git a/native/mtproto-rs/src/auth/auth_key.rs b/native/mtproto-rs/src/auth/auth_key.rs index 7b3bc8896..b7978f68f 100644 --- a/native/mtproto-rs/src/auth/auth_key.rs +++ b/native/mtproto-rs/src/auth/auth_key.rs @@ -3,8 +3,8 @@ use bnum::types::{I128, I256}; use tellers_mtproto::codec::{Boxed, BoxedDecode, Decoder, Encoder, Limits}; use tellers_mtproto::transport::{ - PQInnerDataDcConstructor, ReqDhParamsRequest, ReqPqMultiRequest, ResPq, ServerDhInnerData, - ServerDhParams, SetClientDhParamsAnswer, SetClientDhParamsRequest, + PQInnerDataDcConstructor, PQInnerDataTempDcConstructor, ReqDhParamsRequest, ReqPqMultiRequest, + ResPq, ServerDhInnerData, ServerDhParams, SetClientDhParamsAnswer, SetClientDhParamsRequest, }; use tellers_mtproto_crypto::{ aes_ige_decrypt, aes_ige_encrypt, derive_auth_key, factor_pq, fill_random, initial_server_salt, @@ -128,6 +128,7 @@ struct AuthAdapter { q: Option>, g_b: Option>, auth_key: Option>, + expires_in: Option, } impl AuthAdapter { @@ -148,6 +149,7 @@ impl AuthAdapter { q: None, g_b: None, auth_key: None, + expires_in: None, }) } } @@ -210,17 +212,29 @@ impl AuthorizationAdapter for AuthAdapter { let new_nonce = i256_from_bytes(new_nonce_bytes); self.new_nonce = Some(new_nonce.clone()); - let inner = PQInnerDataDcConstructor { - pq: res.pq.clone(), - p: p_bytes, - q: q_bytes, - nonce: self.nonce.clone(), - server_nonce: res.server_nonce.clone(), - new_nonce, - dc: inner_data_dc(self.dc_id), - }; - let inner_bytes = encode_boxed(&inner) - .map_err(|e| tellers_mtproto_engine::Error::Authorization(e.to_string()))?; + let inner_bytes = if let Some(expires_in) = self.expires_in { + encode_boxed(&PQInnerDataTempDcConstructor { + pq: res.pq.clone(), + p: p_bytes, + q: q_bytes, + nonce: self.nonce.clone(), + server_nonce: res.server_nonce.clone(), + new_nonce, + dc: inner_data_dc(self.dc_id), + expires_in, + }) + } else { + encode_boxed(&PQInnerDataDcConstructor { + pq: res.pq.clone(), + p: p_bytes, + q: q_bytes, + nonce: self.nonce.clone(), + server_nonce: res.server_nonce.clone(), + new_nonce, + dc: inner_data_dc(self.dc_id), + }) + } + .map_err(|e| tellers_mtproto_engine::Error::Authorization(e.to_string()))?; let key = parse_rsa_public_key(&self.rsa_pem) .map_err(|e| tellers_mtproto_engine::Error::Authorization(e.to_string()))?; @@ -427,14 +441,34 @@ pub fn create_auth_key( create_auth_key_with_pem(conn, framing, snapshot, rsa_pem()) } +pub fn create_temp_auth_key( + conn: &mut dyn Connection, + framing: &mut PaddedIntermediate, + snapshot: &mut Snapshot, + expires_in: i32, +) -> Result<(), MtprotoError> { + let mut adapter = AuthAdapter::with_pem(snapshot.dc_id, rsa_pem().to_string())?; + adapter.expires_in = Some(expires_in); + run_exchange(adapter, conn, framing, snapshot) +} + pub fn create_auth_key_with_pem( conn: &mut dyn Connection, framing: &mut PaddedIntermediate, snapshot: &mut Snapshot, pem: &str, ) -> Result<(), MtprotoError> { - let clock = SystemClock; let adapter = AuthAdapter::with_pem(snapshot.dc_id, pem.to_string())?; + run_exchange(adapter, conn, framing, snapshot) +} + +fn run_exchange( + adapter: AuthAdapter, + conn: &mut dyn Connection, + framing: &mut PaddedIntermediate, + snapshot: &mut Snapshot, +) -> Result<(), MtprotoError> { + let clock = SystemClock; let mut auth = Authorization::new(adapter); let body = auth diff --git a/native/mtproto-rs/src/auth/bind_temp.rs b/native/mtproto-rs/src/auth/bind_temp.rs new file mode 100644 index 000000000..bad4a1170 --- /dev/null +++ b/native/mtproto-rs/src/auth/bind_temp.rs @@ -0,0 +1,470 @@ +//! Bind a temporary auth key so an extra home-DC session does not reuse the +//! permanent key. https://core.telegram.org/api/pfs +//! https://core.telegram.org/method/auth.bindTempAuthKey +//! https://core.telegram.org/mtproto/description_v1 + +use tellers_mtproto::codec::{Boxed, Encoder}; +use tellers_mtproto::transport::BindAuthKeyInnerConstructor; +use tellers_mtproto_crypto::{aes_ige_decrypt, aes_ige_encrypt, auth_key_id, fill_random, sha1}; +use tellers_mtproto_session::{Clock, OsRandom, Snapshot}; +use tellers_mtproto_transport::{Connection, Framing, PaddedIntermediate}; + +use crate::MtprotoError; +use crate::auth::auth_key::create_temp_auth_key; + +const TEMP_EXPIRES_IN: i32 = 24 * 60 * 60; +const BOOL_TRUE: u32 = 0x9972_75b5; +const BOOL_FALSE: u32 = 0xbc79_9737; +const RPC_RESULT: u32 = 0xf35c_6d01; +const RPC_ERROR: u32 = 0x2144_ca19; +const MSG_CONTAINER: u32 = 0x73f1_f8dc; + +pub(crate) fn unix_now() -> i32 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| i32::try_from(d.as_secs()).unwrap_or(i32::MAX)) + .unwrap_or(0) +} + +/// MTProto v1 `encrypted_message` for `auth.bindTempAuthKey`. +/// `msg_key` covers the envelope before alignment padding. +pub(crate) fn encrypt_bind_message( + perm_key: &[u8], + msg_id: i64, + inner: &BindAuthKeyInnerConstructor, + random_prefix: &[u8; 16], + padding: &[u8], +) -> Result, MtprotoError> { + if perm_key.len() != 256 { + return Err(MtprotoError::Message("permanent key required".into())); + } + let body = encode_inner(inner)?; + if body.len() != 40 { + return Err(MtprotoError::Message("bind inner length".into())); + } + let mut plain = Vec::with_capacity(72 + padding.len()); + plain.extend_from_slice(random_prefix); + plain.extend_from_slice(&msg_id.to_le_bytes()); + plain.extend_from_slice(&0_i32.to_le_bytes()); + plain.extend_from_slice(&40_i32.to_le_bytes()); + plain.extend_from_slice(&body); + let msg_key = v1_msg_key(&plain); + if padding.len() > 15 || (plain.len() + padding.len()) % 16 != 0 { + return Err(MtprotoError::Message("bind padding".into())); + } + plain.extend_from_slice(padding); + let (aes_key, aes_iv) = v1_aes(&msg_key, perm_key); + let encrypted = aes_ige_encrypt(&plain, &aes_key, &aes_iv) + .map_err(|e| MtprotoError::Message(e.to_string()))?; + let mut out = Vec::with_capacity(24 + encrypted.len()); + out.extend_from_slice(&auth_key_id(perm_key).to_le_bytes()); + out.extend_from_slice(&msg_key); + out.extend_from_slice(&encrypted); + Ok(out) +} + +fn v1_msg_key(unpadded: &[u8]) -> [u8; 16] { + sha1(unpadded)[4..20].try_into().expect("sha1 slice") +} + +fn v1_aes(msg_key: &[u8; 16], auth_key: &[u8]) -> ([u8; 32], [u8; 32]) { + let sha_a = sha1_parts(&[msg_key, &auth_key[0..32]]); + let sha_b = sha1_parts(&[&auth_key[32..48], msg_key, &auth_key[48..64]]); + let sha_c = sha1_parts(&[&auth_key[64..96], msg_key]); + let sha_d = sha1_parts(&[msg_key, &auth_key[96..128]]); + let mut key = [0_u8; 32]; + key[..8].copy_from_slice(&sha_a[..8]); + key[8..20].copy_from_slice(&sha_b[8..20]); + key[20..].copy_from_slice(&sha_c[4..16]); + let mut iv = [0_u8; 32]; + iv[..12].copy_from_slice(&sha_a[8..20]); + iv[12..20].copy_from_slice(&sha_b[..8]); + iv[20..24].copy_from_slice(&sha_c[16..20]); + iv[24..].copy_from_slice(&sha_d[..8]); + (key, iv) +} + +fn sha1_parts(parts: &[&[u8]]) -> [u8; 20] { + let mut buf = Vec::new(); + for part in parts { + buf.extend_from_slice(part); + } + sha1(&buf) +} + +fn encode_inner(inner: &BindAuthKeyInnerConstructor) -> Result, MtprotoError> { + let mut encoder = Encoder::new(); + inner + .encode_boxed(&mut encoder) + .map_err(|e| MtprotoError::Message(e.to_string()))?; + Ok(encoder.into_bytes()) +} + +struct SystemClock; +impl Clock for SystemClock { + fn unix_micros(&self) -> i64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| i64::try_from(d.as_micros()).unwrap_or(i64::MAX)) + .unwrap_or(0) + } +} + +pub(crate) struct BoundTemp { + pub snapshot: Snapshot, + pub expires_at: i32, +} + +/// Handshake a temporary key and bind it. Leaves the permanent key untouched. +pub(crate) fn negotiate_bound_temp( + perm_key: &[u8], + dc_id: i32, + time_offset_micros: i64, +) -> Result { + if perm_key.len() != 256 { + return Err(MtprotoError::Message("permanent key required".into())); + } + let expires_at = unix_now().saturating_add(TEMP_EXPIRES_IN); + let mut snapshot = + Snapshot::new(dc_id, &mut OsRandom).map_err(|e| MtprotoError::Message(e.to_string()))?; + snapshot.time_offset_micros = time_offset_micros; + let mut conn = crate::tcp::connect_obfuscated_dc(crate::rpc::dc_endpoints(dc_id)) + .map_err(|e| MtprotoError::Message(e.to_string()))?; + let mut framing = PaddedIntermediate::default(); + create_temp_auth_key(&mut conn, &mut framing, &mut snapshot, TEMP_EXPIRES_IN)?; + let temp_key = snapshot + .auth_key + .clone() + .ok_or_else(|| MtprotoError::Message("temp key missing".into()))?; + if temp_key == perm_key { + return Err(MtprotoError::Message( + "temp key matched permanent key".into(), + )); + } + let clock = SystemClock; + let msg_id = snapshot + .next_message_id(&clock) + .map_err(|e| MtprotoError::Message(e.to_string()))?; + let mut nonce_bytes = [0_u8; 8]; + fill_random(&mut nonce_bytes).map_err(|e| MtprotoError::Message(e.to_string()))?; + let mut prefix = [0_u8; 16]; + fill_random(&mut prefix).map_err(|e| MtprotoError::Message(e.to_string()))?; + let inner = BindAuthKeyInnerConstructor { + nonce: i64::from_le_bytes(nonce_bytes), + temp_auth_key_id: auth_key_id(&temp_key), + perm_auth_key_id: auth_key_id(perm_key), + temp_session_id: snapshot.session_id, + expires_at, + }; + let encrypted = encrypt_bind_message(perm_key, msg_id, &inner, &prefix, &[0_u8; 8])?; + let request = tellers_mtproto::latest::api::AuthBindTempAuthKeyRequest { + perm_auth_key_id: inner.perm_auth_key_id, + nonce: inner.nonce, + expires_at, + encrypted_message: encrypted, + }; + let body = crate::rpc::encode_boxed_bytes(&request)?; + let seq = snapshot + .next_sequence(true) + .map_err(|e| MtprotoError::Message(e.to_string()))?; + send_v2( + &mut conn, + &mut framing, + &temp_key, + &mut snapshot, + msg_id, + seq, + &body, + )?; + read_bind_result( + &mut conn, + &mut framing, + &temp_key, + snapshot.session_id, + msg_id, + )?; + drop(conn); + Ok(BoundTemp { + snapshot, + expires_at, + }) +} + +fn send_v2( + conn: &mut dyn Connection, + framing: &mut PaddedIntermediate, + auth_key: &[u8], + snapshot: &mut Snapshot, + message_id: i64, + sequence: i32, + body: &[u8], +) -> Result<(), MtprotoError> { + let pad_len = v2_pad_len(body.len()); + let mut padding = vec![0_u8; pad_len]; + fill_random(&mut padding).map_err(|e| MtprotoError::Message(e.to_string()))?; + let packet = tellers_mtproto_crypto::encrypt_message( + auth_key, + tellers_mtproto_crypto::MessageToEncrypt { + server_salt: snapshot.server_salt, + session_id: snapshot.session_id, + message_id, + sequence, + body, + padding: &padding, + }, + tellers_mtproto_crypto::Direction::ClientToServer, + ) + .map_err(|e| MtprotoError::Message(e.to_string()))?; + let framed = framing + .encode(&packet) + .map_err(|e| MtprotoError::Message(e.to_string()))?; + conn.send(&framed) + .map_err(|e| MtprotoError::Message(e.to_string())) +} + +fn v2_pad_len(body_len: usize) -> usize { + let mut pad = 12; + while (32 + body_len + pad) % 16 != 0 { + pad += 1; + } + pad +} + +fn read_bind_result( + conn: &mut dyn Connection, + framing: &mut PaddedIntermediate, + temp_key: &[u8], + session_id: i64, + msg_id: i64, +) -> Result<(), MtprotoError> { + let mut buf = vec![0_u8; 64 * 1024]; + let mut input = Vec::new(); + for _ in 0..8 { + loop { + match framing.decode(&input) { + Ok(packet) => { + let consumed = packet.consumed; + let payload = crate::rpc::trim_padded_mtproto_packet(&packet.payload); + input.drain(..consumed.min(input.len())); + if let Some(control) = tellers_mtproto_transport::payload_control(payload) { + let _ = control; + break; + } + let message = tellers_mtproto_crypto::decrypt_message( + temp_key, + payload, + tellers_mtproto_crypto::Direction::ServerToClient, + session_id, + 1024 * 1024, + ) + .map_err(|e| MtprotoError::Message(e.to_string()))?; + if bind_accepted(&message.body, msg_id)? { + return Ok(()); + } + break; + } + Err(tellers_mtproto_transport::Error::Incomplete { .. }) => { + let n = conn + .receive(&mut buf) + .map_err(|e| MtprotoError::Message(e.to_string()))?; + if n == 0 { + return Err(MtprotoError::Message("temp bind connection closed".into())); + } + input.extend_from_slice(&buf[..n]); + } + Err(e) => return Err(MtprotoError::Message(e.to_string())), + } + } + } + Err(MtprotoError::Message("temp bind had no result".into())) +} + +fn bind_accepted(body: &[u8], msg_id: i64) -> Result { + if body.len() < 4 { + return Ok(false); + } + let ctor = u32::from_le_bytes(body[..4].try_into().expect("ctor")); + if ctor == MSG_CONTAINER { + return container_accepted(body, msg_id); + } + if ctor == RPC_RESULT { + return rpc_result_accepted(body, msg_id); + } + Ok(false) +} + +fn container_accepted(body: &[u8], msg_id: i64) -> Result { + if body.len() < 8 { + return Ok(false); + } + let count = u32::from_le_bytes(body[4..8].try_into().expect("count")) as usize; + let mut offset = 8; + for _ in 0..count { + if offset + 16 > body.len() { + return Ok(false); + } + offset += 8 + 4; + let len = u32::from_le_bytes(body[offset..offset + 4].try_into().expect("len")) as usize; + offset += 4; + let end = offset.saturating_add(len); + if end > body.len() { + return Ok(false); + } + if bind_accepted(&body[offset..end], msg_id)? { + return Ok(true); + } + offset = end; + } + Ok(false) +} + +fn rpc_result_accepted(body: &[u8], msg_id: i64) -> Result { + if body.len() < 16 { + return Ok(false); + } + let req = i64::from_le_bytes(body[4..12].try_into().expect("req")); + if req != msg_id { + return Ok(false); + } + let inner = u32::from_le_bytes(body[12..16].try_into().expect("inner")); + if inner == BOOL_TRUE { + return Ok(true); + } + if inner == BOOL_FALSE { + return Err(MtprotoError::Message("temp bind rejected".into())); + } + if inner == RPC_ERROR { + let code = if body.len() >= 20 { + i32::from_le_bytes(body[16..20].try_into().expect("code")) + } else { + 0 + }; + return Err(MtprotoError::Message(format!("temp bind {code}"))); + } + Err(MtprotoError::Message("temp bind result".into())) +} + +fn promote_main(client: &crate::client::Client, perm: &[u8], dc: i32, offset: i64) -> bool { + let now = unix_now(); + { + let main = client.main.lock(); + if main.temp_expires_at > now && main.snapshot.auth_key.as_deref() != Some(perm) { + return true; + } + } + let Ok(session) = negotiate_bound_temp(perm, dc, offset) else { + eprintln!("monogram.api main temp bind failed"); + return false; + }; + let mut main = client.main.lock(); + main.snapshot = session.snapshot; + main.temp_expires_at = session.expires_at; + main.transport = None; + main.pending_push = Default::default(); + main.last_difference = None; + true +} + +fn bind_lane( + lane: &crate::client::lanes::Lane, + perm: &[u8], + dc: i32, + offset: i64, + now: i32, +) -> bool { + let Some(mut io) = lane.io.try_lock() else { + return false; + }; + if io.temp_expires_at > now + && io.snapshot.auth_key.as_deref() != Some(perm) + && io.snapshot.auth_key.as_ref().map(Vec::len) == Some(256) + { + return true; + } + match negotiate_bound_temp(perm, dc, offset) { + Ok(session) => { + io.snapshot = session.snapshot; + io.temp_expires_at = session.expires_at; + io.transport = None; + io.pending_push = Default::default(); + io.last_difference = None; + true + } + Err(_) => { + eprintln!("monogram.api temp session bind failed"); + false + } + } +} + +/// After `help.getConfig`, bind telegram extra sockets. +/// Read lanes stay within `tmp_sessions`. Upload lanes 1..4 are file sockets. +/// Failure leaves the single main connection and upload lane 0 in place. +pub(crate) fn prepare_extra_main_sessions(handle: u64) { + let Ok(client) = crate::client::get_client(handle) else { + return; + }; + let (perm, dc, offset) = { + let data = client.data.lock(); + if data.user_id.is_none() || data.session_dead { + crate::scheduler::set_bound_extra_sessions(0); + return; + } + let Some(key) = data.home_auth_key.clone() else { + crate::scheduler::set_bound_extra_sessions(0); + return; + }; + (key, data.home_dc, data.home_time_offset) + }; + let now = unix_now(); + let _ = promote_main(&client, &perm, dc, offset); + for lane in &client.upload { + let _ = bind_lane(lane, &perm, dc, offset, now); + } + let allowance = crate::scheduler::main_session_allowance(); + if allowance <= 1 { + crate::scheduler::set_bound_extra_sessions(0); + return; + } + let want = usize::try_from(allowance.saturating_sub(1)).unwrap_or(0); + let want = want.min(crate::scheduler::READ_LANES); + let mut bound = 0_usize; + for lane in client.rpc.iter().take(want) { + if bind_lane(lane, &perm, dc, offset, now) { + bound += 1; + } + } + crate::scheduler::set_bound_extra_sessions(bound); +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn bind_message_uses_v1_msg_key_without_padding() { + let perm = vec![3_u8; 256]; + let inner = BindAuthKeyInnerConstructor { + nonce: 7, + temp_auth_key_id: 8, + perm_auth_key_id: auth_key_id(&perm), + temp_session_id: 9, + expires_at: 1_700_000_000, + }; + let prefix = [1_u8; 16]; + let packet = encrypt_bind_message(&perm, 100, &inner, &prefix, &[2_u8; 8]).unwrap(); + assert_eq!( + i64::from_le_bytes(packet[..8].try_into().unwrap()), + auth_key_id(&perm) + ); + let msg_key: [u8; 16] = packet[8..24].try_into().unwrap(); + let (aes_key, aes_iv) = v1_aes(&msg_key, &perm); + let plain = aes_ige_decrypt(&packet[24..], &aes_key, &aes_iv).unwrap(); + assert_eq!(&plain[..16], &prefix); + assert_eq!(i64::from_le_bytes(plain[16..24].try_into().unwrap()), 100); + assert_eq!(i32::from_le_bytes(plain[24..28].try_into().unwrap()), 0); + assert_eq!(i32::from_le_bytes(plain[28..32].try_into().unwrap()), 40); + assert_eq!(v1_msg_key(&plain[..plain.len() - 8]), msg_key); + assert_eq!(&plain[plain.len() - 8..], &[2_u8; 8]); + } +} diff --git a/native/mtproto-rs/src/auth/mod.rs b/native/mtproto-rs/src/auth/mod.rs index 9567c0349..3bd5d4540 100644 --- a/native/mtproto-rs/src/auth/mod.rs +++ b/native/mtproto-rs/src/auth/mod.rs @@ -1,3 +1,4 @@ pub(crate) mod auth_key; pub(crate) mod auth_rpc; +pub(crate) mod bind_temp; pub(crate) mod srp; diff --git a/native/mtproto-rs/src/client/auth.rs b/native/mtproto-rs/src/client/auth.rs index bd20a8e99..db04c5484 100644 --- a/native/mtproto-rs/src/client/auth.rs +++ b/native/mtproto-rs/src/client/auth.rs @@ -53,7 +53,9 @@ pub fn connect(handle: u64) -> Result<(), MtprotoError> { } recover_imported_user(state)?; Ok(()) - }) + })?; + crate::auth::bind_temp::prepare_extra_main_sessions(handle); + Ok(()) } fn recover_imported_user(state: &mut ClientState) -> Result<(), MtprotoError> { diff --git a/native/mtproto-rs/src/client/dispatch.rs b/native/mtproto-rs/src/client/dispatch.rs index df885c9f1..e50fff050 100644 --- a/native/mtproto-rs/src/client/dispatch.rs +++ b/native/mtproto-rs/src/client/dispatch.rs @@ -590,8 +590,10 @@ pub(crate) fn save_items_with( let _policy = crate::transfer_policy::bind(Arc::clone(&client.policy)); let home = client.main.lock().snapshot.clone(); crate::media::check_transfer_flood(home.dc_id, crate::media::TransferClass::Upload)?; - let mut io = client.upload.lock(); - if sync_upload_snapshot(&mut io.snapshot, &home) { + let mut io = super::lanes::lock_upload_lane(client); + if io.temp_expires_at <= crate::auth::bind_temp::unix_now() + && sync_upload_snapshot(&mut io.snapshot, &home) + { io.transport = None; } let mut transport = io.transport.take(); diff --git a/native/mtproto-rs/src/client/lanes.rs b/native/mtproto-rs/src/client/lanes.rs index 07dc5e170..85aa188e1 100644 --- a/native/mtproto-rs/src/client/lanes.rs +++ b/native/mtproto-rs/src/client/lanes.rs @@ -29,6 +29,8 @@ pub(crate) struct SessionIo { pub(crate) last_difference: Option, pub(crate) snapshot: Snapshot, pub(crate) transport: Option, + /// Unix time when this lane's temporary auth key expires. Zero means none. + pub(crate) temp_expires_at: i32, } /// A session plus the gate that decides which request class gets it next. @@ -175,13 +177,28 @@ pub(crate) fn close_idle_file_transports(client: &Client, now: std::time::Instan for lane in &client.media { close(&lane.io); } - close(&client.upload); + for lane in &client.upload { + close(&lane.io); + } if !pending { client.media_open.store(1, Ordering::Release); } pending } +/// Prefer a free temporary-key upload socket. Lane 0 is the permanent-key fallback. +pub(crate) fn lock_upload_lane(client: &Client) -> parking_lot::MutexGuard<'_, SessionIo> { + let now = crate::auth::bind_temp::unix_now(); + for lane in client.upload.iter().skip(1) { + if let Some(io) = lane.io.try_lock() { + if io.temp_expires_at > now { + return io; + } + } + } + client.upload[0].io.lock() +} + pub(crate) fn schedule_file_cleanup(client: &Arc) { if client.file_cleanup_running.swap(true, Ordering::AcqRel) { return; @@ -248,9 +265,8 @@ pub(crate) fn flush_idle_main_acks( false } -/// Second media-DC TCP if a lane is free. The held first lease makes that gate -/// busy, so this never returns the same lane. Home-DC downloads must not call this. -#[allow(dead_code)] +/// Second file TCP if a lane is free. The held first lease makes that gate +/// busy, so this never returns the same lane. Caller must bind a separate temp key. pub(crate) fn try_lock_second_media_lane(client: &Client) -> Option> { let class = scheduler::current_class(); if scheduler::family(class) != scheduler::LaneFamily::Media { diff --git a/native/mtproto-rs/src/client/media_download.rs b/native/mtproto-rs/src/client/media_download.rs index d377469be..90476ee51 100644 --- a/native/mtproto-rs/src/client/media_download.rs +++ b/native/mtproto-rs/src/client/media_download.rs @@ -14,6 +14,10 @@ use crate::tcp; use super::*; +fn home_lease_ok(state: &ClientState, home_session_id: i64) -> bool { + state.snapshot.session_id == home_session_id || state.perm_session_id == home_session_id +} + pub fn download_wallpaper( handle: u64, id: i64, @@ -43,7 +47,7 @@ pub fn download_wallpaper( ) { Err(error) if media::is_file_reference_error(&error) => { let refreshed = with_client_mut(handle, |state| { - if state.snapshot.session_id != home.session_id { + if !home_lease_ok(state, home.session_id) { return Err(expired_session_lease()); } let media = call_with_migrate(state, |state| { @@ -394,7 +398,7 @@ pub fn download_message_media_range( || thumb_file_id_can_refresh(kind, &media, message_id, &err) => { let refreshed = with_client_mut(handle, |state| { - if state.snapshot.session_id != home.session_id { + if !home_lease_ok(state, home.session_id) { return Err(expired_session_lease()); } crate::rpc::with_rpc_timeout_secs(20, || { @@ -786,7 +790,7 @@ pub(crate) fn download_media_range_on_home_session( |_init_first, requests, on_chunk| { media::check_transfer_flood(dc, media::TransferClass::Download)?; with_client_mut(handle, |state| { - if state.snapshot.session_id != session_id { + if !home_lease_ok(state, session_id) { return Err(expired_session_lease()); } ensure_ready(state)?; @@ -921,33 +925,63 @@ pub(crate) fn download_media_range_on_lane_dc( let _span = crate::perf::span("media_file"); let staged = media::StagedDownload::new(dest)?; let mut io = lock_media_lane(client)?; - // A file owns one session and pipelines at most eight parts on it. - let mut io2: Option> = None; if !session_lease_valid(&client.data.lock(), home.session_id) { return Err(expired_session_lease()); } - match prepare_media_lane_snapshot(&mut io.snapshot, home, target_dc) { - MediaLanePrep::NeedExport => { - io.snapshot = with_client_mut(handle, |state| { - if state.snapshot.session_id != home.session_id { - return Err(expired_session_lease()); - } - copy_authorization_to_dc(&mut state.snapshot, api_id, target_dc) - })?; - io.transport = None; - } - MediaLanePrep::Replaced => io.transport = None, - MediaLanePrep::Reuse => {} - } - if let Some(second) = io2.as_mut() { - match prepare_media_lane_snapshot(&mut second.snapshot, home, target_dc) { + let mut lane_perm = None; + let already_temp = io.temp_expires_at > crate::auth::bind_temp::unix_now() + && io.snapshot.auth_key.as_ref() != home.auth_key.as_ref(); + if !already_temp { + match prepare_media_lane_snapshot(&mut io.snapshot, home, target_dc) { MediaLanePrep::NeedExport => { - second.snapshot = io.snapshot.clone(); - second.transport = None; + io.snapshot = with_client_mut(handle, |state| { + if !home_lease_ok(state, home.session_id) { + return Err(expired_session_lease()); + } + copy_authorization_to_dc(&mut state.snapshot, api_id, target_dc) + })?; + io.transport = None; } - MediaLanePrep::Replaced => second.transport = None, + MediaLanePrep::Replaced => io.transport = None, MediaLanePrep::Reuse => {} } + if let Some(key) = io.snapshot.auth_key.clone() { + lane_perm = Some(key.clone()); + let offset = io.snapshot.time_offset_micros; + let dc = io.snapshot.dc_id; + if let Ok(bound) = crate::auth::bind_temp::negotiate_bound_temp(&key, dc, offset) { + io.snapshot = bound.snapshot; + io.temp_expires_at = bound.expires_at; + io.transport = None; + } + } + } + if target_dc == home.dc_id && io.snapshot.auth_key.as_ref() == home.auth_key.as_ref() { + drop(io); + return download_media_range_on_home_session( + handle, + client, + home.session_id, + api_id, + media, + dest, + offset, + ); + } + let mut io2 = try_lock_second_media_lane(client); + let second_bound = if let (Some(second), Some(key)) = (io2.as_mut(), lane_perm.as_deref()) { + crate::auth::bind_temp::negotiate_bound_temp(key, target_dc, home.time_offset_micros) + .ok() + .map(|bound| { + second.snapshot = bound.snapshot; + second.temp_expires_at = bound.expires_at; + second.transport = None; + }) + } else { + None + }; + if second_bound.is_none() { + io2 = None; } let mut media_snap = io.snapshot.clone(); let mut slot = io.transport.take(); @@ -1016,7 +1050,7 @@ pub(crate) fn download_media_range_on_lane_dc( && is_unrecoverable_session(err) { let _ = with_client_mut(handle, |state| { - if state.snapshot.session_id != home.session_id { + if !home_lease_ok(state, home.session_id) { return Err(expired_session_lease()); } mark_session_dead(state, err); diff --git a/native/mtproto-rs/src/client/mod.rs b/native/mtproto-rs/src/client/mod.rs index f379bd3c3..b1074000f 100644 --- a/native/mtproto-rs/src/client/mod.rs +++ b/native/mtproto-rs/src/client/mod.rs @@ -7,7 +7,7 @@ mod auth; mod dispatch; mod extras; pub(crate) mod extras_rpc; -mod lanes; +pub(crate) mod lanes; mod media_download; mod persist; pub(crate) mod push_rpc; @@ -127,8 +127,7 @@ pub(crate) struct Client { pub(crate) rpc: [Lane; scheduler::READ_LANES], /// File RPCs on separate sessions. pub(crate) media: crate::SmallVec<[Lane; scheduler::MAX_MEDIA_LANES]>, - /// One upload session. Not the main sender and not a download lane. - pub(crate) upload: Mutex, + pub(crate) upload: [Lane; scheduler::UPLOAD_LANES], pub(crate) file_cleanup_running: AtomicBool, pub(crate) media_open: std::sync::atomic::AtomicUsize, pub(crate) media_gate: scheduler::LaneGate, @@ -159,6 +158,10 @@ pub(crate) struct ClientState { pub(crate) session_dead_reason: Option, pub(crate) test_dc: bool, pub(crate) last_inline: Option, + /// Set while the live snapshot uses a PFS temp key. Disk writes keep this key. + pub(crate) perm_auth_key: Option>, + pub(crate) perm_salt: i64, + pub(crate) perm_session_id: i64, } pub(crate) fn prefer_newer_cursor( @@ -268,7 +271,6 @@ pub fn create_client(api_id: i32, api_hash: String, session_path: String) -> u64 }; crate::rpc::set_use_test_dc(test_dc); let media_snapshot = fork_session(&snapshot); - let upload_snapshot = fork_session(&snapshot); let rpc_snapshot = fork_session(&snapshot); let rpc_snapshot_b = fork_session(&snapshot); let media_snapshots: Vec = (0..scheduler::MAX_MEDIA_LANES) @@ -280,6 +282,15 @@ pub fn create_client(api_id: i32, api_hash: String, session_path: String) -> u64 } }) .collect(); + let upload_lanes = core::array::from_fn(|_| { + Lane::new(SessionIo { + pending_push: Default::default(), + last_difference: None, + snapshot: fork_session(&snapshot), + transport: None, + temp_expires_at: 0, + }) + }); let handle = NEXT_HANDLE.fetch_add(1, Ordering::Relaxed); CLIENTS.lock().insert( handle, @@ -326,6 +337,7 @@ pub fn create_client(api_id: i32, api_hash: String, session_path: String) -> u64 last_difference: None, snapshot, transport: None, + temp_expires_at: 0, }), rpc: [ Lane::new(SessionIo { @@ -333,12 +345,14 @@ pub fn create_client(api_id: i32, api_hash: String, session_path: String) -> u64 last_difference: None, snapshot: rpc_snapshot, transport: None, + temp_expires_at: 0, }), Lane::new(SessionIo { pending_push: Default::default(), last_difference: None, snapshot: rpc_snapshot_b, transport: None, + temp_expires_at: 0, }), ], media: media_snapshots @@ -349,15 +363,11 @@ pub fn create_client(api_id: i32, api_hash: String, session_path: String) -> u64 pending_push: Default::default(), last_difference: None, transport: None, + temp_expires_at: 0, }) }) .collect(), - upload: Mutex::new(SessionIo { - pending_push: Default::default(), - last_difference: None, - snapshot: upload_snapshot, - transport: None, - }), + upload: upload_lanes, file_cleanup_running: AtomicBool::new(false), media_open: std::sync::atomic::AtomicUsize::new(1), media_gate: scheduler::LaneGate::new(), @@ -515,6 +525,11 @@ pub(crate) fn with_client_mut( session_dead_reason: d.session_dead_reason.clone(), test_dc: d.test_dc, last_inline: d.last_inline.clone(), + perm_auth_key: (io.temp_expires_at > crate::auth::bind_temp::unix_now()) + .then(|| d.home_auth_key.clone()) + .flatten(), + perm_salt: d.home_salt, + perm_session_id: d.home_session_id, } }; crate::rpc::set_use_test_dc(state.test_dc); @@ -591,10 +606,15 @@ pub(crate) fn with_client_mut( None }; d.home_dc = state.snapshot.dc_id; - d.home_session_id = state.snapshot.session_id; - d.home_auth_key = state.snapshot.auth_key.clone(); - d.home_salt = state.snapshot.server_salt; d.home_time_offset = state.snapshot.time_offset_micros; + if state.session_dead || state.snapshot.auth_key != before_snapshot.auth_key { + io.temp_expires_at = 0; + } + if io.temp_expires_at <= crate::auth::bind_temp::unix_now() { + d.home_session_id = state.snapshot.session_id; + d.home_auth_key = state.snapshot.auth_key.clone(); + d.home_salt = state.snapshot.server_salt; + } d.test_dc = state.test_dc; d.last_inline = state.last_inline; persist_needed = result.is_ok() @@ -610,10 +630,13 @@ pub(crate) fn with_client_mut( if identity_changed { for lane in client.rpc.iter() { if let Some(mut rpc) = lane.io.try_lock() { - rpc.snapshot = fork_session(&io.snapshot); + rpc.temp_expires_at = 0; rpc.transport = None; rpc.pending_push = Default::default(); rpc.last_difference = None; + if io.temp_expires_at <= crate::auth::bind_temp::unix_now() { + rpc.snapshot = fork_session(&io.snapshot); + } } } } @@ -658,13 +681,20 @@ pub(crate) fn with_read_lane( // Reads wait for a read lane; the `with_client_mut` paths below are authorization // repairs, not contention fallbacks. let mut io = acquire_read_lane(&client)?; - match prepare_media_lane_snapshot(&mut io.snapshot, &home, home.dc_id) { - MediaLanePrep::NeedExport => { - drop(io); - return with_client_mut(handle, f); - } - MediaLanePrep::Replaced => io.transport = None, - MediaLanePrep::Reuse => {} + let now = crate::auth::bind_temp::unix_now(); + if io.temp_expires_at > 0 && io.temp_expires_at <= now { + io.temp_expires_at = 0; + io.transport = None; + drop(io); + scheduler::set_bound_extra_sessions(0); + return with_client_mut(handle, f); + } + let bound_temp = io.temp_expires_at > now + && io.snapshot.auth_key.as_ref() != home.auth_key.as_ref() + && io.snapshot.auth_key.as_ref().map(Vec::len) == Some(256); + if !bound_temp { + drop(io); + return with_client_mut(handle, f); } let home_session_id = home.session_id; let home_dc = home.dc_id; @@ -693,6 +723,11 @@ pub(crate) fn with_read_lane( session_dead_reason: d.session_dead_reason.clone(), test_dc: d.test_dc, last_inline: d.last_inline.clone(), + perm_auth_key: (io.temp_expires_at > crate::auth::bind_temp::unix_now()) + .then(|| d.home_auth_key.clone()) + .flatten(), + perm_salt: d.home_salt, + perm_session_id: d.home_session_id, }) } }; @@ -726,8 +761,12 @@ pub(crate) fn with_read_lane( if let Some(metadata) = crate::rpc::take_new_session_metadata() { state.new_session = Some(metadata); } - let diverged = state.snapshot.dc_id != home_dc || state.snapshot.auth_key != home_auth; + let still_temp = io.temp_expires_at > crate::auth::bind_temp::unix_now() + && state.snapshot.auth_key.as_ref() != home_auth.as_ref(); + let diverged = + state.snapshot.dc_id != home_dc || (state.snapshot.auth_key != home_auth && !still_temp); if diverged { + io.temp_expires_at = 0; io.snapshot = fork_session(&home); io.transport = None; drop(io); diff --git a/native/mtproto-rs/src/client/persist.rs b/native/mtproto-rs/src/client/persist.rs index bf70a2a6c..851054916 100644 --- a/native/mtproto-rs/src/client/persist.rs +++ b/native/mtproto-rs/src/client/persist.rs @@ -51,13 +51,23 @@ thread_local! { pub(crate) static EXTRA_READ_LANE: Cell = const { Cell::new(false) }; } +pub(crate) fn durable_snapshot(state: &ClientState) -> tellers_mtproto_session::Snapshot { + let mut snapshot = state.snapshot.clone(); + if let Some(key) = &state.perm_auth_key { + snapshot.auth_key = Some(key.clone()); + snapshot.server_salt = state.perm_salt; + snapshot.session_id = state.perm_session_id; + } + snapshot +} + pub(crate) fn persist(state: &ClientState) -> Result<(), MtprotoError> { // Extra-lane snapshot_id must never replace the main session on disk. if EXTRA_READ_LANE.with(Cell::get) { return Ok(()); } let session = ClientSession { - snapshot: state.snapshot.clone(), + snapshot: durable_snapshot(state), user_id: state.user_id, peers: state.peers.clone(), updates: state.updates.clone(), @@ -97,8 +107,16 @@ pub(crate) fn persist_updates_data(client: &Client, session_id: i64) -> Result<( if !session_lease_valid(&d, session_id) { return Err(expired_session_lease()); } + let mut snapshot = main.snapshot.clone(); + if main.temp_expires_at > crate::auth::bind_temp::unix_now() { + if let Some(key) = d.home_auth_key.clone() { + snapshot.auth_key = Some(key); + snapshot.server_salt = d.home_salt; + snapshot.session_id = d.home_session_id; + } + } let session = ClientSession { - snapshot: main.snapshot.clone(), + snapshot, user_id: d.user_id, peers: d.peers.clone(), updates: d.updates.clone(), diff --git a/native/mtproto-rs/src/client/push_rpc.rs b/native/mtproto-rs/src/client/push_rpc.rs index 68e080c7e..73598e34d 100644 --- a/native/mtproto-rs/src/client/push_rpc.rs +++ b/native/mtproto-rs/src/client/push_rpc.rs @@ -101,7 +101,9 @@ fn settings_dto(settings: &PeerNotifySettings) -> NotifySettingsDto { return NotifySettingsDto::default(); }; NotifySettingsDto { - show_previews: body.show_previews.as_ref() + show_previews: body + .show_previews + .as_ref() .map(|flag| matches!(flag.as_ref(), Bool::BoolTrue(_))), silent: as_bool(&body.silent), mute_until: body.mute_until.unwrap_or(0), @@ -125,7 +127,11 @@ pub(crate) fn exception_from_update(update: &Update) -> Option ("chats", 0_i64, None), NotifyPeer::NotifyBroadcasts(_) => ("broadcasts", 0_i64, None), NotifyPeer::NotifyPeer(peer) => ("peer", peer_chat_id(peer.peer.as_ref()), None), - NotifyPeer::NotifyForumTopic(topic) => ("peer", peer_chat_id(topic.peer.as_ref()), Some(topic.top_msg_id)), + NotifyPeer::NotifyForumTopic(topic) => ( + "peer", + peer_chat_id(topic.peer.as_ref()), + Some(topic.top_msg_id), + ), _ => return None, }; Some(NotifyExceptionDto { diff --git a/native/mtproto-rs/src/client/session.rs b/native/mtproto-rs/src/client/session.rs index c6cf4a74d..d6a8c1fd7 100644 --- a/native/mtproto-rs/src/client/session.rs +++ b/native/mtproto-rs/src/client/session.rs @@ -274,7 +274,9 @@ pub(crate) fn drop_all_transports(handle: u64) { for lane in client.media.iter() { lane.io.lock().transport = None; } - client.upload.lock().transport = None; + for lane in &client.upload { + lane.io.lock().transport = None; + } } } diff --git a/native/mtproto-rs/src/client/updates.rs b/native/mtproto-rs/src/client/updates.rs index 14802068a..f83029131 100644 --- a/native/mtproto-rs/src/client/updates.rs +++ b/native/mtproto-rs/src/client/updates.rs @@ -147,13 +147,15 @@ pub fn drain_updates(handle: u64) -> Result, MtprotoError> { d.lazy_channel_updates && !d.lazy_channel_recovery.is_empty() && d.lazy_retry_at <= now, ) }; - apply_home_auth( - &mut io.snapshot, - home_dc, - home_auth.clone(), - home_salt, - home_off, - ); + if io.temp_expires_at <= crate::auth::bind_temp::unix_now() { + apply_home_auth( + &mut io.snapshot, + home_dc, + home_auth.clone(), + home_salt, + home_off, + ); + } let mut needs_difference = io.transport.is_none() || io .last_difference diff --git a/native/mtproto-rs/src/media/download.rs b/native/mtproto-rs/src/media/download.rs index 6bb94b4dd..d87eca10f 100644 --- a/native/mtproto-rs/src/media/download.rs +++ b/native/mtproto-rs/src/media/download.rs @@ -18,7 +18,8 @@ pub(crate) const DEFAULT_CHUNK: i32 = 128 * 1024; pub(crate) const CHUNK: i32 = DEFAULT_CHUNK; pub(crate) const STREAM_WINDOW_CHUNKS: usize = 4; pub(crate) const FAST_CHUNK: i32 = 512 * 1024; -pub const DOWNLOAD_SESSION_IDLE: std::time::Duration = std::time::Duration::from_secs(2); +/// Telegram keeps foreground file sockets up and pings about every 3 minutes. +pub const DOWNLOAD_SESSION_IDLE: std::time::Duration = std::time::Duration::from_secs(180); /// https://core.telegram.org/api/files #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -763,7 +764,6 @@ pub(crate) fn download_media_range_batched_streaming_capped( if out.seek(SeekFrom::Start(part_offset as u64)).is_ok() && out.write_all(bytes).is_ok() { - let _ = out.flush(); written_end = written_end.max(part_offset + bytes.len() as i64); if counted_offsets.insert(part_offset) { downloaded_bytes += bytes.len() as i64; diff --git a/native/mtproto-rs/src/media/index.rs b/native/mtproto-rs/src/media/index.rs index 6ba324638..58f6fef8a 100644 --- a/native/mtproto-rs/src/media/index.rs +++ b/native/mtproto-rs/src/media/index.rs @@ -330,13 +330,17 @@ pub fn media_ref_peer_photo( has_video: bool, stripped_thumb: Option<&[u8]>, ) -> MediaRef { - let (thumb_cache_key, thumb_location) = match stripped_thumb.and_then(crate::stripped_jpeg::expand_stripped_jpeg) { - Some(bytes) => ( - Some(format!("{}:thumb", cache_key)), - Some(MediaLocation::Inline { bytes, extension: "jpg".into() }), - ), - None => (None, None), - }; + let (thumb_cache_key, thumb_location) = + match stripped_thumb.and_then(crate::stripped_jpeg::expand_stripped_jpeg) { + Some(bytes) => ( + Some(format!("{}:thumb", cache_key)), + Some(MediaLocation::Inline { + bytes, + extension: "jpg".into(), + }), + ), + None => (None, None), + }; MediaRef { file_size: None, kind: if has_video { diff --git a/native/mtproto-rs/src/media/page_plain.rs b/native/mtproto-rs/src/media/page_plain.rs index 66d491d9c..59f698a9e 100644 --- a/native/mtproto-rs/src/media/page_plain.rs +++ b/native/mtproto-rs/src/media/page_plain.rs @@ -390,12 +390,7 @@ fn push_table( let tellers_mtproto::latest::api::PageTableCell::PageTableCell(cell) = cell else { continue; }; - cells.push( - cell.text - .as_deref() - .map(rich_plain) - .unwrap_or_default(), - ); + cells.push(cell.text.as_deref().map(rich_plain).unwrap_or_default()); } if !cells.is_empty() { grid.push(cells); @@ -586,14 +581,13 @@ fn push_rich(text: &RichText, out: &mut FormattedText) { mark(out, start, "url", None); } RichText::TextCustomEmoji(t) => { - let alt = if t.alt.is_empty() { "\u{FFFC}" } else { t.alt.as_str() }; + let alt = if t.alt.is_empty() { + "\u{FFFC}" + } else { + t.alt.as_str() + }; out.text.push_str(alt); - mark( - out, - start, - "custom_emoji", - Some(t.document_id.to_string()), - ); + mark(out, start, "custom_emoji", Some(t.document_id.to_string())); } RichText::TextSubscript(t) => { push_rich(&t.text, out); diff --git a/native/mtproto-rs/src/rpc.rs b/native/mtproto-rs/src/rpc.rs index f6359f211..d2f5dc41e 100644 --- a/native/mtproto-rs/src/rpc.rs +++ b/native/mtproto-rs/src/rpc.rs @@ -29,6 +29,7 @@ pub(crate) use invoke::{ invoke_raw_with_retry_factory, ping_existing_rtt, receive_updates, }; pub(crate) use main_policy::ACK_DELAY; +pub(crate) use timeout::trim_padded_mtproto_packet; #[cfg(test)] use crate::MtprotoError; @@ -81,7 +82,7 @@ pub(crate) use timeout::{ idle_needs_liveness_probe, idle_reused_socket, is_mid_frame, keepalive_probe_failed, leftover_frame_grace, live_transport_stale, note_inbound_liveness, recv_wait_deadline, rpc_attempt_budget, rpc_timeout_message, rpc_timeout_secs, subscribed_read_deadline, - timeout_idle_needs_probe, trim_padded_mtproto_packet, + timeout_idle_needs_probe, }; #[cfg(test)] diff --git a/native/mtproto-rs/src/scheduler.rs b/native/mtproto-rs/src/scheduler.rs index 3d2ef39da..35e743e65 100644 --- a/native/mtproto-rs/src/scheduler.rs +++ b/native/mtproto-rs/src/scheduler.rs @@ -103,6 +103,8 @@ pub enum LaneFamily { /// [`extra_main_sessions`]. pub const READ_LANES: usize = 2; +pub const UPLOAD_LANES: usize = 4; + /// Most parallel main sessions we will ever open on one non-media DC. pub const MAX_MAIN_SESSIONS: i32 = 8; @@ -138,11 +140,27 @@ pub fn main_session_allowance_known() -> bool { MAIN_SESSION_ALLOWANCE.load(std::sync::atomic::Ordering::Relaxed) >= 1 } -/// Extra main sessions remain disabled until each has a bound temporary PFS key. -/// A server allowance alone does not authorize reusing the permanent key across -/// parallel main sessions: https://core.telegram.org/api/datacenter#parallel-sessions +/// Extra main sessions that currently hold a bound temporary key. +/// Zero until [`set_bound_extra_sessions`] records a successful bind. +static BOUND_EXTRA_SESSIONS: std::sync::atomic::AtomicUsize = + std::sync::atomic::AtomicUsize::new(0); + +/// Records how many extra home sessions have a bound temporary key. +/// +/// A server allowance alone does not authorize reusing the permanent key. +/// https://core.telegram.org/api/datacenter#parallel-sessions +pub fn set_bound_extra_sessions(count: usize) { + BOUND_EXTRA_SESSIONS.store(count.min(READ_LANES), std::sync::atomic::Ordering::Relaxed); +} + +/// Extra main sessions that may be used. Stays zero until a temp key is bound. pub fn extra_main_sessions() -> usize { - 0 + let bound = BOUND_EXTRA_SESSIONS.load(std::sync::atomic::Ordering::Relaxed); + if bound == 0 { + return 0; + } + let extras = usize::try_from(main_session_allowance().saturating_sub(1)).unwrap_or(0); + bound.min(extras).min(READ_LANES) } /// Home-DC read lanes the current allowance pays for. diff --git a/native/mtproto-rs/src/updates/updates_rpc.rs b/native/mtproto-rs/src/updates/updates_rpc.rs index 40d5d9bf4..2178b9033 100644 --- a/native/mtproto-rs/src/updates/updates_rpc.rs +++ b/native/mtproto-rs/src/updates/updates_rpc.rs @@ -427,7 +427,10 @@ pub(crate) fn collect_other_updates<'a>( events.push(UpdateEventDto::Ignored { kind: format!( "notify_settings:{}:{}:{}:{}", - ex.peer_kind, ex.chat_id, ex.mute_until, ex.topic_id.unwrap_or(0) + ex.peer_kind, + ex.chat_id, + ex.mute_until, + ex.topic_id.unwrap_or(0) ), }); } diff --git a/native/mtproto-rs/src/upload/upload_rpc.rs b/native/mtproto-rs/src/upload/upload_rpc.rs index 65532ae52..3ad312de9 100644 --- a/native/mtproto-rs/src/upload/upload_rpc.rs +++ b/native/mtproto-rs/src/upload/upload_rpc.rs @@ -26,7 +26,7 @@ use tellers_mtproto_session::Snapshot; use crate::api_invoke; use crate::media::MediaIndex; use crate::messages::{self, input_reply_to_thread, random_id}; -use crate::peers::{self, input_peer_from_cached, CachedPeer}; +use crate::peers::{self, CachedPeer, input_peer_from_cached}; use crate::{MessageDto, MtprotoError, UploadItemDto}; pub const FILE_PART_SMALL: usize = 32 * 1024; diff --git a/native/mtproto-rs/src/vpx.rs b/native/mtproto-rs/src/vpx.rs index 51a30ed8d..4e3081025 100644 --- a/native/mtproto-rs/src/vpx.rs +++ b/native/mtproto-rs/src/vpx.rs @@ -251,12 +251,13 @@ pub fn decode_vpx_alpha_packet( } let dst_w = out_w as usize; let dst_h = out_h as usize; - let mut alpha = vec![ - 0u8; - dst_w.checked_mul(dst_h).ok_or_else(|| MtprotoError::Message( - "vpx alpha frame is too large".into() - ))? - ]; + let mut alpha = + vec![ + 0u8; + dst_w + .checked_mul(dst_h) + .ok_or_else(|| MtprotoError::Message("vpx alpha frame is too large".into()))? + ]; for row in 0..dst_h { let src_row = row * h / dst_h; for col in 0..dst_w { @@ -381,7 +382,7 @@ fn clip(value: i32) -> u8 { #[cfg(test)] mod tests { - use super::{fitted_output, lookup_arc, VPX_DECODER_THREADS, VPX_MAX_OUTPUT}; + use super::{VPX_DECODER_THREADS, VPX_MAX_OUTPUT, fitted_output, lookup_arc}; use crate::{HashMap, HashMapExt}; use parking_lot::Mutex; use std::sync::Arc; @@ -392,7 +393,10 @@ mod tests { values.insert(1u64, Arc::new(5u32)); let map = Mutex::new(values); let value = lookup_arc(&map, 1).expect("present"); - assert!(map.try_lock().is_some(), "map lock must be free after lookup"); + assert!( + map.try_lock().is_some(), + "map lock must be free after lookup" + ); assert_eq!(*value, 5); } @@ -404,7 +408,10 @@ mod tests { #[test] fn output_stays_inside_one_megabyte() { assert_eq!(fitted_output(512, 512), (512, 512)); - assert_eq!(fitted_output(2048, 1024), (VPX_MAX_OUTPUT, VPX_MAX_OUTPUT / 2)); + assert_eq!( + fitted_output(2048, 1024), + (VPX_MAX_OUTPUT, VPX_MAX_OUTPUT / 2) + ); let (w, h) = fitted_output(2048, 2048); assert!(w <= VPX_MAX_OUTPUT && h <= VPX_MAX_OUTPUT); assert!(w as usize * h as usize * 4 <= 1024 * 1024); diff --git a/native/mtproto-rs/tests/unit/client_mgr_channel_tests.rs b/native/mtproto-rs/tests/unit/client_mgr_channel_tests.rs index cf7eb06bd..98b9f444c 100644 --- a/native/mtproto-rs/tests/unit/client_mgr_channel_tests.rs +++ b/native/mtproto-rs/tests/unit/client_mgr_channel_tests.rs @@ -79,6 +79,9 @@ fn begin_updates_preserves_the_unapplied_cursor_in_persistence() { session_dead_reason: None, test_dc: false, last_inline: None, + perm_auth_key: None, + perm_salt: 0, + perm_session_id: 0, }; begin_updates(&mut state); diff --git a/native/mtproto-rs/tests/unit/client_mgr_lane_tests.rs b/native/mtproto-rs/tests/unit/client_mgr_lane_tests.rs index 18f8b4aa5..909d35856 100644 --- a/native/mtproto-rs/tests/unit/client_mgr_lane_tests.rs +++ b/native/mtproto-rs/tests/unit/client_mgr_lane_tests.rs @@ -423,6 +423,55 @@ fn idle_work_uses_home_session_without_pfs() { assert!(!used_read_lane); } +#[test] +fn bound_temp_read_lane_does_not_reuse_the_permanent_key() { + let _allowance = allow_read_lanes(); + struct Reset; + impl Drop for Reset { + fn drop(&mut self) { + crate::scheduler::set_bound_extra_sessions(0); + } + } + let _reset = Reset; + let path = + std::env::temp_dir().join(format!("monogram-temp-read-{}.session", std::process::id())); + let handle = create_client(1, "hash".into(), path.to_string_lossy().into()); + let (home_id, home_auth) = authorize_test_client(handle); + let client = get_client(handle).expect("client"); + let temp_key = vec![9_u8; 256]; + let temp_id = { + let mut io = client.rpc[0].io.lock(); + let mut snap = Snapshot::new(io.snapshot.dc_id, &mut OsRandom).expect("temp"); + snap.auth_key = Some(temp_key.clone()); + snap.server_salt = 99; + let id = snap.session_id; + io.snapshot = snap; + io.temp_expires_at = crate::auth::bind_temp::unix_now().saturating_add(3600); + io.transport = None; + id + }; + crate::scheduler::set_bound_extra_sessions(1); + let mut seen = 0_i64; + with_read_lane(handle, |state| { + seen = state.snapshot.session_id; + assert_eq!( + state.snapshot.auth_key.as_deref(), + Some(temp_key.as_slice()) + ); + assert_ne!(state.snapshot.session_id, home_id); + assert!(crate::api_invoke::invoking_without_updates()); + Ok(()) + }) + .expect("temp read"); + assert_eq!(seen, temp_id); + let data = client.data.lock(); + assert_eq!(data.home_auth_key.as_deref(), Some(home_auth.as_slice())); + assert_eq!(data.home_session_id, home_id); + drop(data); + destroy_client(handle); + let _ = std::fs::remove_file(path); +} + #[test] fn upload_batches_take_the_main_lane_once_per_batch() { let _allowance = crate::scheduler::ALLOWANCE_TEST_LOCK @@ -557,7 +606,7 @@ fn file_parts_use_one_upload_session_and_metadata_stays_on_main() { Ok(()) }) .expect("recorded parts"); - let upload_id = client.upload.lock().snapshot.session_id; + let upload_id = client.upload[0].io.lock().snapshot.session_id; assert_ne!(upload_id, main_id); assert!(!media_ids.contains(&upload_id)); assert!(!rpc_ids.contains(&upload_id)); @@ -593,14 +642,14 @@ fn idle_file_cleanup_closes_sockets_and_preserves_busy_lanes_and_main() { media.last_io = now; upload.last_io = now; client.media[0].io.lock().transport = Some(media); - client.upload.lock().transport = Some(upload); + client.upload[0].io.lock().transport = Some(upload); client.main.lock().transport = Some(main); assert!(close_idle_file_transports(&client, now)); let busy = client.media[0].io.lock(); let later = now + crate::media::DOWNLOAD_SESSION_IDLE; assert!(close_idle_file_transports(&client, later)); assert!(busy.transport.is_some()); - assert!(client.upload.lock().transport.is_none()); + assert!(client.upload[0].io.lock().transport.is_none()); drop(busy); assert!(!close_idle_file_transports(&client, later)); assert!(client.media[0].io.lock().transport.is_none()); @@ -709,7 +758,7 @@ fn expired_upload_rtt_does_not_require_a_probe_connection_or_inherit_download_ca }) .unwrap(); assert_eq!(staging[0].in_flight(), 8); - assert!(client.upload.lock().transport.is_none()); + assert!(client.upload[0].io.lock().transport.is_none()); destroy_client(handle); std::fs::remove_file(path).ok(); std::fs::remove_file(file).ok(); @@ -738,7 +787,7 @@ fn failed_rtt_probe_on_an_existing_upload_connection_does_not_fail_saved_parts() let (peer, _) = listener.accept().unwrap(); drop(peer); { - let mut io = client.upload.lock(); + let mut io = client.upload[0].io.lock(); io.snapshot = fork_session(&home); io.transport = Some(transport); } @@ -763,7 +812,7 @@ fn failed_rtt_probe_on_an_existing_upload_connection_does_not_fail_saved_parts() assert_eq!(saved, 1); assert_eq!(staging.len(), 1); assert!(client.policy.upload_rtt.lock().is_none()); - assert!(client.upload.lock().transport.is_none()); + assert!(client.upload[0].io.lock().transport.is_none()); destroy_client(handle); std::fs::remove_file(path).ok(); std::fs::remove_file(file).ok(); diff --git a/native/mtproto-rs/tests/unit/client_mgr_logout_tests.rs b/native/mtproto-rs/tests/unit/client_mgr_logout_tests.rs index 776d14b6f..3cba129bb 100644 --- a/native/mtproto-rs/tests/unit/client_mgr_logout_tests.rs +++ b/native/mtproto-rs/tests/unit/client_mgr_logout_tests.rs @@ -28,6 +28,9 @@ fn logout_clears_unauthorized_local_session_without_network() { session_dead_reason: None, test_dc: false, last_inline: None, + perm_auth_key: None, + perm_salt: 0, + perm_session_id: 0, }; persist(&state).expect("seed session"); assert!(path.exists()); @@ -136,6 +139,7 @@ fn cancelled_request_exits_while_lane_is_still_owned_by_another_call() { last_difference: None, snapshot: Snapshot::new(2, &mut OsRandom).unwrap(), transport: None, + temp_expires_at: 0, })); let held = lane.lock(); let worker_lane = lane.clone(); diff --git a/native/mtproto-rs/tests/unit/media_page_tests.rs b/native/mtproto-rs/tests/unit/media_page_tests.rs index ce84a7c57..8f3e9c77c 100644 --- a/native/mtproto-rs/tests/unit/media_page_tests.rs +++ b/native/mtproto-rs/tests/unit/media_page_tests.rs @@ -138,12 +138,24 @@ fn paragraph_of(text: RichText) -> PageBlock { #[test] fn rich_message_fields_survive_flattening() { let marked = paragraph_of(RichText::TextMarked( - tellers_mtproto::latest::api::TextMarkedConstructor { text: plain("secret") }, + tellers_mtproto::latest::api::TextMarkedConstructor { + text: plain("secret"), + }, )); let marked_text = page_blocks_formatted(std::iter::once(&marked)); assert_eq!(marked_text.text, "secret"); - assert!(marked_text.entities.iter().any(|entity| entity.kind == "marked")); - assert!(marked_text.entities.iter().all(|entity| entity.kind != "spoiler")); + assert!( + marked_text + .entities + .iter() + .any(|entity| entity.kind == "marked") + ); + assert!( + marked_text + .entities + .iter() + .all(|entity| entity.kind != "spoiler") + ); let email = paragraph_of(RichText::TextEmail( tellers_mtproto::latest::api::TextEmailConstructor { @@ -227,17 +239,13 @@ fn rich_message_fields_survive_flattening() { let date_text = page_blocks_formatted(std::iter::once(&date)); assert_eq!(date_text.text, "fallback"); assert_eq!(date_text.entities[0].kind, "date"); - assert_eq!( - date_text.entities[0].url.as_deref(), - Some("1710000000|8") - ); + assert_eq!(date_text.entities[0].url.as_deref(), Some("1710000000|8")); } #[test] fn divider_table_and_block_math_keep_structure() { - let divider = PageBlock::PageBlockDivider( - tellers_mtproto::latest::api::PageBlockDividerConstructor {}, - ); + let divider = + PageBlock::PageBlockDivider(tellers_mtproto::latest::api::PageBlockDividerConstructor {}); let divided = page_blocks_formatted(std::iter::once(÷r)); assert_eq!(divided.text, "---"); assert_eq!(divided.entities[0].kind, "rule"); @@ -268,8 +276,8 @@ fn divider_table_and_block_math_keep_structure() { }, ) }; - let table = PageBlock::PageBlockTable( - tellers_mtproto::latest::api::PageBlockTableConstructor { + let table = + PageBlock::PageBlockTable(tellers_mtproto::latest::api::PageBlockTableConstructor { flags: 0, bordered: None, striped: None, @@ -279,8 +287,7 @@ fn divider_table_and_block_math_keep_structure() { row(vec![cell("A", true), cell("B", true)]), row(vec![cell("1", false), cell("2", false)]), ])), - }, - ); + }); let formatted = page_blocks_formatted(std::iter::once(&table)); assert!(formatted.text.contains("| A | B |")); assert!(formatted.text.contains("| 1 | 2 |")); @@ -289,17 +296,13 @@ fn divider_table_and_block_math_keep_structure() { .iter() .find(|entity| entity.kind == "table") .expect("table entity"); - assert_eq!( - table_entity.url.as_deref(), - Some("A\u{1f}B\u{1e}1\u{1f}2") - ); + assert_eq!(table_entity.url.as_deref(), Some("A\u{1f}B\u{1e}1\u{1f}2")); assert!(formatted.entities.iter().all(|entity| entity.kind != "pre")); - let block_math = PageBlock::PageBlockMath( - tellers_mtproto::latest::api::PageBlockMathConstructor { + let block_math = + PageBlock::PageBlockMath(tellers_mtproto::latest::api::PageBlockMathConstructor { source: "\\frac{1}{2}".into(), - }, - ); + }); let math = page_blocks_formatted(std::iter::once(&block_math)); assert_eq!(math.entities[0].kind, "math"); assert_eq!(math.entities[0].url.as_deref(), Some("block")); @@ -308,53 +311,46 @@ fn divider_table_and_block_math_keep_structure() { #[test] fn unsupported_rich_media_is_not_invented() { - let video = PageBlock::PageBlockVideo( - tellers_mtproto::latest::api::PageBlockVideoConstructor { + let video = + PageBlock::PageBlockVideo(tellers_mtproto::latest::api::PageBlockVideoConstructor { flags: 0, autoplay: None, loop_: None, spoiler: None, video_id: 7, caption: empty_caption(), - }, - ); - let audio = PageBlock::PageBlockAudio( - tellers_mtproto::latest::api::PageBlockAudioConstructor { + }); + let audio = + PageBlock::PageBlockAudio(tellers_mtproto::latest::api::PageBlockAudioConstructor { audio_id: 8, caption: empty_caption(), - }, - ); - let collage = PageBlock::PageBlockCollage( - tellers_mtproto::latest::api::PageBlockCollageConstructor { + }); + let collage = + PageBlock::PageBlockCollage(tellers_mtproto::latest::api::PageBlockCollageConstructor { items: Box::new(boxed_vector(Vec::::new())), caption: empty_caption(), - }, - ); + }); let slideshow = PageBlock::PageBlockSlideshow( tellers_mtproto::latest::api::PageBlockSlideshowConstructor { items: Box::new(boxed_vector(Vec::::new())), caption: empty_caption(), }, ); - let map = PageBlock::PageBlockMap( - tellers_mtproto::latest::api::PageBlockMapConstructor { - geo: Box::new(tellers_mtproto::latest::api::GeoPoint::GeoPointEmpty( - tellers_mtproto::latest::api::GeoPointEmptyConstructor {}, - )), - zoom: 1, - w: 10, - h: 10, - caption: empty_caption(), - }, - ); + let map = PageBlock::PageBlockMap(tellers_mtproto::latest::api::PageBlockMapConstructor { + geo: Box::new(tellers_mtproto::latest::api::GeoPoint::GeoPointEmpty( + tellers_mtproto::latest::api::GeoPointEmptyConstructor {}, + )), + zoom: 1, + w: 10, + h: 10, + caption: empty_caption(), + }); for block in [&video, &audio, &collage, &slideshow, &map] { let formatted = page_blocks_formatted(std::iter::once(block)); assert!(formatted.text.is_empty(), "{}", formatted.text); - assert!( - formatted.entities.iter().all(|entity| { - entity.kind != "photo" && entity.kind != "video" && entity.kind != "audio" - }) - ); + assert!(formatted.entities.iter().all(|entity| { + entity.kind != "photo" && entity.kind != "video" && entity.kind != "audio" + })); } } @@ -362,5 +358,10 @@ fn unsupported_rich_media_is_not_invented() { fn plain_dashes_are_not_a_divider() { let formatted = page_blocks_formatted(std::iter::once(¶graph("---"))); assert_eq!(formatted.text, "---"); - assert!(formatted.entities.iter().all(|entity| entity.kind != "rule")); + assert!( + formatted + .entities + .iter() + .all(|entity| entity.kind != "rule") + ); } diff --git a/native/mtproto-rs/tests/unit/media_thumbs_tests.rs b/native/mtproto-rs/tests/unit/media_thumbs_tests.rs index 1766dc520..502688664 100644 --- a/native/mtproto-rs/tests/unit/media_thumbs_tests.rs +++ b/native/mtproto-rs/tests/unit/media_thumbs_tests.rs @@ -473,5 +473,7 @@ fn peer_photo_with_stripped_thumb() { Some(&raw_stripped), ); assert_eq!(photo.thumb_cache_key, Some("avatar:7:thumb".into())); - assert!(matches!(photo.thumb_location, Some(MediaLocation::Inline { extension, .. }) if extension == "jpg")); + assert!( + matches!(photo.thumb_location, Some(MediaLocation::Inline { extension, .. }) if extension == "jpg") + ); } diff --git a/native/mtproto-rs/tests/unit/scheduler_tests.rs b/native/mtproto-rs/tests/unit/scheduler_tests.rs index 4ce0b40fc..041559ace 100644 --- a/native/mtproto-rs/tests/unit/scheduler_tests.rs +++ b/native/mtproto-rs/tests/unit/scheduler_tests.rs @@ -64,6 +64,7 @@ fn unknown_allowance_keeps_a_single_main_session() { let _serial = ALLOWANCE_TEST_LOCK .lock() .unwrap_or_else(|e| e.into_inner()); + set_bound_extra_sessions(0); let original = MAIN_SESSION_ALLOWANCE.load(std::sync::atomic::Ordering::Relaxed); MAIN_SESSION_ALLOWANCE.store(-1, std::sync::atomic::Ordering::Relaxed); assert!(!main_session_allowance_known()); @@ -78,6 +79,7 @@ fn tmp_sessions_bounds_extra_main_sessions() { let _serial = ALLOWANCE_TEST_LOCK .lock() .unwrap_or_else(|e| e.into_inner()); + set_bound_extra_sessions(0); let original = MAIN_SESSION_ALLOWANCE.load(std::sync::atomic::Ordering::Relaxed); // Absent `tmp_sessions` means one main session: no parallel home-DC lanes. assert_eq!(set_main_session_allowance(None), 1); @@ -91,6 +93,14 @@ fn tmp_sessions_bounds_extra_main_sessions() { assert_eq!(read_lanes(), 0); assert_eq!(set_main_session_allowance(Some(2)), 2); assert_eq!(read_lanes(), 0); + set_bound_extra_sessions(2); + assert_eq!(extra_main_sessions(), 1); + assert_eq!(read_lanes(), 1); + set_main_session_allowance(Some(8)); + set_bound_extra_sessions(2); + assert_eq!(extra_main_sessions(), 2); + set_bound_extra_sessions(0); + assert_eq!(extra_main_sessions(), 0); MAIN_SESSION_ALLOWANCE.store(original, std::sync::atomic::Ordering::Relaxed); } From 629c7685ba352e4463c8b91b6b5099b109746c1c Mon Sep 17 00:00:00 2001 From: Artur Skubei <41114720+gdlbo@users.noreply.github.com> Date: Thu, 8 Oct 2026 19:31:57 +0300 Subject: [PATCH 2/6] feat(media): download off the main socket --- .../org/monogram/core/ui/DownloadSettings.kt | 16 +-- .../monogram/core/ui/DownloadSettingsTest.kt | 25 ++--- .../feature/settings/ui/SettingsContent.kt | 1 - .../feature/settings/ui/SettingsData.kt | 20 ---- .../src/main/res/values-es/strings.xml | 3 - .../src/main/res/values-hy/strings.xml | 3 - .../src/main/res/values-ja/strings.xml | 3 - .../src/main/res/values-pt-rBR/strings.xml | 3 - .../src/main/res/values-ru/strings.xml | 3 - .../src/main/res/values-sk/strings.xml | 3 - .../src/main/res/values-tr/strings.xml | 3 - .../src/main/res/values-uk/strings.xml | 3 - .../src/main/res/values-zh/strings.xml | 3 - .../settings/src/main/res/values/strings.xml | 3 - native/mtproto-rs/src/client/lanes.rs | 2 + .../mtproto-rs/src/client/media_download.rs | 98 ++++++++++--------- native/mtproto-rs/src/client/mod.rs | 5 + native/mtproto-rs/src/tcp.rs | 6 +- native/mtproto-rs/src/transfer_policy.rs | 49 ++++++++-- .../tests/unit/client_mgr_logout_tests.rs | 1 + .../tests/unit/transfer_policy_tests.rs | 46 +++++++++ 21 files changed, 162 insertions(+), 137 deletions(-) diff --git a/core/ui/src/main/java/org/monogram/core/ui/DownloadSettings.kt b/core/ui/src/main/java/org/monogram/core/ui/DownloadSettings.kt index 3ccb2b28c..f98543f11 100644 --- a/core/ui/src/main/java/org/monogram/core/ui/DownloadSettings.kt +++ b/core/ui/src/main/java/org/monogram/core/ui/DownloadSettings.kt @@ -147,7 +147,6 @@ data class AutoDownloadPreset( data class DownloadState( val lanes: Int = 8, val speedUpUploads: Boolean = false, - val speedUpDownloads: Boolean = false, val premium: Boolean = false, val wifi: AutoDownloadPreset = AutoDownloadPreset.WIFI, val mobile: AutoDownloadPreset = AutoDownloadPreset.MOBILE, @@ -156,12 +155,12 @@ data class DownloadState( val autoplayVideos: Boolean = true, val activeNetwork: AutoDownloadNetwork = AutoDownloadNetwork.Wifi, ) { - val effectiveLanes: Int get() = lanes.coerceIn(1, if (premium && speedUpDownloads) 8 else 2) - val parts: Int get() = if (!premium) 2 else if (speedUpDownloads) 8 else 4 + val effectiveLanes: Int get() = lanes.coerceIn(1, if (premium) 8 else 2) + val parts: Int get() = if (premium) 8 else 4 val concurrency: DownloadConcurrency get() = DownloadConcurrency(lanes = effectiveLanes, parts = parts) val filePartKib: Int get() = 512 - val downloadChunkKib: Int get() = if (premium && speedUpDownloads) 512 else 128 + val downloadChunkKib: Int get() = if (premium) 512 else 128 fun presetFor(network: AutoDownloadNetwork): AutoDownloadPreset = when (network) { AutoDownloadNetwork.Wifi -> wifi @@ -173,7 +172,6 @@ data class DownloadState( object DownloadSettings { private const val PREFS = "monogram_download" private const val KEY_SPEED_UP_UPLOADS = "speed_up_uploads" - private const val KEY_SPEED_UP_DOWNLOADS = "speed_up_downloads" private const val KEY_WIFI = "autodownload_wifi" private const val KEY_MOBILE = "autodownload_mobile" private const val KEY_ROAMING = "autodownload_roaming" @@ -203,7 +201,6 @@ object DownloadSettings { val prefs = app.getSharedPreferences(PREFS, Context.MODE_PRIVATE) mutable.value = DownloadState( speedUpUploads = prefs.getBoolean(KEY_SPEED_UP_UPLOADS, false), - speedUpDownloads = prefs.getBoolean(KEY_SPEED_UP_DOWNLOADS, false), wifi = AutoDownloadPreset.decode(prefs.getString(KEY_WIFI, null), AutoDownloadPreset.WIFI), mobile = AutoDownloadPreset.decode(prefs.getString(KEY_MOBILE, null), AutoDownloadPreset.MOBILE), roaming = AutoDownloadPreset.decode(prefs.getString(KEY_ROAMING, null), AutoDownloadPreset.ROAMING), @@ -225,12 +222,6 @@ object DownloadSettings { apply?.invoke(mutable.value) } - fun setSpeedUpDownloads(enabled: Boolean) { - mutable.update { it.copy(speedUpDownloads = enabled) } - persist() - apply?.invoke(mutable.value) - } - fun setPremium(premium: Boolean) { if (mutable.value.premium == premium) return mutable.update { it.copy(premium = premium) } @@ -309,7 +300,6 @@ object DownloadSettings { val value = mutable.value prefs.edit { putBoolean(KEY_SPEED_UP_UPLOADS, value.speedUpUploads) - putBoolean(KEY_SPEED_UP_DOWNLOADS, value.speedUpDownloads) putString(KEY_WIFI, value.wifi.encode()) putString(KEY_MOBILE, value.mobile.encode()) putString(KEY_ROAMING, value.roaming.encode()) diff --git a/core/ui/src/test/java/org/monogram/core/ui/DownloadSettingsTest.kt b/core/ui/src/test/java/org/monogram/core/ui/DownloadSettingsTest.kt index b4c31af1a..603c319cc 100644 --- a/core/ui/src/test/java/org/monogram/core/ui/DownloadSettingsTest.kt +++ b/core/ui/src/test/java/org/monogram/core/ui/DownloadSettingsTest.kt @@ -13,22 +13,17 @@ class DownloadSettingsTest { } @Test - fun downloadBoostRespectsPremiumAndConfiguredLaneCap() { - for (boost in listOf(false, true)) { - val free = DownloadState(speedUpDownloads = boost, premium = false) - assertEquals(DownloadConcurrency(2, 2), free.concurrency) - assertEquals(128, free.downloadChunkKib) - assertEquals(512, free.filePartKib) - } - val premium = DownloadState(premium = true) - assertEquals(DownloadConcurrency(2, 4), premium.concurrency) - assertEquals(128, premium.downloadChunkKib) - val boosted = premium.copy(speedUpDownloads = true, lanes = 7) - assertEquals(DownloadConcurrency(7, 8), boosted.concurrency) - assertEquals(512, boosted.downloadChunkKib) + fun premiumDownloadUsesTheFastDefaults() { + val free = DownloadState(premium = false, lanes = 99) + assertEquals(DownloadConcurrency(2, 4), free.concurrency) + assertEquals(128, free.downloadChunkKib) + assertEquals(512, free.filePartKib) + val premium = DownloadState(premium = true, lanes = 7) + assertEquals(DownloadConcurrency(7, 8), premium.concurrency) + assertEquals(512, premium.downloadChunkKib) assertEquals(512, premium.filePartKib) - assertEquals(1, DownloadState(lanes = 1).effectiveLanes) - assertEquals(8, boosted.copy(lanes = 99).effectiveLanes) + assertEquals(1, DownloadState(lanes = 1, premium = true).effectiveLanes) + assertEquals(8, DownloadState(lanes = 99, premium = true).effectiveLanes) } @Test diff --git a/feature/settings/src/main/java/org/monogram/feature/settings/ui/SettingsContent.kt b/feature/settings/src/main/java/org/monogram/feature/settings/ui/SettingsContent.kt index 76076ad59..c3154b8c3 100644 --- a/feature/settings/src/main/java/org/monogram/feature/settings/ui/SettingsContent.kt +++ b/feature/settings/src/main/java/org/monogram/feature/settings/ui/SettingsContent.kt @@ -352,7 +352,6 @@ fun SettingsContent( cacheMessage = state.cacheMessage, loading = state.loading || state.loggingOut, download = download, - onSpeedUpDownloads = DownloadSettings::setSpeedUpDownloads, onOpenAutoDownload = { network -> component.openPage( SettingsPage.AutoDownload(autoDownloadNetworkKey(network)), diff --git a/feature/settings/src/main/java/org/monogram/feature/settings/ui/SettingsData.kt b/feature/settings/src/main/java/org/monogram/feature/settings/ui/SettingsData.kt index 2e1504561..289b6711a 100644 --- a/feature/settings/src/main/java/org/monogram/feature/settings/ui/SettingsData.kt +++ b/feature/settings/src/main/java/org/monogram/feature/settings/ui/SettingsData.kt @@ -18,7 +18,6 @@ import androidx.compose.material.icons.automirrored.outlined.InsertDriveFile import androidx.compose.material.icons.automirrored.outlined.KeyboardArrowRight import androidx.compose.material.icons.outlined.BugReport import androidx.compose.material.icons.outlined.CellTower -import androidx.compose.material.icons.outlined.CloudDownload import androidx.compose.material.icons.outlined.DeleteSweep import androidx.compose.material.icons.outlined.EmojiEmotions import androidx.compose.material.icons.outlined.Gif @@ -169,7 +168,6 @@ internal fun LazyListScope.dataItems( cacheMessage: String?, loading: Boolean, download: DownloadState, - onSpeedUpDownloads: (Boolean) -> Unit, onOpenAutoDownload: (AutoDownloadNetwork) -> Unit, onClear: () -> Unit, onClearChat: (Long) -> Unit, @@ -304,24 +302,6 @@ internal fun LazyListScope.dataItems( ) } item { Spacer(Modifier.height(8.dp)) } - item { SectionHeader(stringResource(R.string.settings_data_transfers)) } - item { - SettingsTile( - icon = Icons.Outlined.CloudDownload, - title = stringResource(R.string.settings_speed_up_downloads), - subtitle = stringResource(R.string.settings_speed_up_downloads_sub), - iconColor = MaterialTheme.colorScheme.primary, - position = ItemPosition.STANDALONE, - onClick = { onSpeedUpDownloads(!download.speedUpDownloads) }, - trailingContent = { - Switch( - checked = download.speedUpDownloads, - onCheckedChange = onSpeedUpDownloads, - ) - }, - ) - } - item { Spacer(Modifier.height(8.dp)) } item { val photos = cacheByKind[FileCache.KIND_PHOTOS] ?: 0L val videos = cacheByKind[FileCache.KIND_VIDEOS] ?: 0L diff --git a/feature/settings/src/main/res/values-es/strings.xml b/feature/settings/src/main/res/values-es/strings.xml index 8fc4b794e..2ee327cea 100644 --- a/feature/settings/src/main/res/values-es/strings.xml +++ b/feature/settings/src/main/res/values-es/strings.xml @@ -126,9 +126,6 @@ Última conexión, llamadas y reenvío siguen los ajustes de privacidad de Telegram — abre tu perfil para más detalles Datos y almacenamiento Descarga automática y caché - Transferencias - Acelerar descargas - Descargar fotos, videos y archivos más rápido con buena conexión Descarga automática de medios Con Wi-Fi Con datos móviles diff --git a/feature/settings/src/main/res/values-hy/strings.xml b/feature/settings/src/main/res/values-hy/strings.xml index 58896c16a..b6b88408e 100644 --- a/feature/settings/src/main/res/values-hy/strings.xml +++ b/feature/settings/src/main/res/values-hy/strings.xml @@ -126,9 +126,6 @@ Վերջին այցը, զանգերը և վերահասցեավորումը հետևում են Telegram-ի գաղտնիության կարգավորումներին — մանրամասները պրոֆիլում Տվյալներ և պահոց Ինքնաբեռնում և քեշ - Փոխանցումներ - Արագացնել ներբեռնումը - Լուսանկարները, տեսանյութերը և ֆայլերը լավ կապի դեպքում ավելի արագ են ներբեռնվում Մեդիայի ինքնաբեռնում Wi-Fi-ով Բջջային ցանցով diff --git a/feature/settings/src/main/res/values-ja/strings.xml b/feature/settings/src/main/res/values-ja/strings.xml index 0ded19ae7..83c7ebee8 100644 --- a/feature/settings/src/main/res/values-ja/strings.xml +++ b/feature/settings/src/main/res/values-ja/strings.xml @@ -126,9 +126,6 @@ 最終オンライン、通話、転送は Telegram のプライバシー設定に従います — 詳細はプロフィールをご覧ください データとストレージ 自動ダウンロードとキャッシュ - 転送 - ダウンロードを高速化 - 良好な接続時に写真、動画、ファイルをより速くダウンロードします メディアの自動ダウンロード Wi-Fi接続時 モバイル通信時 diff --git a/feature/settings/src/main/res/values-pt-rBR/strings.xml b/feature/settings/src/main/res/values-pt-rBR/strings.xml index f9848ec22..a16506493 100644 --- a/feature/settings/src/main/res/values-pt-rBR/strings.xml +++ b/feature/settings/src/main/res/values-pt-rBR/strings.xml @@ -126,9 +126,6 @@ Última vez, chamadas e encaminhamento seguem as configurações de privacidade do Telegram — abra seu perfil para detalhes Dados e armazenamento Baixar automaticamente e cache - Transferências - Acelerar downloads - Baixar fotos, vídeos e arquivos mais rápido em conexões boas Download automático de mídia No Wi-Fi Nos dados móveis diff --git a/feature/settings/src/main/res/values-ru/strings.xml b/feature/settings/src/main/res/values-ru/strings.xml index fe2575ec7..fdc2b89d9 100644 --- a/feature/settings/src/main/res/values-ru/strings.xml +++ b/feature/settings/src/main/res/values-ru/strings.xml @@ -126,9 +126,6 @@ Статус в сети, звонки и пересылка зависят от настроек конфиденциальности Telegram — подробности в профиле Данные и память Автозагрузка и кэш - Передача - Ускорить скачивание - Скачивать фото, видео и файлы быстрее при хорошем соединении Автозагрузка медиа Через Wi-Fi Через мобильную сеть diff --git a/feature/settings/src/main/res/values-sk/strings.xml b/feature/settings/src/main/res/values-sk/strings.xml index 5a5d58f61..20c758a15 100644 --- a/feature/settings/src/main/res/values-sk/strings.xml +++ b/feature/settings/src/main/res/values-sk/strings.xml @@ -126,9 +126,6 @@ Naposledy videný, hovory a preposielanie sa riadia nastaveniami súkromia Telegramu — podrobnosti nájdete v profile Údaje a úložisko Automatické sťahovanie a vyrovnávacia pamäť - Prenosy - Zrýchliť sťahovanie - Sťahujte fotky, videá a súbory rýchlejšie pri dobrom pripojení Automatické sťahovanie médií Cez Wi-Fi Cez mobilné dáta diff --git a/feature/settings/src/main/res/values-tr/strings.xml b/feature/settings/src/main/res/values-tr/strings.xml index fa3361c60..09df18eb0 100644 --- a/feature/settings/src/main/res/values-tr/strings.xml +++ b/feature/settings/src/main/res/values-tr/strings.xml @@ -126,9 +126,6 @@ Son görülme, aramalar ve iletme Telegram gizlilik ayarlarına bağlıdır — ayrıntılar için profilinizi açın Veri ve depolama Otomatik indirme ve önbellek - Aktarımlar - İndirmeleri hızlandır - İyi bir bağlantıda fotoğrafları, videoları ve dosyaları daha hızlı indirin Otomatik medya indirme Wi-Fi ile Hücresel veri ile diff --git a/feature/settings/src/main/res/values-uk/strings.xml b/feature/settings/src/main/res/values-uk/strings.xml index 09d543c18..326c058d0 100644 --- a/feature/settings/src/main/res/values-uk/strings.xml +++ b/feature/settings/src/main/res/values-uk/strings.xml @@ -126,9 +126,6 @@ Статус, дзвінки та пересилання залежать від налаштувань приватності Telegram — деталі у профілі Дані та пам\'ять Автозавантаження та кеш - Передавання - Прискорити завантаження - Завантажувати фото, відео та файли швидше за стабільного з\'єднання Автозавантаження медіа Через Wi-Fi Через мобільну мережу diff --git a/feature/settings/src/main/res/values-zh/strings.xml b/feature/settings/src/main/res/values-zh/strings.xml index b94db3d35..5b1277a84 100644 --- a/feature/settings/src/main/res/values-zh/strings.xml +++ b/feature/settings/src/main/res/values-zh/strings.xml @@ -126,9 +126,6 @@ 最后上线时间、通话和转发遵循 Telegram 的隐私设置 — 打开资料查看详情 数据和存储 自动下载与缓存 - 传输 - 加速下载 - 在良好网络连接下更快下载照片、视频和文件 自动下载媒体 使用 Wi-Fi 使用移动网络 diff --git a/feature/settings/src/main/res/values/strings.xml b/feature/settings/src/main/res/values/strings.xml index 391b0e00d..10335c7b8 100644 --- a/feature/settings/src/main/res/values/strings.xml +++ b/feature/settings/src/main/res/values/strings.xml @@ -126,9 +126,6 @@ Last seen, calls, and forwarding follow Telegram\'s privacy settings — open your profile for details Data and Storage Auto-download and cache - Transfers - Speed up downloads - Download photos, videos, and files faster on a good connection Automatic media download Using Wi-Fi Using mobile data diff --git a/native/mtproto-rs/src/client/lanes.rs b/native/mtproto-rs/src/client/lanes.rs index 85aa188e1..b11f236a2 100644 --- a/native/mtproto-rs/src/client/lanes.rs +++ b/native/mtproto-rs/src/client/lanes.rs @@ -31,6 +31,8 @@ pub(crate) struct SessionIo { pub(crate) transport: Option, /// Unix time when this lane's temporary auth key expires. Zero means none. pub(crate) temp_expires_at: i32, + /// Key this lane's temp key was bound from. Not the key on the wire. + pub(crate) perm_key: Option>, } /// A session plus the gate that decides which request class gets it next. diff --git a/native/mtproto-rs/src/client/media_download.rs b/native/mtproto-rs/src/client/media_download.rs index 90476ee51..4342970d8 100644 --- a/native/mtproto-rs/src/client/media_download.rs +++ b/native/mtproto-rs/src/client/media_download.rs @@ -18,6 +18,16 @@ fn home_lease_ok(state: &ClientState, home_session_id: i64) -> bool { state.snapshot.session_id == home_session_id || state.perm_session_id == home_session_id } +fn download_temp_ready(io: &super::lanes::SessionIo, dc: i32, home_key: Option<&[u8]>) -> bool { + io.temp_expires_at > crate::auth::bind_temp::unix_now() + && io.snapshot.dc_id == dc + && io + .snapshot + .auth_key + .as_deref() + .is_some_and(|key| Some(key) != home_key) +} + pub fn download_wallpaper( handle: u64, id: i64, @@ -689,21 +699,9 @@ pub(crate) fn download_media_range_on_lane( loop { media::check_transfer_flood(target_dc, media::TransferClass::Download)?; let admission = client.policy.admit_download(target_dc, large)?; - let first = if target_dc == home.dc_id { - download_media_range_on_home_session( - handle, - client, - home.session_id, - api_id, - media, - dest, - offset, - ) - } else { - download_media_range_on_lane_dc( - handle, client, home, api_id, media, dest, offset, target_dc, - ) - }; + let first = download_media_range_on_lane_dc( + handle, client, home, api_id, media, dest, offset, target_dc, + ); if let Err(ref error) = first { if let Some(dc) = api_invoke::migrate_dc(error) { if dc != target_dc && migrations < 3 { @@ -731,21 +729,9 @@ pub(crate) fn download_media_range_on_lane( Ok(path) => Ok(path), Err(_) => media::with_cdn_supported(false, || { let _admission = client.policy.admit_download(target_dc, large)?; - if target_dc == home.dc_id { - download_media_range_on_home_session( - handle, - client, - home.session_id, - api_id, - media, - dest, - offset, - ) - } else { - download_media_range_on_lane_dc( - handle, client, home, api_id, media, dest, offset, target_dc, - ) - } + download_media_range_on_lane_dc( + handle, client, home, api_id, media, dest, offset, target_dc, + ) }), } } @@ -754,11 +740,8 @@ pub(crate) fn download_media_range_on_lane( } } -/// Home-DC download over the single main session. -/// -/// A second session on this DC would be another main session. -/// https://core.telegram.org/api/datacenter#parallel-sessions -#[allow(dead_code)] +/// Home-DC download on the main socket. Used only when a temp key cannot be bound, +/// so the permanent key is not opened on a second TCP. pub(crate) fn download_media_range_on_home_session( handle: u64, client: &Client, @@ -929,7 +912,9 @@ pub(crate) fn download_media_range_on_lane_dc( return Err(expired_session_lease()); } let mut lane_perm = None; - let already_temp = io.temp_expires_at > crate::auth::bind_temp::unix_now() + let now = crate::auth::bind_temp::unix_now(); + let already_temp = io.temp_expires_at > now + && io.snapshot.dc_id == target_dc && io.snapshot.auth_key.as_ref() != home.auth_key.as_ref(); if !already_temp { match prepare_media_lane_snapshot(&mut io.snapshot, home, target_dc) { @@ -946,6 +931,7 @@ pub(crate) fn download_media_range_on_lane_dc( MediaLanePrep::Reuse => {} } if let Some(key) = io.snapshot.auth_key.clone() { + io.perm_key = Some(key.clone()); lane_perm = Some(key.clone()); let offset = io.snapshot.time_offset_micros; let dc = io.snapshot.dc_id; @@ -968,19 +954,39 @@ pub(crate) fn download_media_range_on_lane_dc( offset, ); } - let mut io2 = try_lock_second_media_lane(client); - let second_bound = if let (Some(second), Some(key)) = (io2.as_mut(), lane_perm.as_deref()) { - crate::auth::bind_temp::negotiate_bound_temp(key, target_dc, home.time_offset_micros) - .ok() - .map(|bound| { - second.snapshot = bound.snapshot; - second.temp_expires_at = bound.expires_at; - second.transport = None; - }) + let bind_key = lane_perm.or(io.perm_key.clone()).or_else(|| { + (target_dc == home.dc_id) + .then(|| home.auth_key.clone()) + .flatten() + }); + let mut io2 = if media::download_in_flight() > 1 { + try_lock_second_media_lane(client) } else { None }; - if second_bound.is_none() { + let second_ready = match (io2.as_mut(), bind_key.as_deref()) { + (Some(second), _) if download_temp_ready(second, target_dc, home.auth_key.as_deref()) => { + true + } + (Some(second), Some(key)) => { + match crate::auth::bind_temp::negotiate_bound_temp( + key, + target_dc, + home.time_offset_micros, + ) { + Ok(bound) => { + second.perm_key = Some(key.to_vec()); + second.snapshot = bound.snapshot; + second.temp_expires_at = bound.expires_at; + second.transport = None; + true + } + Err(_) => false, + } + } + _ => false, + }; + if !second_ready { io2 = None; } let mut media_snap = io.snapshot.clone(); diff --git a/native/mtproto-rs/src/client/mod.rs b/native/mtproto-rs/src/client/mod.rs index b1074000f..6549427e4 100644 --- a/native/mtproto-rs/src/client/mod.rs +++ b/native/mtproto-rs/src/client/mod.rs @@ -289,6 +289,7 @@ pub fn create_client(api_id: i32, api_hash: String, session_path: String) -> u64 snapshot: fork_session(&snapshot), transport: None, temp_expires_at: 0, + perm_key: None, }) }); let handle = NEXT_HANDLE.fetch_add(1, Ordering::Relaxed); @@ -338,6 +339,7 @@ pub fn create_client(api_id: i32, api_hash: String, session_path: String) -> u64 snapshot, transport: None, temp_expires_at: 0, + perm_key: None, }), rpc: [ Lane::new(SessionIo { @@ -346,6 +348,7 @@ pub fn create_client(api_id: i32, api_hash: String, session_path: String) -> u64 snapshot: rpc_snapshot, transport: None, temp_expires_at: 0, + perm_key: None, }), Lane::new(SessionIo { pending_push: Default::default(), @@ -353,6 +356,7 @@ pub fn create_client(api_id: i32, api_hash: String, session_path: String) -> u64 snapshot: rpc_snapshot_b, transport: None, temp_expires_at: 0, + perm_key: None, }), ], media: media_snapshots @@ -364,6 +368,7 @@ pub fn create_client(api_id: i32, api_hash: String, session_path: String) -> u64 last_difference: None, transport: None, temp_expires_at: 0, + perm_key: None, }) }) .collect(), diff --git a/native/mtproto-rs/src/tcp.rs b/native/mtproto-rs/src/tcp.rs index 571c64e1a..58e91d80c 100644 --- a/native/mtproto-rs/src/tcp.rs +++ b/native/mtproto-rs/src/tcp.rs @@ -114,9 +114,7 @@ impl TcpConnection { pub fn set_io_timeout_ms(&mut self, ms: u64) { let timeout = Duration::from_millis(ms.max(50)); self.read_timeout = timeout; - let _ = self - .stream - .set_read_timeout(Some(timeout.min(Duration::from_millis(250)))); + let _ = self.stream.set_read_timeout(Some(timeout)); let _ = self.stream.set_write_timeout(Some(timeout)); } @@ -179,7 +177,7 @@ impl ObfuscatedTcp { pub fn set_io_timeout_ms(&mut self, ms: u64) { self.read_timeout = Duration::from_millis(ms.max(50)); - let read = self.read_timeout.min(Duration::from_millis(250)); + let read = self.read_timeout; let _ = match &mut self.inner { NativeTransportConnection::Obfuscated(connection) => connection.set_io_timeout_ms(ms), NativeTransportConnection::Http(connection) => connection.set_io_timeout_ms(ms), diff --git a/native/mtproto-rs/src/transfer_policy.rs b/native/mtproto-rs/src/transfer_policy.rs index fa1b3dec5..b11589a44 100644 --- a/native/mtproto-rs/src/transfer_policy.rs +++ b/native/mtproto-rs/src/transfer_policy.rs @@ -26,7 +26,7 @@ pub(crate) struct TransferPolicy { progress: RwLock>, pub(crate) floods: parking_lot::Mutex, pub(crate) queue_limits: parking_lot::Mutex<(usize, usize)>, - downloads: parking_lot::Mutex>, + downloads: parking_lot::Mutex, download_wake: parking_lot::Condvar, pub(crate) upload_rtt: parking_lot::Mutex>, @@ -50,7 +50,7 @@ impl TransferPolicy { progress: RwLock::new(None), floods: parking_lot::Mutex::new(crate::media::FloodScope::new()), queue_limits: parking_lot::Mutex::new((5, 2)), - downloads: parking_lot::Mutex::new(crate::HashMap::new()), + downloads: parking_lot::Mutex::new(ActiveDownloads::default()), download_wake: parking_lot::Condvar::new(), upload_rtt: parking_lot::Mutex::new(None), config_due: std::sync::atomic::AtomicU64::new(0), @@ -82,15 +82,34 @@ impl TransferPolicy { dc: i32, large: bool, ) -> Result { + let priority = crate::scheduler::current_class().priority(); let mut active = self.downloads.lock(); + let ticket = active.next_ticket; + active.next_ticket = active.next_ticket.wrapping_add(1); + active.waiters.push(DownloadWaiter { + ticket, + priority, + dc, + large, + }); loop { - crate::request_control::check() - .map_err(|e| crate::MtprotoError::Message(e.to_string()))?; + if let Err(error) = crate::request_control::check() { + active.waiters.retain(|waiter| waiter.ticket != ticket); + self.download_wake.notify_all(); + return Err(crate::MtprotoError::Message(error.to_string())); + } let limits = *self.queue_limits.lock(); let limit = if large { limits.1.min(2) } else { limits.0 }; - let count = active.entry((dc, large)).or_insert(0); - if *count < limit.max(1) { - *count += 1; + let count = active.counts.get(&(dc, large)).copied().unwrap_or(0); + let best = active + .waiters + .iter() + .filter(|waiter| waiter.dc == dc && waiter.large == large) + .min_by_key(|waiter| (waiter.priority, waiter.ticket)) + .map(|waiter| waiter.ticket); + if count < limit.max(1) && best == Some(ticket) { + *active.counts.entry((dc, large)).or_insert(0) += 1; + active.waiters.retain(|waiter| waiter.ticket != ticket); return Ok(DownloadAdmission { policy: self.clone(), dc, @@ -218,13 +237,27 @@ pub(crate) struct DownloadAdmission { impl Drop for DownloadAdmission { fn drop(&mut self) { let mut active = self.policy.downloads.lock(); - if let Some(count) = active.get_mut(&(self.dc, self.large)) { + if let Some(count) = active.counts.get_mut(&(self.dc, self.large)) { *count = count.saturating_sub(1); } self.policy.download_wake.notify_all(); } } +#[derive(Default)] +struct ActiveDownloads { + counts: crate::HashMap<(i32, bool), usize>, + waiters: Vec, + next_ticket: u64, +} + +struct DownloadWaiter { + ticket: u64, + priority: u8, + dc: i32, + large: bool, +} + fn valid_chunk(size: i32) -> bool { matches!(size, 131_072 | 262_144 | 524_288) } diff --git a/native/mtproto-rs/tests/unit/client_mgr_logout_tests.rs b/native/mtproto-rs/tests/unit/client_mgr_logout_tests.rs index 3cba129bb..8c995f2c4 100644 --- a/native/mtproto-rs/tests/unit/client_mgr_logout_tests.rs +++ b/native/mtproto-rs/tests/unit/client_mgr_logout_tests.rs @@ -140,6 +140,7 @@ fn cancelled_request_exits_while_lane_is_still_owned_by_another_call() { snapshot: Snapshot::new(2, &mut OsRandom).unwrap(), transport: None, temp_expires_at: 0, + perm_key: None, })); let held = lane.lock(); let worker_lane = lane.clone(); diff --git a/native/mtproto-rs/tests/unit/transfer_policy_tests.rs b/native/mtproto-rs/tests/unit/transfer_policy_tests.rs index 8a6b26953..fe8abb7bb 100644 --- a/native/mtproto-rs/tests/unit/transfer_policy_tests.rs +++ b/native/mtproto-rs/tests/unit/transfer_policy_tests.rs @@ -67,6 +67,52 @@ fn ffi_media_dispatch_profiles_match_kotlin_values_and_reset_after_request() { assert_eq!(crate::media::request_profile(), None); } +#[test] +fn interactive_download_preempts_a_queued_background_download() { + let policy = Arc::new(TransferPolicy::stock()); + *policy.queue_limits.lock() = (1, 1); + let held = policy.admit_download(2, false).unwrap(); + let (order_tx, order_rx) = std::sync::mpsc::channel(); + let (release_bg_tx, release_bg_rx) = std::sync::mpsc::channel::<()>(); + let background = policy.clone(); + let bg_order = order_tx.clone(); + let bg = std::thread::spawn(move || { + crate::scheduler::with_class(crate::scheduler::RequestClass::BackgroundMedia, || { + let _lease = background.admit_download(2, false).unwrap(); + bg_order.send("background").unwrap(); + let _ = release_bg_rx.recv_timeout(std::time::Duration::from_secs(2)); + }); + }); + std::thread::sleep(std::time::Duration::from_millis(40)); + let (release_fg_tx, release_fg_rx) = std::sync::mpsc::channel::<()>(); + let interactive = policy.clone(); + let fg = std::thread::spawn(move || { + crate::scheduler::with_class(crate::scheduler::RequestClass::InteractiveMedia, || { + let _lease = interactive.admit_download(2, false).unwrap(); + order_tx.send("interactive").unwrap(); + let _ = release_fg_rx.recv_timeout(std::time::Duration::from_secs(2)); + }); + }); + std::thread::sleep(std::time::Duration::from_millis(40)); + drop(held); + assert_eq!( + order_rx + .recv_timeout(std::time::Duration::from_secs(1)) + .ok(), + Some("interactive") + ); + release_fg_tx.send(()).unwrap(); + assert_eq!( + order_rx + .recv_timeout(std::time::Duration::from_secs(1)) + .ok(), + Some("background") + ); + release_bg_tx.send(()).unwrap(); + fg.join().unwrap(); + bg.join().unwrap(); +} + #[test] fn server_queue_limits_block_only_the_matching_dc_and_file_queue() { let policy = Arc::new(TransferPolicy::stock()); From 6a44a6410beeeed818cbdd4e61642d026b513d75 Mon Sep 17 00:00:00 2001 From: Artur Skubei <41114720+gdlbo@users.noreply.github.com> Date: Thu, 8 Oct 2026 19:44:34 +0300 Subject: [PATCH 3/6] feat: remove speed-up uploads setting --- .../main/java/org/monogram/core/ui/DownloadSettings.kt | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/core/ui/src/main/java/org/monogram/core/ui/DownloadSettings.kt b/core/ui/src/main/java/org/monogram/core/ui/DownloadSettings.kt index f98543f11..5881be0b0 100644 --- a/core/ui/src/main/java/org/monogram/core/ui/DownloadSettings.kt +++ b/core/ui/src/main/java/org/monogram/core/ui/DownloadSettings.kt @@ -146,7 +146,6 @@ data class AutoDownloadPreset( data class DownloadState( val lanes: Int = 8, - val speedUpUploads: Boolean = false, val premium: Boolean = false, val wifi: AutoDownloadPreset = AutoDownloadPreset.WIFI, val mobile: AutoDownloadPreset = AutoDownloadPreset.MOBILE, @@ -171,7 +170,6 @@ data class DownloadState( object DownloadSettings { private const val PREFS = "monogram_download" - private const val KEY_SPEED_UP_UPLOADS = "speed_up_uploads" private const val KEY_WIFI = "autodownload_wifi" private const val KEY_MOBILE = "autodownload_mobile" private const val KEY_ROAMING = "autodownload_roaming" @@ -200,7 +198,6 @@ object DownloadSettings { this.apply = apply val prefs = app.getSharedPreferences(PREFS, Context.MODE_PRIVATE) mutable.value = DownloadState( - speedUpUploads = prefs.getBoolean(KEY_SPEED_UP_UPLOADS, false), wifi = AutoDownloadPreset.decode(prefs.getString(KEY_WIFI, null), AutoDownloadPreset.WIFI), mobile = AutoDownloadPreset.decode(prefs.getString(KEY_MOBILE, null), AutoDownloadPreset.MOBILE), roaming = AutoDownloadPreset.decode(prefs.getString(KEY_ROAMING, null), AutoDownloadPreset.ROAMING), @@ -216,12 +213,6 @@ object DownloadSettings { fun activePreset(): AutoDownloadPreset = mutable.value.presetFor(network.value) - fun setSpeedUpUploads(enabled: Boolean) { - mutable.update { it.copy(speedUpUploads = enabled) } - persist() - apply?.invoke(mutable.value) - } - fun setPremium(premium: Boolean) { if (mutable.value.premium == premium) return mutable.update { it.copy(premium = premium) } @@ -299,7 +290,6 @@ object DownloadSettings { val prefs = appContext?.getSharedPreferences(PREFS, Context.MODE_PRIVATE) ?: return val value = mutable.value prefs.edit { - putBoolean(KEY_SPEED_UP_UPLOADS, value.speedUpUploads) putString(KEY_WIFI, value.wifi.encode()) putString(KEY_MOBILE, value.mobile.encode()) putString(KEY_ROAMING, value.roaming.encode()) From ba4a4b02f5ec88eef446c507f14473f9282f8b11 Mon Sep 17 00:00:00 2001 From: Artur Skubei <41114720+gdlbo@users.noreply.github.com> Date: Thu, 8 Oct 2026 20:17:58 +0300 Subject: [PATCH 4/6] feat(media): handle cancelled and incomplete media downloads --- .../feature/dialog/ui/MessageMedia.kt | 17 ++++++++++++++--- native/mtproto-rs/src/media/download.rs | 19 +++++++++++++++++-- 2 files changed, 31 insertions(+), 5 deletions(-) diff --git a/feature/dialog/src/main/java/org/monogram/feature/dialog/ui/MessageMedia.kt b/feature/dialog/src/main/java/org/monogram/feature/dialog/ui/MessageMedia.kt index 94f48dd0b..1de749891 100644 --- a/feature/dialog/src/main/java/org/monogram/feature/dialog/ui/MessageMedia.kt +++ b/feature/dialog/src/main/java/org/monogram/feature/dialog/ui/MessageMedia.kt @@ -202,6 +202,7 @@ fun MessageMedia( return@LaunchedEffect } val displayPriority = MediaPriority.VISIBLE + var fullCancelled = false suspend fun fetch(kind: MediaFetchKind): File? = withContext(Dispatchers.IO) { var networkTries = 0 repeat(16) { @@ -216,7 +217,10 @@ fun MessageMedia( when (result) { is Outcome.Ok -> return@withContext result.value is Outcome.Err -> { - if (result.message == "cancelled") return@withContext null + if (result.message == "cancelled") { + if (kind == MediaFetchKind.Full) fullCancelled = true + return@withContext null + } if (kind == MediaFetchKind.Display && result.message == "no display size") { return@withContext null } @@ -275,8 +279,11 @@ fun MessageMedia( if (needFull && fullFile == null) { fullFailed = false fullFile = fetch(MediaFetchKind.Full) - fullFailed = fullFile == null && (playing || requestFull) - failed = fullFile == null && thumbFile == null && (playing || requestFull) + fullFailed = fullFile == null && !fullCancelled && (playing || requestFull) + failed = fullFile == null && + thumbFile == null && + !fullCancelled && + (playing || requestFull) } } previewFetchDone = true @@ -378,6 +385,8 @@ fun MessageMedia( onCancel = { wantFull = false playing = false + fullFailed = false + failed = false mediaRepository?.cancel(fullKey) }, modifier = modifier @@ -414,6 +423,8 @@ fun MessageMedia( }, onCancel = { wantFull = false + fullFailed = false + failed = false mediaRepository?.cancel(fullKey) }, modifier = modifier diff --git a/native/mtproto-rs/src/media/download.rs b/native/mtproto-rs/src/media/download.rs index d87eca10f..5a114ff67 100644 --- a/native/mtproto-rs/src/media/download.rs +++ b/native/mtproto-rs/src/media/download.rs @@ -881,7 +881,6 @@ pub(crate) fn download_media_range_batched_streaming_capped( .map_err(|e| MtprotoError::Message(e.to_string()))?; out.write_all(&bytes) .map_err(|e| MtprotoError::Message(e.to_string()))?; - let _ = out.flush(); written_end = written_end.max(part_offset + bytes.len() as i64); if counted_offsets.insert(part_offset) { downloaded_bytes += bytes.len() as i64; @@ -931,9 +930,25 @@ pub(crate) fn download_media_range_batched_streaming_capped( if last_part || stream { break 'download; } - next = next_to_request; + // A caller may ignore the refill request returned by `on_chunk`. + // Continue at the first hole instead of the cursor those refills advanced. + next = first_missing_offset( + &counted_offsets, + offset.unwrap_or(0), + next_to_request, + chunk, + ); } } // 'download + if !stream { + if let Some(expected) = media.file_size { + if expected > 0 && written_end < expected { + return Err(MtprotoError::Message(format!( + "download incomplete: received {written_end} of {expected} bytes" + ))); + } + } + } out.set_len(written_end as u64) .map_err(|e| MtprotoError::Message(e.to_string()))?; out.flush() From 55f5858624b323c3982dc21b7d5485fd235c4a68 Mon Sep 17 00:00:00 2001 From: Artur Skubei <41114720+gdlbo@users.noreply.github.com> Date: Thu, 8 Oct 2026 22:40:44 +0300 Subject: [PATCH 5/6] feat(media): tune media downloads and show transfer speed --- .../org/monogram/core/ui/DownloadSettings.kt | 6 +- .../monogram/core/ui/DownloadSettingsTest.kt | 7 +- .../feature/dialog/ui/DialogDownloads.kt | 11 +- .../monogram/feature/dialog/ui/FileBubbles.kt | 4 +- .../feature/dialog/ui/MediaDuration.kt | 47 +++++ .../dialog/src/main/res/values-es/strings.xml | 1 + .../dialog/src/main/res/values-hy/strings.xml | 1 + .../dialog/src/main/res/values-ja/strings.xml | 1 + .../src/main/res/values-pt-rBR/strings.xml | 1 + .../dialog/src/main/res/values-ru/strings.xml | 1 + .../dialog/src/main/res/values-sk/strings.xml | 1 + .../dialog/src/main/res/values-tr/strings.xml | 1 + .../dialog/src/main/res/values-uk/strings.xml | 1 + .../dialog/src/main/res/values-zh/strings.xml | 1 + .../dialog/src/main/res/values/strings.xml | 1 + .../feature/dialog/ui/MediaDurationTest.kt | 8 + .../mtproto-rs/src/client/media_download.rs | 174 +++++++++++++----- native/mtproto-rs/src/media/download.rs | 121 ++++++++++-- native/mtproto-rs/src/request_control.rs | 21 +++ native/mtproto-rs/src/rpc.rs | 4 +- native/mtproto-rs/src/rpc/invoke.rs | 4 + native/mtproto-rs/src/transfer_policy.rs | 11 +- .../tests/unit/client_mgr_lane_tests.rs | 110 +++++++++++ .../unit/media_download_profile_tests.rs | 4 +- .../tests/unit/media_download_tests.rs | 166 +++++++++++++++++ .../tests/unit/request_control_tests.rs | 24 +++ .../unit/rpc_invoke_updates_salt_tests.rs | 79 ++++++++ .../tests/unit/transfer_policy_tests.rs | 45 +++++ 28 files changed, 769 insertions(+), 87 deletions(-) diff --git a/core/ui/src/main/java/org/monogram/core/ui/DownloadSettings.kt b/core/ui/src/main/java/org/monogram/core/ui/DownloadSettings.kt index 5881be0b0..c86f8b9d5 100644 --- a/core/ui/src/main/java/org/monogram/core/ui/DownloadSettings.kt +++ b/core/ui/src/main/java/org/monogram/core/ui/DownloadSettings.kt @@ -145,7 +145,7 @@ data class AutoDownloadPreset( } data class DownloadState( - val lanes: Int = 8, + val lanes: Int = 2, val premium: Boolean = false, val wifi: AutoDownloadPreset = AutoDownloadPreset.WIFI, val mobile: AutoDownloadPreset = AutoDownloadPreset.MOBILE, @@ -155,11 +155,11 @@ data class DownloadState( val activeNetwork: AutoDownloadNetwork = AutoDownloadNetwork.Wifi, ) { val effectiveLanes: Int get() = lanes.coerceIn(1, if (premium) 8 else 2) - val parts: Int get() = if (premium) 8 else 4 + val parts: Int get() = 4 val concurrency: DownloadConcurrency get() = DownloadConcurrency(lanes = effectiveLanes, parts = parts) val filePartKib: Int get() = 512 - val downloadChunkKib: Int get() = if (premium) 512 else 128 + val downloadChunkKib: Int get() = if (premium) 512 else 256 fun presetFor(network: AutoDownloadNetwork): AutoDownloadPreset = when (network) { AutoDownloadNetwork.Wifi -> wifi diff --git a/core/ui/src/test/java/org/monogram/core/ui/DownloadSettingsTest.kt b/core/ui/src/test/java/org/monogram/core/ui/DownloadSettingsTest.kt index 603c319cc..679270746 100644 --- a/core/ui/src/test/java/org/monogram/core/ui/DownloadSettingsTest.kt +++ b/core/ui/src/test/java/org/monogram/core/ui/DownloadSettingsTest.kt @@ -13,13 +13,14 @@ class DownloadSettingsTest { } @Test - fun premiumDownloadUsesTheFastDefaults() { + fun premiumDownloadKeepsLargeChunksWithConservativeConcurrency() { val free = DownloadState(premium = false, lanes = 99) assertEquals(DownloadConcurrency(2, 4), free.concurrency) - assertEquals(128, free.downloadChunkKib) + assertEquals(256, free.downloadChunkKib) assertEquals(512, free.filePartKib) val premium = DownloadState(premium = true, lanes = 7) - assertEquals(DownloadConcurrency(7, 8), premium.concurrency) + assertEquals(DownloadConcurrency(7, 4), premium.concurrency) + assertEquals(DownloadConcurrency(2, 4), DownloadState(premium = true).concurrency) assertEquals(512, premium.downloadChunkKib) assertEquals(512, premium.filePartKib) assertEquals(1, DownloadState(lanes = 1, premium = true).effectiveLanes) diff --git a/feature/dialog/src/main/java/org/monogram/feature/dialog/ui/DialogDownloads.kt b/feature/dialog/src/main/java/org/monogram/feature/dialog/ui/DialogDownloads.kt index d3351fb79..e8d0c61e7 100644 --- a/feature/dialog/src/main/java/org/monogram/feature/dialog/ui/DialogDownloads.kt +++ b/feature/dialog/src/main/java/org/monogram/feature/dialog/ui/DialogDownloads.kt @@ -95,15 +95,8 @@ internal fun DialogDownloadsSheet( ) val total = item.totalBytes Text( - text = if (total != null && total > 0L) { - stringResource( - R.string.dialog_media_download_progress, - formatFileSize(bytes), - formatFileSize(total), - ) - } else { - stringResource(R.string.dialog_media_loading) - }, + text = downloadProgressText(bytes, total) + ?: stringResource(R.string.dialog_media_loading), style = MaterialTheme.typography.labelMedium, color = MaterialTheme.colorScheme.onSurfaceVariant, ) diff --git a/feature/dialog/src/main/java/org/monogram/feature/dialog/ui/FileBubbles.kt b/feature/dialog/src/main/java/org/monogram/feature/dialog/ui/FileBubbles.kt index 115134738..e11aa3ce4 100644 --- a/feature/dialog/src/main/java/org/monogram/feature/dialog/ui/FileBubbles.kt +++ b/feature/dialog/src/main/java/org/monogram/feature/dialog/ui/FileBubbles.kt @@ -191,7 +191,7 @@ internal fun DocumentBubble( openFailed -> stringResource(R.string.dialog_media_open_failed) failed && file == null -> stringResource(R.string.dialog_media_failed) uploading && loading -> stringResource(R.string.dialog_media_uploading) - loading && file == null -> formatDownloadProgress(downloadedBytes, size) + loading && file == null -> downloadProgressText(downloadedBytes, size) else -> size?.takeIf { it > 0 }?.let(::formatFileSize) } sizeText?.let { @@ -329,7 +329,7 @@ internal fun AudioBubble( ) Text( when { - downloading -> formatDownloadProgress(downloadedBytes, fileSize) + downloading -> downloadProgressText(downloadedBytes, fileSize) ?: stringResource(R.string.dialog_media_loading) failedHere -> stringResource(R.string.dialog_media_failed) diff --git a/feature/dialog/src/main/java/org/monogram/feature/dialog/ui/MediaDuration.kt b/feature/dialog/src/main/java/org/monogram/feature/dialog/ui/MediaDuration.kt index 2854fc58d..816958f58 100644 --- a/feature/dialog/src/main/java/org/monogram/feature/dialog/ui/MediaDuration.kt +++ b/feature/dialog/src/main/java/org/monogram/feature/dialog/ui/MediaDuration.kt @@ -1,5 +1,19 @@ package org.monogram.feature.dialog.ui +import android.os.SystemClock +import androidx.compose.runtime.Composable +import androidx.compose.runtime.LaunchedEffect +import androidx.compose.runtime.getValue +import androidx.compose.runtime.mutableLongStateOf +import androidx.compose.runtime.remember +import androidx.compose.runtime.rememberUpdatedState +import androidx.compose.runtime.setValue +import androidx.compose.ui.res.stringResource +import androidx.lifecycle.Lifecycle +import androidx.lifecycle.compose.LocalLifecycleOwner +import androidx.lifecycle.repeatOnLifecycle +import kotlinx.coroutines.delay +import org.monogram.feature.dialog.R import org.monogram.network.http.MediaPriority import org.monogram.core.models.isStickerFileName import org.monogram.core.ui.AutoDownloadPreset @@ -28,6 +42,39 @@ internal fun formatDownloadProgress(bytes: Long, total: Long?): String? { } } +internal fun downloadBytesPerSecond(previous: Long, current: Long, elapsedMs: Long): Long = + if (elapsedMs <= 0L || current < previous) 0L + else ((current - previous).toDouble() * 1000.0 / elapsedMs).toLong() + +@Composable +internal fun downloadProgressText(bytes: Long, total: Long?, active: Boolean = true): String? { + val latestBytes by rememberUpdatedState(bytes) + var speed by remember(active) { mutableLongStateOf(0L) } + val lifecycle = LocalLifecycleOwner.current.lifecycle + LaunchedEffect(active, lifecycle) { + if (!active) return@LaunchedEffect + lifecycle.repeatOnLifecycle(Lifecycle.State.STARTED) { + speed = 0L + var previousBytes = latestBytes + var previousTime = SystemClock.elapsedRealtime() + while (true) { + delay(1000L) + val now = SystemClock.elapsedRealtime() + val received = latestBytes + speed = downloadBytesPerSecond(previousBytes, received, now - previousTime) + previousBytes = received + previousTime = now + } + } + } + val progress = formatDownloadProgress(bytes, total) ?: return null + return if (active) { + stringResource(R.string.dialog_media_progress_speed, progress, formatFileSize(speed)) + } else { + progress + } +} + internal fun formatMediaDuration(seconds: Int): String { val total = seconds.coerceAtLeast(0) val hours = total / 3600 diff --git a/feature/dialog/src/main/res/values-es/strings.xml b/feature/dialog/src/main/res/values-es/strings.xml index 61eb72b67..641da6b65 100644 --- a/feature/dialog/src/main/res/values-es/strings.xml +++ b/feature/dialog/src/main/res/values-es/strings.xml @@ -352,4 +352,5 @@ Lenguaje (opcional) Ampliar editor %1$s a las %2$s + %1$s · %2$s/s diff --git a/feature/dialog/src/main/res/values-hy/strings.xml b/feature/dialog/src/main/res/values-hy/strings.xml index a830fde39..88d163629 100644 --- a/feature/dialog/src/main/res/values-hy/strings.xml +++ b/feature/dialog/src/main/res/values-hy/strings.xml @@ -352,4 +352,5 @@ Լեզու (ըստ ցանկության) Ընդլայնել խմբագիրը %1$s, %2$s + %1$s · %2$s/վ diff --git a/feature/dialog/src/main/res/values-ja/strings.xml b/feature/dialog/src/main/res/values-ja/strings.xml index 26c7bf4f7..91ba2f50d 100644 --- a/feature/dialog/src/main/res/values-ja/strings.xml +++ b/feature/dialog/src/main/res/values-ja/strings.xml @@ -341,4 +341,5 @@ 言語（任意） エディターを展開 %1$s %2$s + %1$s · %2$s/秒 diff --git a/feature/dialog/src/main/res/values-pt-rBR/strings.xml b/feature/dialog/src/main/res/values-pt-rBR/strings.xml index 8833770f5..3047338ae 100644 --- a/feature/dialog/src/main/res/values-pt-rBR/strings.xml +++ b/feature/dialog/src/main/res/values-pt-rBR/strings.xml @@ -352,4 +352,5 @@ Linguagem (opcional) Expandir editor %1$s às %2$s + %1$s · %2$s/s diff --git a/feature/dialog/src/main/res/values-ru/strings.xml b/feature/dialog/src/main/res/values-ru/strings.xml index 70dc6dd4b..61334a505 100644 --- a/feature/dialog/src/main/res/values-ru/strings.xml +++ b/feature/dialog/src/main/res/values-ru/strings.xml @@ -374,4 +374,5 @@ Язык (необязательно) Развернуть редактор %1$s в %2$s + %1$s · %2$s/с diff --git a/feature/dialog/src/main/res/values-sk/strings.xml b/feature/dialog/src/main/res/values-sk/strings.xml index 5df657b78..be7ed354b 100644 --- a/feature/dialog/src/main/res/values-sk/strings.xml +++ b/feature/dialog/src/main/res/values-sk/strings.xml @@ -374,4 +374,5 @@ Jazyk (voliteľné) Rozbaliť editor %1$s o %2$s + %1$s · %2$s/s diff --git a/feature/dialog/src/main/res/values-tr/strings.xml b/feature/dialog/src/main/res/values-tr/strings.xml index 76c884a60..5c9aa4f7f 100644 --- a/feature/dialog/src/main/res/values-tr/strings.xml +++ b/feature/dialog/src/main/res/values-tr/strings.xml @@ -352,4 +352,5 @@ Dil (isteğe bağlı) Düzenleyiciyi genişlet %1$s saat %2$s + %1$s · %2$s/sn diff --git a/feature/dialog/src/main/res/values-uk/strings.xml b/feature/dialog/src/main/res/values-uk/strings.xml index d944fbdc7..6ca35a301 100644 --- a/feature/dialog/src/main/res/values-uk/strings.xml +++ b/feature/dialog/src/main/res/values-uk/strings.xml @@ -374,4 +374,5 @@ Мова (необов’язково) Розгорнути редактор %1$s о %2$s + %1$s · %2$s/с diff --git a/feature/dialog/src/main/res/values-zh/strings.xml b/feature/dialog/src/main/res/values-zh/strings.xml index 33213b4a4..5e5353d83 100644 --- a/feature/dialog/src/main/res/values-zh/strings.xml +++ b/feature/dialog/src/main/res/values-zh/strings.xml @@ -341,4 +341,5 @@ 语言（可选） 展开编辑器 %1$s %2$s + %1$s · %2$s/秒 diff --git a/feature/dialog/src/main/res/values/strings.xml b/feature/dialog/src/main/res/values/strings.xml index f2530403e..645bfe85d 100644 --- a/feature/dialog/src/main/res/values/strings.xml +++ b/feature/dialog/src/main/res/values/strings.xml @@ -353,4 +353,5 @@ Language (optional) Expand editor %1$s at %2$s + %1$s · %2$s/s diff --git a/feature/dialog/src/test/java/org/monogram/feature/dialog/ui/MediaDurationTest.kt b/feature/dialog/src/test/java/org/monogram/feature/dialog/ui/MediaDurationTest.kt index bf7a1364f..c9ac56f5d 100644 --- a/feature/dialog/src/test/java/org/monogram/feature/dialog/ui/MediaDurationTest.kt +++ b/feature/dialog/src/test/java/org/monogram/feature/dialog/ui/MediaDurationTest.kt @@ -163,6 +163,14 @@ class MediaDurationTest { assertEquals("2 KB", formatFileSize(2048)) } + @Test + fun downloadSpeedUsesElapsedTimeAndResetsAfterRestart() { + assertEquals(2048L, downloadBytesPerSecond(1024, 4096, 1500)) + assertEquals(0L, downloadBytesPerSecond(4096, 4096, 1000)) + assertEquals(0L, downloadBytesPerSecond(4096, 1024, 1000)) + assertEquals(0L, downloadBytesPerSecond(1024, 4096, 0)) + } + @Test fun downloadProgressUsesKnownTotal() { assertEquals(0f, downloadProgressFraction(0, 2000)!!, 0.001f) diff --git a/native/mtproto-rs/src/client/media_download.rs b/native/mtproto-rs/src/client/media_download.rs index 4342970d8..5a04f9ea3 100644 --- a/native/mtproto-rs/src/client/media_download.rs +++ b/native/mtproto-rs/src/client/media_download.rs @@ -812,12 +812,106 @@ fn invoke_lane_batch( slot: &mut Option, api_id: i32, requests: Vec, + on_chunk: &mut dyn FnMut( + usize, + Result, + ) -> Option, ) -> Result>, MtprotoError> { with_client_transport(client, slot, || { crate::api_invoke::invoke_api_batch_without_updates_streaming::< _, tellers_mtproto::latest::api::UploadFile, - >(snapshot, api_id, requests, &mut |_, _| None) + >(snapshot, api_id, requests, on_chunk) + }) +} + +pub(crate) fn invoke_parallel_media_batches( + batches: [Vec<(usize, Req)>; 2], + invoke_a: impl FnOnce( + Vec, + &mut dyn FnMut(usize, Result) -> Option, + ) -> Result>, MtprotoError> + + Send, + invoke_b: impl FnOnce( + Vec, + &mut dyn FnMut(usize, Result) -> Option, + ) -> Result>, MtprotoError> + + Send, + on_chunk: &mut dyn FnMut(usize, Result) -> Option, +) -> Result>, MtprotoError> { + use std::sync::mpsc::{Sender, sync_channel}; + + fn run_lane( + batch: Vec<(usize, Req)>, + completed: Sender<( + usize, + Result, + std::sync::mpsc::SyncSender>, + )>, + invoke: impl FnOnce( + Vec, + &mut dyn FnMut(usize, Result) -> Option, + ) -> Result>, MtprotoError>, + ) -> Result<(), MtprotoError> { + let (mut indices, requests): (Vec<_>, Vec<_>) = batch.into_iter().unzip(); + invoke(requests, &mut |index, result| { + let global_index = *indices.get(index)?; + let (reply, refill) = sync_channel(0); + completed.send((global_index, result, reply)).ok()?; + let (next_index, request) = refill.recv().ok()??; + indices.push(next_index); + Some(request) + }) + .map(|_| ()) + } + + let initial_len = batches.iter().map(Vec::len).sum(); + let mut results: Vec>> = + (0..initial_len).map(|_| None).collect(); + let failed = std::sync::atomic::AtomicBool::new(false); + std::thread::scope(|scope| { + let [a, b] = batches; + let failed_ref = &failed; + let (completed, chunks) = std::sync::mpsc::channel(); + let completed_b = completed.clone(); + let even = scope.spawn(move || { + let result = run_lane(a, completed, invoke_a); + if result.is_err() { + failed_ref.store(true, Ordering::Release); + } + result + }); + let odd = scope.spawn(move || { + let result = run_lane(b, completed_b, invoke_b); + if result.is_err() { + failed_ref.store(true, Ordering::Release); + } + result + }); + for (index, result, reply) in chunks { + if result.is_err() { + failed.store(true, Ordering::Release); + } + let refill = on_chunk(index, result.clone()) + .filter(|_| !failed.load(Ordering::Acquire)) + .map(|request| { + let next_index = results.len(); + results.push(None); + (next_index, request) + }); + results[index] = Some(result); + let _ = reply.send(refill); + } + even.join() + .map_err(|_| MtprotoError::Message("even media lane".into()))??; + odd.join() + .map_err(|_| MtprotoError::Message("odd media lane".into()))??; + Ok(results + .into_iter() + .map(|result| { + result.unwrap_or_else(|| Err(MtprotoError::Message("missing media part".into()))) + }) + .collect()) }) } @@ -850,49 +944,33 @@ fn invoke_media_batch_even_odd( } let snap_b = snap_b.unwrap(); let slot_b = slot_b.unwrap(); - let mut snap_a_t = snap_a.clone(); - let mut snap_b_t = snap_b.clone(); - let mut slot_a_t = slot_a.take(); - let mut slot_b_t = slot_b.take(); - let joined = std::thread::scope(|scope| { - let even_h = scope - .spawn(|| invoke_lane_batch(client, &mut snap_a_t, &mut slot_a_t, api_id, even_reqs)); - let odd_h = scope - .spawn(|| invoke_lane_batch(client, &mut snap_b_t, &mut slot_b_t, api_id, odd_reqs)); - (even_h.join(), odd_h.join()) - }); - *snap_a = snap_a_t; - *snap_b = snap_b_t; - *slot_a = slot_a_t; - *slot_b = slot_b_t; - let even_res = joined - .0 - .map_err(|_| MtprotoError::Message("even media lane".into()))??; - let odd_res = joined - .1 - .map_err(|_| MtprotoError::Message("odd media lane".into()))??; - let mut merged: Vec>> = - (0..requests.len()).map(|_| None).collect(); - for (slot, result) in even_idx.into_iter().zip(even_res) { - merged[slot] = Some(result); - } - for (slot, result) in odd_idx.into_iter().zip(odd_res) { - merged[slot] = Some(result); - } - let mut out = Vec::with_capacity(merged.len()); - for (index, item) in merged.into_iter().enumerate() { - let item = item.unwrap_or_else(|| Err(MtprotoError::Message("missing media part".into()))); - match &item { - Ok(file) => { - let _ = on_chunk(index, Ok(file.clone())); - } - Err(err) => { - let _ = on_chunk(index, Err(err.clone())); - } - } - out.push(item); - } - Ok(out) + let timeout = crate::rpc::rpc_timeout_secs(); + let class = crate::scheduler::current_class(); + let request_id = crate::request_control::current_id(); + let run = |snapshot: &mut tellers_mtproto_session::Snapshot, + slot: &mut Option, + requests, + on_chunk: &mut dyn FnMut(_, _) -> _| { + let _policy = crate::transfer_policy::bind(client.policy.clone()); + crate::request_control::with_shared_binding(request_id, || { + crate::scheduler::with_class(class, || { + crate::rpc::with_rpc_timeout_secs(timeout, || { + invoke_lane_batch(client, snapshot, slot, api_id, requests, on_chunk) + }) + }) + }) + }; + let result = invoke_parallel_media_batches( + [ + even_idx.into_iter().zip(even_reqs).collect(), + odd_idx.into_iter().zip(odd_reqs).collect(), + ], + |requests, on_chunk| run(snap_a, slot_a, requests, on_chunk), + |requests, on_chunk| run(snap_b, slot_b, requests, on_chunk), + on_chunk, + ); + crate::request_control::detach_sockets(); + result } pub(crate) fn download_media_range_on_lane_dc( @@ -959,7 +1037,7 @@ pub(crate) fn download_media_range_on_lane_dc( .then(|| home.auth_key.clone()) .flatten() }); - let mut io2 = if media::download_in_flight() > 1 { + let mut io2 = if should_open_second_media_lane(target_dc, media::download_in_flight()) { try_lock_second_media_lane(client) } else { None @@ -1068,6 +1146,12 @@ pub(crate) fn download_media_range_on_lane_dc( publish_media(client, home.session_id, staged, dest) } +pub(crate) fn should_open_second_media_lane(dc_id: i32, in_flight: usize) -> bool { + in_flight > 1 + && crate::scheduler::active_media_lanes() > 1 + && !crate::media::download_flood_active(dc_id) +} + fn follow_cdn_redirect( handle: u64, api_id: i32, diff --git a/native/mtproto-rs/src/media/download.rs b/native/mtproto-rs/src/media/download.rs index 5a114ff67..5199da9c4 100644 --- a/native/mtproto-rs/src/media/download.rs +++ b/native/mtproto-rs/src/media/download.rs @@ -147,7 +147,13 @@ pub(crate) fn configured_download_window(profile: DownloadProfile) -> DownloadWi let maximum = download_window(profile); let policy = crate::transfer_policy::current(); DownloadWindow { - chunk: maximum.chunk.min(policy.chunk_size()), + chunk: if policy.is_chunk_owned() { + maximum.chunk.min(policy.chunk_size()) + } else if maximum.chunk >= FAST_CHUNK { + maximum.chunk + } else { + maximum.chunk.min(policy.chunk_size()) + }, in_flight: maximum.in_flight.min(policy.pipeline_parts()), } } @@ -258,7 +264,16 @@ pub enum TransferClass { #[derive(Clone, Debug)] pub struct FloodScope { - slow_until: crate::HashMap<(i32, TransferClass), (DownloadWindow, std::time::Instant, bool)>, + slow_until: crate::HashMap<(i32, TransferClass), FloodBackoff>, +} + +#[derive(Clone, Debug)] +struct FloodBackoff { + window: DownloadWindow, + until: std::time::Instant, + premium: bool, + recover_after: std::time::Instant, + successes: usize, } impl FloodScope { @@ -278,23 +293,42 @@ impl FloodScope { now: std::time::Instant, current: DownloadWindow, ) { - let slowed = if premium { + let key = (dc_id, class); + let previous = self.slow_until.get(&key); + let slowed = if class == TransferClass::Download { + DownloadWindow { + chunk: current.chunk, + in_flight: if premium { + 1 + } else { + (previous.map_or(current.in_flight, |state| state.window.in_flight) / 2) + .clamp(1, 2) + }, + } + } else if premium { DownloadWindow { chunk: DEFAULT_CHUNK, - in_flight: 2, + in_flight: 1, } } else { - current + DownloadWindow { + chunk: current.chunk.min(DEFAULT_CHUNK), + in_flight: current.in_flight.min(4), + } }; let seconds = seconds.max(1); - let deadline = now + flood_wait_duration(seconds); - if let Some((_, existing_deadline, _)) = self.slow_until.get(&(dc_id, class)) { - if *existing_deadline > deadline { - return; - } - } - self.slow_until - .insert((dc_id, class), (slowed, deadline, premium)); + let deadline = (now + flood_wait_duration(seconds)) + .max(previous.map_or(now, |state| state.until)); + let premium = previous + .filter(|state| state.until > now) + .map_or(premium, |state| state.premium || premium); + self.slow_until.insert(key, FloodBackoff { + window: slowed, + until: deadline, + premium, + recover_after: deadline + std::time::Duration::from_secs(30), + successes: 0, + }); } pub fn window( @@ -307,25 +341,47 @@ impl FloodScope { let expired = self .slow_until .get(&(dc_id, class)) - .is_some_and(|(_, until, _)| now >= *until); - if expired { + .is_some_and(|state| now >= state.until); + if expired && class != TransferClass::Download { self.slow_until.remove(&(dc_id, class)); return fallback; } self.slow_until .get(&(dc_id, class)) - .map(|(window, _, _)| *window) + .map(|state| DownloadWindow { + chunk: state.window.chunk.min(fallback.chunk), + in_flight: state.window.in_flight.min(fallback.in_flight), + }) .unwrap_or(fallback) } + pub(crate) fn recovering_download(&self, dc_id: i32) -> bool { + self.slow_until.contains_key(&(dc_id, TransferClass::Download)) + } + + pub(crate) fn download_succeeded(&mut self, dc_id: i32, now: std::time::Instant) { + let Some(state) = self.slow_until.get_mut(&(dc_id, TransferClass::Download)) else { + return; + }; + if now < state.until { + return; + } + state.successes = state.successes.saturating_add(1); + if state.successes >= 32 && now >= state.recover_after && state.window.in_flight < 4 { + state.window.in_flight += 1; + state.successes = 0; + state.recover_after = now + std::time::Duration::from_secs(30); + } + } + pub(crate) fn pending_error( &self, dc_id: i32, class: TransferClass, now: std::time::Instant, ) -> Option { - let (_, until, premium) = self.slow_until.get(&(dc_id, class))?; - let remaining = until.checked_duration_since(now)?; + let state = self.slow_until.get(&(dc_id, class))?; + let remaining = state.until.checked_duration_since(now)?; if remaining.is_zero() { return None; } @@ -334,7 +390,7 @@ impl FloodScope { .saturating_add(u64::from(remaining.subsec_nanos() > 0)); Some(MtprotoError::Message(format!( "RPC 420: {}_{seconds}", - if *premium { + if state.premium { "FLOOD_PREMIUM_WAIT" } else { "FLOOD_WAIT" @@ -343,6 +399,13 @@ impl FloodScope { } } +pub(crate) fn download_flood_active(dc_id: i32) -> bool { + crate::transfer_policy::current() + .floods + .lock() + .recovering_download(dc_id) +} + pub(crate) fn note_transfer_flood(dc_id: i32, class: TransferClass, err: &MtprotoError) { let Some(seconds) = flood_wait_secs(err) else { return; @@ -705,6 +768,12 @@ pub(crate) fn download_media_range_batched_streaming_capped( let mut first_batch = true; let target_str = cancellation_path.display().to_string(); 'download: loop { + width = crate::transfer_policy::current().floods.lock().window( + dc_id, + TransferClass::Download, + std::time::Instant::now(), + DownloadWindow { chunk: resolve_chunk(offset), in_flight: parts_in_flight.max(1) }, + ).in_flight.min(if retried_limit { width } else { parts_in_flight.max(1) }); let chunk = resolve_chunk(Some(next).filter(|_| next > 0).or(offset)); let batch = if stream { bound_download_window() @@ -744,6 +813,7 @@ pub(crate) fn download_media_range_batched_streaming_capped( let mut written_parts = vec![false; offsets.len()]; let mut short_or_empty = vec![false; offsets.len()]; let mut last_part = false; + let mut stop_refill = false; let mut offset_by_index = offsets.clone(); let mut next_to_request = next + (batch as i64) * i64::from(chunk); @@ -756,6 +826,9 @@ pub(crate) fn download_media_range_batched_streaming_capped( return None; } let part_offset = *offset_by_index.get(index)?; + if !matches!(&res, Ok(UploadFile::UploadFile(_))) { + stop_refill = true; + } let mut hit_short = false; if let Ok(UploadFile::UploadFile(file)) = &res { let bytes = &file.bytes; @@ -767,6 +840,14 @@ pub(crate) fn download_media_range_batched_streaming_capped( written_end = written_end.max(part_offset + bytes.len() as i64); if counted_offsets.insert(part_offset) { downloaded_bytes += bytes.len() as i64; + crate::transfer_policy::current().floods.lock() + .download_succeeded(dc_id, std::time::Instant::now()); + if crate::transfer_policy::current().floods.lock().window( + dc_id, TransferClass::Download, std::time::Instant::now(), + DownloadWindow { chunk, in_flight: parts_in_flight.max(1) }, + ).in_flight != width { + stop_refill = true; + } notify_progress(&target_str, downloaded_bytes, 0); } if index < written_parts.len() { @@ -782,7 +863,7 @@ pub(crate) fn download_media_range_batched_streaming_capped( } } } - if stream || hit_short { + if stream || hit_short || last_part || stop_refill { return None; } if refill_limit.is_some_and(|limit| refills_used >= limit) { diff --git a/native/mtproto-rs/src/request_control.rs b/native/mtproto-rs/src/request_control.rs index 50d7a9eb4..3b8c3d2f1 100644 --- a/native/mtproto-rs/src/request_control.rs +++ b/native/mtproto-rs/src/request_control.rs @@ -15,6 +15,7 @@ static REQUESTS: LazyLock>>> = LazyLock::new(|| Mutex::new(HashMap::new())); thread_local! { static CURRENT: RefCell<(u64, Option>)> = const { RefCell::new((0, None)) }; + static SHARED: std::cell::Cell = const { std::cell::Cell::new(false) }; } pub fn create() -> u64 { @@ -39,6 +40,22 @@ pub fn bind(id: u64) -> u64 { CURRENT.with(|cell| cell.replace((id, control)).0) } +pub(crate) fn current_id() -> u64 { + CURRENT.with(|cell| cell.borrow().0) +} + +pub(crate) fn with_shared_binding(id: u64, body: impl FnOnce() -> T) -> T { + struct Restore(u64, bool); + impl Drop for Restore { + fn drop(&mut self) { + bind(self.0); + SHARED.with(|cell| cell.set(self.1)); + } + } + let _restore = Restore(bind(id), SHARED.with(|cell| cell.replace(true))); + body() +} + pub fn cancel(id: u64) { let control = REQUESTS.lock().get(&id).cloned(); if let Some(control) = control { @@ -51,6 +68,10 @@ pub fn release(id: u64) { } pub(crate) fn detach_sockets() { + // A finished parallel lane must not unregister the other lane's socket. + if SHARED.with(|cell| cell.get()) { + return; + } CURRENT.with(|cell| { if let Some(control) = &cell.borrow().1 { control.detach_sockets(); diff --git a/native/mtproto-rs/src/rpc.rs b/native/mtproto-rs/src/rpc.rs index d2f5dc41e..961607760 100644 --- a/native/mtproto-rs/src/rpc.rs +++ b/native/mtproto-rs/src/rpc.rs @@ -22,6 +22,7 @@ pub use inbound::{ pub(crate) use inbound::{gzip_if_smaller, ungzip_if_needed}; pub use invoke::invoke_raw; pub use live::{LiveTransport, drop_live_transport, with_live_transport}; +pub(crate) use timeout::rpc_timeout_secs; pub use timeout::with_rpc_timeout_secs; pub(crate) use invoke::{ @@ -81,8 +82,7 @@ pub(crate) use timeout::{ KEEPALIVE_PING_SECS, fail_fast_idle, idle_after_complete_frames, idle_empty_first_byte, idle_needs_liveness_probe, idle_reused_socket, is_mid_frame, keepalive_probe_failed, leftover_frame_grace, live_transport_stale, note_inbound_liveness, recv_wait_deadline, - rpc_attempt_budget, rpc_timeout_message, rpc_timeout_secs, subscribed_read_deadline, - timeout_idle_needs_probe, + rpc_attempt_budget, rpc_timeout_message, subscribed_read_deadline, timeout_idle_needs_probe, }; #[cfg(test)] diff --git a/native/mtproto-rs/src/rpc/invoke.rs b/native/mtproto-rs/src/rpc/invoke.rs index 1f4477eb5..4a35e0a1e 100644 --- a/native/mtproto-rs/src/rpc/invoke.rs +++ b/native/mtproto-rs/src/rpc/invoke.rs @@ -637,6 +637,7 @@ pub(crate) fn invoke_batch_until_results_streaming bool { + self.chunk_owned.load(Ordering::Relaxed) + } + pub(crate) fn own_file_part_kib(&self, kib: i32) { self.file_part_owned.store(true, Ordering::Relaxed); self.assign_file_part_kib(kib); diff --git a/native/mtproto-rs/tests/unit/client_mgr_lane_tests.rs b/native/mtproto-rs/tests/unit/client_mgr_lane_tests.rs index 909d35856..822ce12e9 100644 --- a/native/mtproto-rs/tests/unit/client_mgr_lane_tests.rs +++ b/native/mtproto-rs/tests/unit/client_mgr_lane_tests.rs @@ -1,4 +1,21 @@ use super::*; +use crate::media; + +#[test] +fn configured_single_media_lane_disables_second_download_socket() { + let policy = std::sync::Arc::new(crate::transfer_policy::TransferPolicy::stock()); + let _guard = crate::transfer_policy::bind(policy.clone()); + policy.own_media_lanes(1); + assert!(!should_open_second_media_lane(2, 8)); + policy.own_media_lanes(2); + assert!(should_open_second_media_lane(2, 8)); + policy.floods.lock().park(2, media::TransferClass::Download, 1, false, + std::time::Instant::now() - std::time::Duration::from_secs(5), + media::DownloadWindow { chunk: media::FAST_CHUNK, in_flight: 8 }); + assert!(media::check_transfer_flood(2, media::TransferClass::Download).is_ok()); + assert!(!should_open_second_media_lane(2, 8)); + assert!(should_open_second_media_lane(4, 8)); +} #[test] fn download_sessions_grow_from_one_to_eight_and_reset_after_idle() { @@ -540,6 +557,99 @@ fn even_odd_media_offsets_split_by_chunk() { assert_eq!(odd, vec![1, 3]); } +#[test] +fn parallel_media_refills_fast_lane_before_slow_lane_completes() { + use std::sync::mpsc::channel; + use std::time::Duration; + let (refilled, wait_for_refill) = channel(); + let caller = std::thread::current().id(); + let mut delivered = Vec::new(); + let results = media_download::invoke_parallel_media_batches( + [vec![(0, 0usize)], vec![(1, 1usize)]], + |mut requests, on_chunk| { + let mut results = Vec::new(); + let mut index = 0; + while index < requests.len() { + if index == 1 { + refilled.send(()).unwrap(); + } + let value = requests[index]; + if let Some(next) = on_chunk(index, Ok(value)) { + requests.push(next); + } + results.push(Ok(value)); + index += 1; + } + Ok(results) + }, + move |requests, on_chunk| { + wait_for_refill + .recv_timeout(Duration::from_secs(2)) + .expect("fast lane must refill while slow lane is outstanding"); + assert_eq!(requests, vec![1]); + assert!(on_chunk(0, Ok(1)).is_none()); + Ok(vec![Ok(1)]) + }, + &mut |index, result| { + assert_eq!(std::thread::current().id(), caller); + delivered.push((index, result.unwrap())); + (index == 0).then_some(2) + }, + ) + .unwrap(); + assert_eq!( + results.into_iter().collect::, _>>().unwrap(), + vec![0, 1, 2] + ); + delivered.sort_unstable(); + assert_eq!(delivered, vec![(0, 0), (1, 1), (2, 2)]); +} + +#[test] +fn parallel_media_preserves_lane_error_and_delivers_other_lane() { + let mut delivered = Vec::new(); + let result = media_download::invoke_parallel_media_batches( + [vec![(0, 0usize)], vec![(1, 1usize)]], + |_, _| Err(MtprotoError::Message("lane failed".into())), + |_, on_chunk| { + assert!(on_chunk(0, Ok(1usize)).is_none()); + Ok(vec![Ok(1)]) + }, + &mut |index, result| { + delivered.push((index, result.unwrap())); + None + }, + ); + assert!(matches!(result, Err(MtprotoError::Message(message)) if message == "lane failed")); + assert_eq!(delivered, vec![(1, 1)]); +} + +#[test] +fn parallel_media_does_not_refill_when_callback_stops() { + let invoke = + |requests: Vec, + on_chunk: &mut dyn FnMut(usize, Result) -> Option| { + for (index, value) in requests.iter().copied().enumerate() { + assert!(on_chunk(index, Ok(value)).is_none()); + } + Ok(requests.into_iter().map(Ok).collect()) + }; + let mut delivered = Vec::new(); + let results = media_download::invoke_parallel_media_batches( + [vec![(0, 0)], vec![(1, 1)]], + invoke, + invoke, + &mut |index, result| { + delivered.push((index, result.unwrap())); + None + }, + ) + .unwrap(); + assert_eq!(results.len(), 2); + delivered.sort_unstable(); + assert_eq!(delivered, vec![(0, 0), (1, 1)]); +} + #[test] fn extra_main_sessions_stay_zero_without_tmp_sessions() { assert_eq!(crate::scheduler::extra_main_sessions(), 0); diff --git a/native/mtproto-rs/tests/unit/media_download_profile_tests.rs b/native/mtproto-rs/tests/unit/media_download_profile_tests.rs index 07b3c60e9..3dc01eb9b 100644 --- a/native/mtproto-rs/tests/unit/media_download_profile_tests.rs +++ b/native/mtproto-rs/tests/unit/media_download_profile_tests.rs @@ -118,7 +118,7 @@ fn flood_parks_one_dc_and_class_for_the_server_interval() { scope .window(2, TransferClass::Download, now, fast) .in_flight, - 2 + 1 ); assert_eq!(scope.window(4, TransferClass::Download, now, fast), fast); assert_eq!( @@ -126,5 +126,5 @@ fn flood_parks_one_dc_and_class_for_the_server_interval() { ordinary ); let later = now + flood_wait_duration(120); - assert_eq!(scope.window(2, TransferClass::Download, later, fast), fast); + assert_eq!(scope.window(2, TransferClass::Download, later, fast).in_flight, 1); } diff --git a/native/mtproto-rs/tests/unit/media_download_tests.rs b/native/mtproto-rs/tests/unit/media_download_tests.rs index 53b36c3d7..ea14d259e 100644 --- a/native/mtproto-rs/tests/unit/media_download_tests.rs +++ b/native/mtproto-rs/tests/unit/media_download_tests.rs @@ -749,6 +749,96 @@ fn completed_part_refills_inflight_window_without_waiting_for_the_batch() { fs::remove_file(&path).ok(); } +#[test] +fn parallel_media_download_preserves_bytes_with_out_of_order_refills() { + let _globals = hold_media_globals(); + let _profile = bind_download_profile(DownloadProfile::Ordinary); + use tellers_mtproto::latest::api::{ + StorageFileType, StorageFileUnknownConstructor, UploadFileConstructor, UploadGetFileRequest, + }; + let media = MediaRef { + file_size: None, + kind: "video".into(), + cache_key: "parallel-refill".into(), + location: MediaLocation::Document { + id: 31, + access_hash: 32, + file_reference: vec![7], + thumb_size: String::new(), + dc_id: 2, + mime_type: "video/mp4".into(), + }, + thumb_cache_key: None, + thumb_location: None, + display_cache_key: None, + display_location: None, + sticker_set_id: None, + sticker_set_access_hash: None, + source_url: None, + }; + let path = + std::env::temp_dir().join(format!("monogram-parallel-refill-{}", std::process::id())); + let mut batches = 0; + let mut file_size = 0; + let result = download_media_range_batched_streaming( + 2, + &media, + &path, + &path, + None, + 4, + |_, requests, on_chunk| { + batches += 1; + file_size = requests[0].limit as usize * 7 + 123; + let mut lanes = [Vec::new(), Vec::new()]; + for (index, request) in requests.into_iter().enumerate() { + lanes[index % 2].push((index, request)); + } + let invoke = |mut requests: Vec, + on_chunk: &mut dyn FnMut( + usize, + Result, + ) -> Option| { + let mut results = Vec::new(); + let mut index = 0; + while index < requests.len() { + let request = &requests[index]; + if request.offset == 0 { + std::thread::sleep(std::time::Duration::from_millis(20)); + } + let start = request.offset as usize; + let len = file_size.saturating_sub(start).min(request.limit as usize); + let file = UploadFile::UploadFile(UploadFileConstructor { + type_: Box::new(StorageFileType::StorageFileUnknown( + StorageFileUnknownConstructor {}, + )), + mtime: 0, + bytes: (start..start + len) + .map(|offset| (offset % 251) as u8) + .collect(), + }); + if let Some(next) = on_chunk(index, Ok(file.clone())) { + requests.push(next); + } + results.push(Ok(file)); + index += 1; + } + Ok(results) + }; + crate::client::invoke_parallel_media_batches(lanes, invoke, invoke, on_chunk) + }, + ); + assert!(result.is_ok(), "parallel download failed: {result:?}"); + assert_eq!(batches, 1); + assert_eq!( + fs::read(&path).unwrap(), + (0..file_size) + .map(|offset| (offset % 251) as u8) + .collect::>() + ); + fs::remove_file(path).unwrap(); +} + #[test] fn one_mib_chunks_are_rejected() { let _globals = hold_media_globals(); @@ -818,9 +908,85 @@ fn flood_wait_returns_without_sleeping_or_retrying_on_the_lane() { ); assert_eq!(calls, 1); assert!(!path.exists()); + let policy = crate::transfer_policy::current(); + *policy.floods.lock() = FloodScope::new(); + policy.floods.lock().park(2, TransferClass::Download, 2, false, + std::time::Instant::now() - std::time::Duration::from_secs(5), + DownloadWindow { chunk: FAST_CHUNK, in_flight: 8 }); + assert!(check_transfer_flood(2, TransferClass::Download).is_ok()); + download_media_range_batched(2, &media, &path, &path, None, 8, |_, requests| { + assert_eq!(requests.len(), 2, "retry must keep the reduced window"); + Ok(requests.into_iter().map(|request| { + Ok(UploadFile::UploadFile(UploadFileConstructor { + type_: Box::new(StorageFileType::StorageFileUnknown(StorageFileUnknownConstructor {})), + mtime: 0, + bytes: vec![9u8; (20i64 - request.offset).max(0) as usize], + })) + }).collect()) + }).unwrap(); + assert_eq!(fs::read(&path).unwrap(), vec![9u8; 20]); fs::remove_file(&path).ok(); } +#[test] +fn ordinary_flood_wait_reduces_aggressive_download_window() { + let now = std::time::Instant::now(); + let current = DownloadWindow { + chunk: FAST_CHUNK, + in_flight: 8, + }; + let mut scope = FloodScope::new(); + + scope.park(2, TransferClass::Download, 2, false, now, current); + assert_eq!( + scope.window(2, TransferClass::Download, now, current), + DownloadWindow { + chunk: FAST_CHUNK, + in_flight: 2, + } + ); + + let mut premium_scope = FloodScope::new(); + premium_scope.park(2, TransferClass::Download, 2, true, now, current); + assert_eq!( + premium_scope.window(2, TransferClass::Download, now, current), + DownloadWindow { + chunk: FAST_CHUNK, + in_flight: 1, + } + ); +} + +#[test] +fn flood_retry_keeps_window_after_deadline_and_recovers_only_with_success() { + let now = std::time::Instant::now(); + let fast = DownloadWindow { chunk: FAST_CHUNK, in_flight: 8 }; + let mut scope = FloodScope::new(); + scope.park(2, TransferClass::Download, 2, false, now, fast); + let retry = now + std::time::Duration::from_secs(3); + assert!(scope.pending_error(2, TransferClass::Download, retry).is_none()); + assert_eq!(scope.window(2, TransferClass::Download, retry, fast).in_flight, 2); + assert!(scope.recovering_download(2)); + assert!(!scope.recovering_download(4)); + scope.park(2, TransferClass::Download, 2, false, retry, fast); + let ready = retry + std::time::Duration::from_secs(2); + assert_eq!(scope.window(2, TransferClass::Download, ready, fast).in_flight, 1); + let later = ready + std::time::Duration::from_secs(31); + assert_eq!(scope.window(2, TransferClass::Download, later, fast).in_flight, 1); + for _ in 0..31 { scope.download_succeeded(2, later); } + assert_eq!(scope.window(2, TransferClass::Download, later, fast).in_flight, 1); + scope.download_succeeded(2, later); + assert_eq!(scope.window(2, TransferClass::Download, later, fast).in_flight, 2); + for _ in 0..32 { scope.download_succeeded(2, later); } + assert_eq!(scope.window(2, TransferClass::Download, later, fast).in_flight, 2); + let stable = later + std::time::Duration::from_secs(30); + scope.download_succeeded(2, stable); + assert_eq!(scope.window(2, TransferClass::Download, stable, fast).in_flight, 3); + assert_eq!(scope.window(4, TransferClass::Download, stable, fast), fast); + let small = DownloadWindow { chunk: DEFAULT_CHUNK, in_flight: 1 }; + assert_eq!(scope.window(2, TransferClass::Download, stable, small), small); +} + #[test] fn rpc_timeout_retries_from_missing_offset_without_dropping_written_parts() { let _globals = hold_media_globals(); diff --git a/native/mtproto-rs/tests/unit/request_control_tests.rs b/native/mtproto-rs/tests/unit/request_control_tests.rs index 7b0c447ae..df72d5958 100644 --- a/native/mtproto-rs/tests/unit/request_control_tests.rs +++ b/native/mtproto-rs/tests/unit/request_control_tests.rs @@ -1,4 +1,28 @@ use super::*; + +#[test] +fn finished_parallel_lane_keeps_other_socket_cancellable() { + use std::io::Read; + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let socket = Arc::new(TcpStream::connect(listener.local_addr().unwrap()).unwrap()); + let (mut server, _) = listener.accept().unwrap(); + server + .set_read_timeout(Some(std::time::Duration::from_secs(1))) + .unwrap(); + let id = create(); + let previous = bind(id); + register(&socket).unwrap(); + std::thread::scope(|scope| { + scope + .spawn(|| with_shared_binding(id, detach_sockets)) + .join() + .unwrap(); + }); + cancel(id); + assert_eq!(server.read(&mut [0; 1]).unwrap(), 0); + bind(previous); + release(id); +} #[test] fn cancellation_is_scoped_and_binding_restores() { let first = create(); diff --git a/native/mtproto-rs/tests/unit/rpc_invoke_updates_salt_tests.rs b/native/mtproto-rs/tests/unit/rpc_invoke_updates_salt_tests.rs index 7e4beb9a3..9a0183418 100644 --- a/native/mtproto-rs/tests/unit/rpc_invoke_updates_salt_tests.rs +++ b/native/mtproto-rs/tests/unit/rpc_invoke_updates_salt_tests.rs @@ -1,5 +1,84 @@ use super::updates_session_should_abort; +#[test] +fn streaming_refill_is_sent_before_waiting_for_another_packet() { + use tellers_mtproto_crypto::{Direction, MessageToEncrypt, encrypt_message}; + use tellers_mtproto_engine::{Engine, ExponentialBackoff}; + use tellers_mtproto_session::{Clock, OsRandom, Snapshot}; + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let mut transport = super::super::framing::open_live_addr( + 2, + &listener.local_addr().unwrap().to_string(), + None, + 1, + ) + .unwrap(); + let (_peer, _) = listener.accept().unwrap(); + let clock = super::SystemClock; + let mut snapshot = Snapshot::new(2, &mut OsRandom).unwrap(); + snapshot.auth_key = Some(vec![7; 256]); + let server_id = ((clock.unix_micros() / 1_000_000) << 32) | 1; + snapshot.last_message_id = server_id + (1_i64 << 32) - 1; + let request_id = snapshot.last_message_id + 4; + let mut body = super::super::inbound::RPC_RESULT.to_le_bytes().to_vec(); + body.extend_from_slice(&request_id.to_le_bytes()); + body.extend_from_slice(&[1; 4]); + let packet = encrypt_message( + snapshot.auth_key.as_ref().unwrap(), + MessageToEncrypt { + server_salt: 0, + session_id: snapshot.session_id, + message_id: server_id, + sequence: 1, + body: &body, + padding: &[3; 16], + }, + Direction::ServerToClient, + ) + .unwrap(); + transport + .input + .extend_from_slice(&(packet.len() as u32).to_le_bytes()); + transport.input.extend_from_slice(&packet); + let mut engine = Engine::new( + snapshot, + ExponentialBackoff { + timeout_micros: 5_000_000, + initial_delay_micros: 0, + max_attempts: 1, + }, + ) + .unwrap(); + let deadline = std::time::Instant::now() + std::time::Duration::from_millis(150); + let mut callbacks = 0; + let result = super::invoke_batch_until_results_streaming( + &mut engine, + &mut transport, + &[vec![1; 4]], + &clock, + deadline, + deadline, + deadline, + false, + &mut false, + |index, response| { + assert_eq!(index, 0); + assert_eq!(response.unwrap(), &[1; 4]); + callbacks += 1; + Some(vec![2; 4]) + }, + ); + assert!( + result.is_err(), + "peer deliberately withholds the refill response" + ); + assert_eq!(callbacks, 1); + assert!( + engine.next_outbound().is_none(), + "refill must be sent before receive blocks" + ); +} + #[test] fn salt_retry_does_not_abort_updates_session() { assert!(!updates_session_should_abort(true)); diff --git a/native/mtproto-rs/tests/unit/transfer_policy_tests.rs b/native/mtproto-rs/tests/unit/transfer_policy_tests.rs index fe8abb7bb..b9889d6ef 100644 --- a/native/mtproto-rs/tests/unit/transfer_policy_tests.rs +++ b/native/mtproto-rs/tests/unit/transfer_policy_tests.rs @@ -52,6 +52,51 @@ fn download_profiles_and_uploads_respect_client_settings() { assert_eq!(thumb.in_flight, 1); } +#[test] +fn flood_recovery_serializes_downloads_across_file_queues_on_one_dc() { + let policy = Arc::new(TransferPolicy::stock()); + policy.floods.lock().park(2, crate::media::TransferClass::Download, 1, false, + std::time::Instant::now() - std::time::Duration::from_secs(5), + crate::media::DownloadWindow { chunk: crate::media::FAST_CHUNK, in_flight: 8 }); + let held = policy.admit_download(2, true).unwrap(); + let other_dc = policy.admit_download(4, false).unwrap(); + let (ready_tx, ready_rx) = std::sync::mpsc::channel(); + let waiting = policy.clone(); + let thread = std::thread::spawn(move || { + let _lease = waiting.admit_download(2, false).unwrap(); + ready_tx.send(()).unwrap(); + }); + assert!(ready_rx.recv_timeout(std::time::Duration::from_millis(100)).is_err()); + drop(held); + ready_rx.recv_timeout(std::time::Duration::from_secs(2)).unwrap(); + thread.join().unwrap(); + drop(other_dc); +} + +#[test] +fn unowned_policy_keeps_large_chunks_but_respects_pipeline_limit() { + let policy = Arc::new(TransferPolicy::stock()); + let _guard = transfer_policy::bind(policy.clone()); + + for profile in [ + crate::media::DownloadProfile::Visible, + crate::media::DownloadProfile::Playing, + crate::media::DownloadProfile::User, + ] { + let window = crate::media::configured_download_window(profile); + assert_eq!(window.chunk, crate::media::FAST_CHUNK); + assert_eq!(window.in_flight, 4); + policy.assign_pipeline_parts(2); + assert_eq!(crate::media::configured_download_window(profile).in_flight, 2); + policy.assign_pipeline_parts(4); + } + + let ordinary = + crate::media::configured_download_window(crate::media::DownloadProfile::Ordinary); + assert_eq!(ordinary.chunk, crate::media::DEFAULT_CHUNK); + assert_eq!(ordinary.in_flight, 4); +} + #[test] fn ffi_media_dispatch_profiles_match_kotlin_values_and_reset_after_request() { for (class, profile) in [ From e37acd171512d50495da3ef9d501f8f5e90b9b0d Mon Sep 17 00:00:00 2001 From: Artur Skubei <41114720+gdlbo@users.noreply.github.com> Date: Thu, 8 Oct 2026 23:04:11 +0300 Subject: [PATCH 6/6] feat(mtproto): overlap second media lane preparation --- .../mtproto-rs/src/client/media_download.rs | 101 ++++++++++++++---- 1 file changed, 78 insertions(+), 23 deletions(-) diff --git a/native/mtproto-rs/src/client/media_download.rs b/native/mtproto-rs/src/client/media_download.rs index 5a04f9ea3..5ab45977b 100644 --- a/native/mtproto-rs/src/client/media_download.rs +++ b/native/mtproto-rs/src/client/media_download.rs @@ -28,6 +28,32 @@ fn download_temp_ready(io: &super::lanes::SessionIo, dc: i32, home_key: Option<& .is_some_and(|key| Some(key) != home_key) } +pub(crate) struct MediaLanePreparation { + request_id: u64, + worker: std::thread::JoinHandle, +} + +impl MediaLanePreparation { + pub(crate) fn start(prepare: impl FnOnce() -> bool + Send + 'static) -> Self { + let request_id = crate::request_control::create(); + let worker = std::thread::spawn(move || { + struct Release(u64); + impl Drop for Release { + fn drop(&mut self) { crate::request_control::release(self.0); } + } + let _release = Release(request_id); + crate::request_control::with_shared_binding(request_id, prepare) + }); + Self { request_id, worker } + } + + pub(crate) fn finished(&self) -> bool { self.worker.is_finished() } +} + +impl Drop for MediaLanePreparation { + fn drop(&mut self) { crate::request_control::cancel(self.request_id); } +} + pub fn download_wallpaper( handle: u64, id: i64, @@ -989,7 +1015,7 @@ pub(crate) fn download_media_range_on_lane_dc( if !session_lease_valid(&client.data.lock(), home.session_id) { return Err(expired_session_lease()); } - let mut lane_perm = None; + let mut first_chunk_span = Some(crate::perf::span("media_first_chunk")); let now = crate::auth::bind_temp::unix_now(); let already_temp = io.temp_expires_at > now && io.snapshot.dc_id == target_dc @@ -1010,7 +1036,6 @@ pub(crate) fn download_media_range_on_lane_dc( } if let Some(key) = io.snapshot.auth_key.clone() { io.perm_key = Some(key.clone()); - lane_perm = Some(key.clone()); let offset = io.snapshot.time_offset_micros; let dc = io.snapshot.dc_id; if let Ok(bound) = crate::auth::bind_temp::negotiate_bound_temp(&key, dc, offset) { @@ -1032,7 +1057,7 @@ pub(crate) fn download_media_range_on_lane_dc( offset, ); } - let bind_key = lane_perm.or(io.perm_key.clone()).or_else(|| { + let bind_key = io.perm_key.clone().or_else(|| { (target_dc == home.dc_id) .then(|| home.auth_key.clone()) .flatten() @@ -1042,31 +1067,43 @@ pub(crate) fn download_media_range_on_lane_dc( } else { None }; - let second_ready = match (io2.as_mut(), bind_key.as_deref()) { - (Some(second), _) if download_temp_ready(second, target_dc, home.auth_key.as_deref()) => { + let second_ready = match io2.as_mut() { + Some(second) if download_temp_ready(second, target_dc, home.auth_key.as_deref()) => { true } - (Some(second), Some(key)) => { - match crate::auth::bind_temp::negotiate_bound_temp( - key, - target_dc, - home.time_offset_micros, - ) { - Ok(bound) => { - second.perm_key = Some(key.to_vec()); - second.snapshot = bound.snapshot; - second.temp_expires_at = bound.expires_at; - second.transport = None; - true - } - Err(_) => false, - } - } _ => false, }; if !second_ready { io2 = None; } + let mut preparation = if io2.is_none() && offset.is_none() + && should_open_second_media_lane(target_dc, media::download_in_flight()) { + bind_key.and_then(|key| get_client(handle).ok().map(|worker_client| { + let class = scheduler::current_class(); + let home_key = home.auth_key.clone(); + let time_offset = home.time_offset_micros; + let expected_session = home.session_id; + MediaLanePreparation::start(move || { + let _policy = crate::transfer_policy::bind(worker_client.policy.clone()); + scheduler::with_class(class, || { + tcp::with_connection_control(&worker_client.connections, || { + let _span = crate::perf::span("media_second_lane_prepare"); + let Some(mut second) = try_lock_second_media_lane(&worker_client) else { return false; }; + if !session_lease_valid(&worker_client.data.lock(), expected_session) { return false; } + if download_temp_ready(&second, target_dc, home_key.as_deref()) { return true; } + let Ok(bound) = crate::auth::bind_temp::negotiate_bound_temp(&key, target_dc, time_offset) else { return false; }; + if crate::request_control::check().is_err() + || !session_lease_valid(&worker_client.data.lock(), expected_session) { return false; } + second.perm_key = Some(key); + second.snapshot = bound.snapshot; + second.temp_expires_at = bound.expires_at; + second.transport = None; + true + }) + }) + }) + })) + } else { None }; let mut media_snap = io.snapshot.clone(); let mut slot = io.transport.take(); let mut media_snap2 = io2.as_ref().map(|lane| lane.snapshot.clone()); @@ -1075,15 +1112,32 @@ pub(crate) fn download_media_range_on_lane_dc( // Timeout is per window, not the whole file. Streaming idle-extends while parts arrive. let path = { let lane_dc = media_snap.dc_id; - media::download_media_range_batched_streaming( + media::download_media_range_batched_streaming_capped( lane_dc, media, staged.path(), dest, offset, media::download_in_flight(), + None, |_init_first, requests, on_chunk| { media::check_transfer_flood(lane_dc, media::TransferClass::Download)?; + if preparation.as_ref().is_some_and(MediaLanePreparation::finished) { + preparation = None; + if should_open_second_media_lane(lane_dc, media::download_in_flight()) { + io2 = try_lock_second_media_lane(client).filter(|lane| + download_temp_ready(lane, lane_dc, home.auth_key.as_deref())); + media_snap2 = io2.as_ref().map(|lane| lane.snapshot.clone()); + slot2 = io2.as_mut().map(|lane| lane.transport.take()); + } + } + let mut stream_chunk = |index, result: Result| { + if matches!(&result, Ok(tellers_mtproto::latest::api::UploadFile::UploadFile(file)) if !file.bytes.is_empty()) { + first_chunk_span.take(); + } + let refill = on_chunk(index, result); + if preparation.as_ref().is_some_and(MediaLanePreparation::finished) { None } else { refill } + }; crate::rpc::with_rpc_timeout_secs(45, || { invoke_media_batch_even_odd( client, @@ -1093,12 +1147,13 @@ pub(crate) fn download_media_range_on_lane_dc( slot2.as_mut(), api_id, requests, - on_chunk, + &mut stream_chunk, ) }) }, ) }; + drop(preparation); io.snapshot = media_snap.clone(); // FILE_REFERENCE / transport errors must not park this TCP for the next file. io.transport = if path.is_ok() { slot } else { None };