优雅退出 TcpListener.incoming()

Question

来自 rust 标准网络库：

let listener = TcpListener::bind(("127.0.0.1", port)).unwrap();

info!("Opened socket on localhost port {}", port);

// accept connections and process them serially
for stream in listener.incoming() {
    break;
}

info!("closed socket");

如何让听众停止倾听？它在 API 中说，当侦听器被删除时，它会停止。但是如果incoming() 是一个阻塞调用，我们如何丢弃它呢？最好不要像 tokio/mio 这样的外部 crate。

Answer 1

您需要使用 set_nonblocking() 方法将 TcpListener 置于非阻塞模式，如下所示：

use std::io;
use std::net::TcpListener;

let listener = TcpListener::bind("127.0.0.1:7878").unwrap();
listener.set_nonblocking(true).expect("Cannot set non-blocking");

for stream in listener.incoming() {
    match stream {
        Ok(s) => {
            // do something with the TcpStream
            handle_connection(s);
        }
        Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
            // Decide if we should exit
            break;
            // Decide if we should try to accept a connection again
            continue;
        }
        Err(e) => panic!("encountered IO error: {}", e),
    }
}

incoming() 调用将立即返回一个 Result 类型，而不是等待连接。如果 Result 为 Ok()，则表示已建立连接，您可以对其进行处理。如果 Result 是 Err(WouldBlock)，这实际上不是错误，只是在 incoming() 检查套接字的那一刻没有挂起的连接。

请注意，在 WillBlock 的情况下，您可能需要在继续之前放置一个 sleep() 或其他东西，否则您的程序将快速轮询incoming() 函数检查连接，从而导致 CPU 使用率很高。

代码示例改编自here

Answer 2

标准库没有为此提供 API，但您可以使用特定于平台的 API 来关闭套接字上的读取，这将导致 incoming 迭代器返回错误。然后，您可以在收到错误时中断处理连接。例如，在 Unix 系统上：

use std::net::TcpListener;
use std::os::unix::io::AsRawFd;
use std::thread;

let listener = TcpListener::bind("localhost:0")?;

let fd = listener.as_raw_fd();

let handle = thread::spawn(move || {
  for connection in listener.incoming() {
    match connection {
      Ok(connection) => /* handle connection */
      Err(_) => break,
  }
});

libc::shutdown(fd, libc::SHUT_RD);

handle.join();

Answer 3

您可以使用eventfd poll 您的套接字，它用于发送信号。 我为此写了一个助手。

let shutdown = EventFd::new();
let listener = TcpListener::bind("0.0.0.0:12345")?;
let incoming = CancellableIncoming::new(&listener, &shutdown);

for stream in incoming {
    // Your logic
}

// While in other thread
shutdown.add(1);  // Light the shutdown signal, now your incoming loop exits gracefully.

use nix;
use nix::poll::{poll, PollFd, PollFlags};
use nix::sys::eventfd::{eventfd, EfdFlags};
use nix::unistd::{close, write};
use std;
use std::net::{TcpListener, TcpStream};
use std::os::unix::io::{AsRawFd, RawFd};

pub struct EventFd {
    fd: RawFd,
}

impl EventFd {
    pub fn new() -> Self {
        EventFd {
            fd: eventfd(0, EfdFlags::empty()).unwrap(),
        }
    }

    pub fn add(&self, v: i64) -> nix::Result<usize> {
        let b = v.to_le_bytes();
        write(self.fd, &b)
    }
}

impl AsRawFd for EventFd {
    fn as_raw_fd(&self) -> RawFd {
        self.fd
    }
}

impl Drop for EventFd {
    fn drop(&mut self) {
        let _ = close(self.fd);
    }
}

// -----
//
pub struct CancellableIncoming<'a> {
    listener: &'a TcpListener,
    eventfd: &'a EventFd,
}

impl<'a> CancellableIncoming<'a> {
    pub fn new(listener: &'a TcpListener, eventfd: &'a EventFd) -> Self {
        Self { listener, eventfd }
    }
}

impl<'a> Iterator for CancellableIncoming<'a> {
    type Item = std::io::Result<TcpStream>;
    fn next(&mut self) -> Option<std::io::Result<TcpStream>> {
        use nix::errno::Errno;

        let fd = self.listener.as_raw_fd();
        let evfd = self.eventfd.as_raw_fd();
        let mut poll_fds = vec![
            PollFd::new(fd, PollFlags::POLLIN),
            PollFd::new(evfd, PollFlags::POLLIN),
        ];

        loop {
            match poll(&mut poll_fds, -1) {
                Ok(_) => break,
                Err(nix::Error::Sys(Errno::EINTR)) => continue,
                _ => panic!("Error polling"),
            }
        }

        if poll_fds[0].revents().unwrap() == PollFlags::POLLIN {
            Some(self.listener.accept().map(|p| p.0))
        } else if poll_fds[1].revents().unwrap() == PollFlags::POLLIN {
            None
        } else {
            panic!("Can't be!");
        }
    }
}

Answer 4

注意设置

listener.set_nonblocking(true).expect("Cannot set non-blocking");

按照@effect's answer 中的建议，自动将所有流上的非阻塞设置为 true。也许你想要这个，也许你不想要；我不得不用

将它们设置回false

stream.set_nonblocking(false);