sark 0.9.0

#![cfg(target_os = "linux")]
#![allow(clippy::too_many_arguments)]

mod support;

use std::collections::VecDeque;
use std::io::{Read, Write};
use std::net::TcpStream;
use std::pin::Pin;
use std::task::Poll;
use std::time::Duration;

use dope_extra::harness::Harness;
use dope_fiber::{Context, Fiber};
use o3::buffer::Shared;
use sark::{Executor, Throughput, driver};
use sark_core::http::Stream;

const CHUNK_TERMINATOR_BYTES: &[u8] = b"0\r\n\r\n";

struct DelayedChunks {
    chunks: VecDeque<Shared>,
    pend_next: bool,
    yielded: u32,
}

impl DelayedChunks {
    fn new<I: IntoIterator<Item = &'static [u8]>>(parts: I) -> Self {
        Self {
            chunks: parts.into_iter().map(Shared::copy_from_slice).collect(),
            pend_next: true,
            yielded: 0,
        }
    }

    fn arm_pend(&mut self, cx: Pin<&mut Context<'_, '_>>) -> Poll<()> {
        if self.pend_next {
            self.pend_next = false;
            cx.wake();
            return Poll::Pending;
        }
        self.pend_next = true;
        Poll::Ready(())
    }
}

impl<'d> Fiber<'d> for DelayedChunks {
    type Output = Option<Shared>;
    fn poll(self: Pin<&mut Self>, cx: Pin<&mut Context<'_, 'd>>) -> Poll<Option<Shared>> {
        let this = self.get_mut();
        if this.chunks.is_empty() {
            return Poll::Ready(None);
        }
        if this.arm_pend(cx).is_pending() {
            return Poll::Pending;
        }
        this.yielded += 1;
        Poll::Ready(this.chunks.pop_front())
    }
}

#[sark_gen::request]
struct EmptyReq {}

#[sark_gen::handler]
fn stream_handler(_req: EmptyReq, _state: &()) -> Stream<DelayedChunks> {
    Stream::new(DelayedChunks::new([
        b"alpha".as_slice(),
        b"-beta".as_slice(),
        b"-gamma".as_slice(),
    ]))
    .header(b"content-type", b"text/plain")
}

sark_gen::define_route! {
    StreamPendingDispatch: () => {
        GET "/stream" => stream(capacity = 32) stream_handler,
    }
}

#[test]
fn pending_chunk_producer_completes_the_response() {
    let bind: std::net::SocketAddr = "127.0.0.1:18921".parse().unwrap();
    let server = support::http_server(bind, Duration::from_secs(10));

    Harness::new(bind)
        .run_with_trigger(
            |_ctx, trigger| {
                let driver_config =
                    driver::Config::for_tcp_profile::<Throughput>(support::MAX_CONNECTIONS);
                let executor = Executor::new(driver_config)?;
                executor.enter(|mut session| {
                    server.clone().serve(
                        &mut session,
                        StreamPendingDispatch::new(
                            (),
                            sark::app::Config {
                                timer_capacity: support::MAX_CONNECTIONS.saturating_mul(2),
                                task_capacity: support::MAX_CONNECTIONS,
                            },
                        ),
                        Some(trigger),
                    )
                })
            },
            |bind| {
                let mut sock = TcpStream::connect(bind).expect("connect");
                sock.set_read_timeout(Some(Duration::from_secs(3))).unwrap();

                sock.write_all(b"GET /stream HTTP/1.1\r\nHost: x\r\nConnection: close\r\n\r\n")
                    .unwrap();
                let mut resp = Vec::new();
                let mut chunk = [0u8; 4096];
                loop {
                    match sock.read(&mut chunk) {
                        Ok(0) => break,
                        Ok(n) => {
                            resp.extend_from_slice(&chunk[..n]);
                            if resp.ends_with(CHUNK_TERMINATOR_BYTES) {
                                break;
                            }
                        }
                        Err(e)
                            if e.kind() == std::io::ErrorKind::WouldBlock
                                || e.kind() == std::io::ErrorKind::TimedOut =>
                        {
                            break;
                        }
                        Err(e) => panic!("read: {e}"),
                    }
                }
                let resp_str = String::from_utf8_lossy(&resp);

                assert!(resp_str.contains("200 OK"), "resp: {resp_str}");
                assert!(
                    resp_str
                        .to_lowercase()
                        .contains("transfer-encoding: chunked"),
                    "resp: {resp_str}"
                );
                assert!(
                    resp.ends_with(CHUNK_TERMINATOR_BYTES),
                    "stream never completed (wedged?): {resp_str}"
                );

                let body_start = resp
                    .windows(4)
                    .position(|w| w == b"\r\n\r\n")
                    .map(|i| i + 4)
                    .expect("header terminator");
                let body = &resp[body_start..];

                let mut cursor = 0usize;
                let mut reassembled = Vec::new();
                loop {
                    let line_end = body[cursor..]
                        .windows(2)
                        .position(|w| w == b"\r\n")
                        .expect("chunk size CRLF")
                        + cursor;
                    let size_str = std::str::from_utf8(&body[cursor..line_end]).expect("hex utf8");
                    let size = usize::from_str_radix(size_str.trim(), 16).expect("hex parse");
                    cursor = line_end + 2;
                    if size == 0 {
                        break;
                    }
                    reassembled.extend_from_slice(&body[cursor..cursor + size]);
                    cursor += size;
                    assert_eq!(&body[cursor..cursor + 2], b"\r\n", "chunk CRLF");
                    cursor += 2;
                }
                assert_eq!(
                    reassembled, b"alpha-beta-gamma",
                    "reassembled body: {resp_str}"
                );
            },
        )
        .expect("harness");
}