SoulStream/crates/soulseek-client/src/upload.rs

//! Upload servicing: when a peer requests one of our shared files, we send
//! it over an F connection. Slot-limited with a simple FIFO queue.
//!
//! Two protocol sub-flows, mirroring the download side (M2):
//!
//! * **Immediate** — peer sends `TransferRequest(download, their_token)`.
//!   With a free slot we reply `TransferResponse(their_token, allowed)`,
//!   then await *their* inbound F connection (keyed by their token), read
//!   the 8-byte offset, and stream the file.
//! * **Queued / proactive** — no slot yet (or a bare `QueueDownload`). We
//!   enqueue; when a slot frees we send `TransferRequest(upload, our_token)`,
//!   await their `TransferResponse`, open the F connection ourselves, read
//!   the offset, and stream.

use std::path::PathBuf;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt};
use tokio::net::TcpStream;
use tokio::sync::oneshot;

use soulseek_proto::peer;

use crate::{conn, ClientError, Event, Inner, Result};

const ACK_TIMEOUT: Duration = Duration::from_secs(30);
const STALL_TIMEOUT: Duration = Duration::from_secs(60);

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum UploadState {
    Queued,
    Transferring,
    Complete,
    Failed { reason: String },
}

#[derive(Debug, Clone)]
pub struct UploadUpdate {
    pub id: u64,
    pub username: String,
    pub filename: String,
    pub state: UploadState,
    pub sent: u64,
    pub total: u64,
}

/// Peer sent `TransferRequest(download, their_token, filename)`.
pub(crate) fn on_transfer_request(
    inner: Arc<Inner>,
    username: String,
    their_token: u32,
    remote_path: String,
) {
    let Some((local, size)) = inner.shares.resolve(&remote_path) else {
        reply(
            &inner,
            &username,
            peer::transfer_response_rejected(their_token, "File not shared"),
        );
        return;
    };

    if inner.has_free_upload_slot() {
        // Accept: they will open the F connection with their token.
        reply(
            &inner,
            &username,
            peer::transfer_response_allowed(their_token, size),
        );
        spawn_immediate(inner, username, remote_path, local, size, their_token);
    } else {
        reply(
            &inner,
            &username,
            peer::transfer_response_rejected(their_token, "Queued"),
        );
        enqueue(&inner, &username, &remote_path);
    }
}

/// Peer sent a bare `QueueDownload(filename)` (no token) — enqueue, or start
/// the proactive flow immediately if a slot is free.
pub(crate) fn on_queue_download(inner: Arc<Inner>, username: String, remote_path: String) {
    let Some((local, size)) = inner.shares.resolve(&remote_path) else {
        reply(
            &inner,
            &username,
            peer::upload_denied(&remote_path, "File not shared"),
        );
        return;
    };
    if inner.has_free_upload_slot() {
        spawn_proactive(inner, username, remote_path, local, size);
    } else {
        enqueue(&inner, &username, &remote_path);
    }
}

/// Report queue place when a peer asks (PlaceInQueueRequest).
pub(crate) fn on_place_query(inner: &Arc<Inner>, username: &str, remote_path: &str) {
    let place = inner
        .upload_queue
        .lock()
        .unwrap()
        .iter()
        .position(|(u, p)| u == username && p == remote_path)
        .map(|i| i as u32 + 1)
        .unwrap_or(0);
    reply(
        inner,
        username,
        peer::place_in_queue_response(remote_path, place),
    );
}

fn enqueue(inner: &Arc<Inner>, username: &str, remote_path: &str) {
    let mut q = inner.upload_queue.lock().unwrap();
    if !q.iter().any(|(u, p)| u == username && p == remote_path) && q.len() < 500 {
        q.push_back((username.to_string(), remote_path.to_string()));
    }
    let place = q.len() as u32;
    drop(q);
    reply(
        inner,
        username,
        peer::place_in_queue_response(remote_path, place),
    );
}

fn reply(inner: &Arc<Inner>, username: &str, frame: Vec<u8>) {
    if let Some(h) = inner.peers.lock().unwrap().get(username) {
        let _ = h.write_tx.send(frame);
    }
}

/// After a slot frees, start the next queued upload (proactive flow).
fn pump_queue(inner: Arc<Inner>) {
    if !inner.has_free_upload_slot() {
        return;
    }
    let next = inner.upload_queue.lock().unwrap().pop_front();
    if let Some((username, remote_path)) = next {
        if let Some((local, size)) = inner.shares.resolve(&remote_path) {
            spawn_proactive(inner, username, remote_path, local, size);
        }
    }
}

// -------- immediate flow (peer opens the F connection) --------

fn spawn_immediate(
    inner: Arc<Inner>,
    username: String,
    remote_path: String,
    local: PathBuf,
    size: u64,
    their_token: u32,
) {
    inner.uploads_active.fetch_add(1, Ordering::Relaxed);
    let id = inner.upload_counter.fetch_add(1, Ordering::Relaxed);

    // Register to receive their inbound F connection (keyed by their token).
    let (conn_tx, conn_rx) = oneshot::channel();
    inner
        .expected_transfers
        .lock()
        .unwrap()
        .insert((username.clone(), their_token), conn_tx);

    tokio::spawn(async move {
        let user_err = username.clone();
        run_upload(
            &inner,
            id,
            &username,
            &remote_path,
            &local,
            size,
            async move {
                let stream = tokio::time::timeout(ACK_TIMEOUT, conn_rx)
                    .await
                    .map_err(|_| ClientError::Timeout)?
                    .map_err(|_| ClientError::PeerUnreachable(user_err))?;
                Ok::<TcpStream, ClientError>(stream)
            },
        )
        .await;
        inner
            .expected_transfers
            .lock()
            .unwrap()
            .remove(&(username.clone(), their_token));
        inner.uploads_active.fetch_sub(1, Ordering::Relaxed);
        pump_queue(inner.clone());
    });
}

// -------- proactive flow (we open the F connection) --------

fn spawn_proactive(
    inner: Arc<Inner>,
    username: String,
    remote_path: String,
    local: PathBuf,
    size: u64,
) {
    inner.uploads_active.fetch_add(1, Ordering::Relaxed);
    let id = inner.upload_counter.fetch_add(1, Ordering::Relaxed);

    tokio::spawn(async move {
        let inner2 = inner.clone();
        let user2 = username.clone();
        let path2 = remote_path.clone();
        run_upload(
            &inner,
            id,
            &username,
            &remote_path,
            &local,
            size,
            async move {
                // Offer the upload and wait for acceptance.
                let handle = conn::get_or_connect_peer(&inner2, &user2).await?;
                let token = inner2.next_token();
                let (ack_tx, ack_rx) = oneshot::channel();
                inner2
                    .transfer_waits
                    .lock()
                    .unwrap()
                    .insert((user2.clone(), token), ack_tx);
                handle
                    .write_tx
                    .send(peer::transfer_request_upload(token, &path2, size))
                    .map_err(|_| ClientError::PeerUnreachable(user2.clone()))?;
                let ack = tokio::time::timeout(ACK_TIMEOUT, ack_rx)
                    .await
                    .map_err(|_| ClientError::Timeout)?
                    .map_err(|_| ClientError::PeerUnreachable(user2.clone()))?;
                if !ack.allowed {
                    return Err(ClientError::TransferRejected(ack.message));
                }
                conn::connect_transfer(&inner2, &user2, token).await
            },
        )
        .await;
        inner.uploads_active.fetch_sub(1, Ordering::Relaxed);
        pump_queue(inner.clone());
    });
}

/// Shared upload driver: obtain the F stream via `get_stream`, read the
/// 8-byte start offset, and stream the file from there.
async fn run_upload<F>(
    inner: &Arc<Inner>,
    id: u64,
    username: &str,
    remote_path: &str,
    local: &PathBuf,
    size: u64,
    get_stream: F,
) where
    F: std::future::Future<Output = Result<TcpStream>>,
{
    let emit = |state: UploadState, sent: u64| {
        let _ = inner.event_tx.send(Event::Upload(UploadUpdate {
            id,
            username: username.to_string(),
            filename: remote_path.to_string(),
            state,
            sent,
            total: size,
        }));
    };
    emit(UploadState::Transferring, 0);

    let result = async {
        let mut stream = get_stream.await?;

        let mut off = [0u8; 8];
        tokio::time::timeout(ACK_TIMEOUT, stream.read_exact(&mut off))
            .await
            .map_err(|_| ClientError::Timeout)??;
        let offset = u64::from_le_bytes(off).min(size);

        let mut file = tokio::fs::File::open(local).await?;
        file.seek(std::io::SeekFrom::Start(offset)).await?;

        let mut sent = offset;
        let mut buf = vec![0u8; 64 * 1024];
        let mut window = Instant::now();
        let mut window_bytes: u64 = 0;
        let mut last_emit = Instant::now();

        while sent < size {
            let n = file.read(&mut buf).await?;
            if n == 0 {
                break;
            }
            tokio::time::timeout(STALL_TIMEOUT, stream.write_all(&buf[..n]))
                .await
                .map_err(|_| ClientError::Timeout)??;
            sent += n as u64;
            window_bytes += n as u64;

            let elapsed = window.elapsed();
            if elapsed >= Duration::from_secs(2) {
                inner.upload_speed.store(
                    (window_bytes as f64 / elapsed.as_secs_f64()) as u32,
                    Ordering::Relaxed,
                );
                window = Instant::now();
                window_bytes = 0;
            }
            if last_emit.elapsed() >= Duration::from_millis(300) {
                last_emit = Instant::now();
                emit(UploadState::Transferring, sent);
            }
        }
        stream.flush().await?;
        Ok::<(), ClientError>(())
    }
    .await;

    match result {
        Ok(()) => emit(UploadState::Complete, size),
        Err(e) => emit(
            UploadState::Failed {
                reason: e.to_string(),
            },
            0,
        ),
    }
}