21. 從零用Rust編寫正反向代理，tokio竟然這樣對待socket！

問蒙服務框架發表於2023-11-01
Rust

wmproxy

wmproxy已用Rust實現http/https代理, socks5代理, 反向代理, 靜態檔案伺服器，四層TCP/UDP轉發，內網穿透，後續將實現websocket代理等，會將實現過程分享出來，感興趣的可以一起造個輪子

專案地址

國內: https://gitee.com/tickbh/wmproxy

github: https://github.com/tickbh/wmproxy

有請主角上場

Socket是集萬千寵愛為一身的王子，在作業系統的王國裡，他負責對外的所有通訊，所以要想溝通鄰國的公主必須經過他，所以大家對他都是萬般友好。

這天一個Rust城市裡的大臣tokio對他發起了邀請，邀請他來參觀嚴謹的邏輯莊園。

tokio莊園

莊園中的各成員對即將到來的王子議論紛紛。

  大管家mio說：“大家都想想等下怎麼向socket王子介紹自己，好讓他配合大家的工作。”

  大管家mio是tokio的基石，一切和王國打交道的都交由他去打理，他是保證莊園高效運轉的關鍵，此刻他準備好了歡迎宴會。

  在宴會上，socket聽說tokio莊園是這座城市非同步執行的重要基石，就很好奇的讓大夥介紹介紹下怎麼工作的。

  莊園主tokio就說：“我是依靠著大管家mio幫我負責處理底層的事，Waker來提醒我有新的事情，PollEvented來幫我管理事件的。下面先讓mio來介紹下。”

  管家mio說：“我負責收集莊園中的所有資訊，他們告訴我他們要關心的什麼比較，比如您的到來（可讀），或者您有什麼話想說（可寫），我會負責和王國的底層進行溝通，我在這個國家用的是epoll，據說在遙遠的另一個國家用的是iocp，如果有相應的需求，我將會通知Waker，由他去提醒莊主來及時的處理，這場宴會也是提前得到通知而進行準備的。”

  通知Waker說：“我所做的事情就是微不足道，我的對接物件是PollEvented，當他關心讀事件，我會向mio去發起poll_read請求，如果此時mio那邊已經知道有新的訊息了，那我就直接把他們讀出來交給民眾Poll::Ready，如果此時還沒有新訊息，那我會告訴管家，有新訊息的時候通知我Poll::Pending，此時我就在這裡等待，直到有新的訊息到達我就通知給民眾。當他關心寫事件，我會向mio請求poll_write請求，後續的和收訊息的一致。現在給你們展示下包裝了一層我的Context和我能換醒的虛表。”

/// 這個在程式碼裡就是經常看到，它就是我的一層淺封裝啦。
pub struct Context<'a> {
    waker: &'a Waker,
    _marker: PhantomData<fn(&'a ()) -> &'a ()>,
}
/// 我透過他來控制回撥，保證喚醒的時候能正確的通知
pub struct RawWakerVTable {
    clone: unsafe fn(*const ()) -> RawWaker,
    wake: unsafe fn(*const ()),
    wake_by_ref: unsafe fn(*const ()),
    drop: unsafe fn(*const ()),
}

  事件PollEvented說：“莊主要處理的事情太多了，而有些事情又需要等待一層層的反饋，他沒法把精力放在一件事情上一直等待，所以就有了我出馬，莊主告訴我他關心什麼事，我就把它記下來，這樣子莊主就可以去處理其它的事，等事件到來的時候我就告訴莊主，這樣子莊主就可以高效的處理所有的事件。”

王子覺得他們說了一堆有點囉嗦

  “帶我看看你們實際的工坊，我要實地考查下。”王子說。

莊主就帶著王子來到了，受理工坊，受理工坊正在處理建立受理點:

TcpListener::bind(addr).await

受理點的內容就是PollEvent：

pub struct TcpListener {
    io: PollEvented<mio::net::TcpListener>,
}

當他接受新的受理者的時候：

pub async fn accept(&self) -> io::Result<(TcpStream, SocketAddr)> {
    let (mio, addr) = self
        .io
        .registration()
        .async_io(Interest::READABLE, || self.io.accept())
        .await?;

    let stream = TcpStream::new(mio)?;
    Ok((stream, addr))
}

他向PollEvent註冊了可讀事情有的時候通知他，此時PollEvent就建立了一個Waker物件，當有符合條件的時候就來告訴他：

/// 建立一個可讀的Future物件
fn readiness_fut(&self, interest: Interest) -> Readiness<'_> {
    Readiness {
        scheduled_io: self,
        state: State::Init,
        waiter: UnsafeCell::new(Waiter {
            pointers: linked_list::Pointers::new(),
            waker: None,
            is_ready: false,
            interest,
            _p: PhantomPinned,
        }),
    }
}

底層有mio和系統互動，一旦有資料就通知Waker，他建在runtime/io/driver.rs上面

fn turn(&mut self, handle: &Handle, max_wait: Option<Duration>) {
    let events = &mut self.events;
    // 高效的監聽埠
    match self.poll.poll(events, max_wait) {
        Ok(_) => {}
    }

    for event in events.iter() {
        let token = event.token();

        if token == TOKEN_WAKEUP {
        } else if token == TOKEN_SIGNAL {
        } else {
            let ready = Ready::from_mio(event);
            let ptr: *const ScheduledIo = token.0 as *const _;
            
            let io: &ScheduledIo = unsafe { &*ptr };
            io.set_readiness(Tick::Set(self.tick), |curr| curr | ready);
            // 有相應的事件，進行喚醒然後通知上層處理相應的事件
            io.wake(ready);

        }
    }

}

然後看到Waker工坊上處理：

pub(crate) fn wake_all(&mut self) {
    assert!(self.curr <= NUM_WAKERS);
    while self.curr > 0 {
        self.curr -= 1;
        let waker = unsafe { ptr::read(self.inner[self.curr].as_mut_ptr()) };
        waker.wake();
    }
}
pub fn wake(self) {
    // 存在的回撥函式及對應的資料，好進行呼叫
    let wake = self.waker.vtable.wake;
    let data = self.waker.data;

    crate::mem::forget(self);

    // 用回撥函式的方式處理剛等待執行的執行緒
    unsafe { (wake)(data) };
}

最後又回到了受理工坊，我們知道了一個新的來源TcpStream的到來，期間在等待的時候，我們可以去處理其它的事情，不至於空有許多人力物力，卻在等我的寶的事情到來沒法快速處理事情。

  王子說：“在只有一個受理的時候你們這麼高效，如果有同時多個受理，又需要在處理完訪問相同的資料，你們又能處理嗎？”

  莊主說：“那麼接下來就讓我帶你參觀下Poll工坊，他用同步的方式可以同時處理多個連結。”

參觀完非同步工坊，莊主又帶著王子來到了

只見Poll工坊大螢幕上就出現了一個例子：

#[tokio::main]
async fn main() -> std::io::Result<()> {
    use std::{future::poll_fn, task::Poll, pin::Pin};
    use tokio::net::TcpListener;
    let mut listeners = vec![];
    // 同時監聽若干個埠
    for i in 1024..1099 {
        listeners.push(TcpListener::bind(format!("127.0.0.1:{}", i)).await?);
    }
    loop {
        let mut pin_listener = Pin::new(&mut listeners);
        // 同時監聽若干個埠，只要有任一個返回則返回資料
        let fun = poll_fn(|cx| {
            for l in &*pin_listener.as_mut() {
                match l.poll_accept(cx) {
                    v @ Poll::Ready(_) => return v,
                    Poll::Pending => {},
                }
            }
            Poll::Pending
        });
        let (conn, addr) = fun.await?;
        println!("receiver conn:{:?} addr:{:?}", conn, addr);
    }
}

他可快速的在函式中同時等待多個埠資料，這種同步的邏輯可以在複雜結構時很方便的書寫程式碼的邏輯。

王子看完後：“現在你的處理能力已經高效又靈活，真正的可甜可鹹了，把我的能力發揮完全又簡化了操作。”

點選 [關注]，[在看]，[點贊] 是對作者最大的支援

相關文章