Function proxy_main

gateway/src/proxy.rs:171–234

(rt: &Runtime, config: &ProxyConfig, proxy: Proxy)

#[inline(never)]
pub async fn proxy_main(rt: &Runtime, config: &ProxyConfig, proxy: Proxy) -> Result<()> {
    let mut tcp_listeners = Vec::new();
    for &port in &config.listen_port {
        let listener = TcpListener::bind((config.listen_addr, port))
            .await
            .with_context(|| format!("failed to bind {}:{}", config.listen_addr, port))?;
        info!("tcp bridge listening on {}:{}", config.listen_addr, port);
        tcp_listeners.push(listener);
    }
    if tcp_listeners.is_empty() {
        bail!("no tcp listen ports configured");
    }

    let poll_counter = AtomicUsize::new(0);
    loop {
        // Accept from any TCP listener via round-robin poll.
        let poll_start = poll_counter.fetch_add(1, Ordering::Relaxed);
        let n = tcp_listeners.len();
        let accepted: std::io::Result<(TcpStream, std::net::SocketAddr)> =
            std::future::poll_fn(|cx| {
                for j in 0..n {
                    let i = (poll_start + j) % n;
                    if let Poll::Ready(result) = tcp_listeners[i].poll_accept(cx) {
                        return Poll::Ready(result);
                    }
                }
                Poll::Pending
            })
            .await;
        match accepted {
            Ok((inbound, from)) => {
                let span = info_span!("conn", id = next_connection_id());
                let _enter = span.enter();
                let conn_entered = EnteredCounter::new(&NUM_CONNECTIONS);

                info!(%from, "new connection");
                let proxy = proxy.clone();
                rt.spawn(
                    async move {
                        let _conn_entered = conn_entered;
                        let timeouts = &proxy.config.proxy.timeouts;
                        let result =
                            timeout(timeouts.total, handle_connection(inbound, proxy)).await;
                        match result {
                            Ok(Ok(_)) => {
                                info!("connection closed");
                            }
                            Ok(Err(e)) => {
                                error!("connection error: {e:#}");
                            }
                            Err(_) => {
                                error!("connection kept too long, force closing");
                            }
                        }
                    }
                    .in_current_span(),
                );
            }