Force non blocking read with TcpStream - sockets

I've got a thread, that maintains a list of sockets, and I'd like to traverse the list, see if there is anything to read, if so - act upon it, if not - move onto the next. The problem is, as soon as I come across the first node, all execution is halted until something comes through on the read.
I'm using std::io::Read::read(&mut self, buf: &mut [u8]) -> Result<usize>
From the doc
This function does not provide any guarantees about whether it blocks waiting for data, but if an object needs to block for a read but cannot it will typically signal this via an Err return value.
Digging into the source, the TcpStream Read implementation is
impl Read for TcpStream {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> { self.0.read(buf) }
}
Which invokes
pub fn read(&mut self, buf: &mut [u8]) -> IoResult<uint> {
let fd = self.fd();
let dolock = || self.lock_nonblocking();
let doread = |nb| unsafe {
let flags = if nb {c::MSG_DONTWAIT} else {0};
libc::recv(fd,
buf.as_mut_ptr() as *mut libc::c_void,
buf.len() as wrlen,
flags) as libc::c_int
};
read(fd, self.read_deadline, dolock, doread)
}
And finally, calls read<T, L, R>(fd: sock_t, deadline: u64, mut lock: L, mut read: R)
Where I can see loops over non blocking reads until data has been retrieved or an error has occurred.
Is there a way to force a non-blocking read with TcpStream?

Updated Answer
It should be noted, that as of Rust 1.9.0, std::net::TcpStream has added functionality:
fn set_nonblocking(&self, nonblocking: bool) -> Result<()>
Original Answer
Couldn't exactly get it with TcpStream, and didn't want to pull in a separate lib for IO operations, so I decided to set the file descriptor as Non-blocking before using it, and executing a system call to read/write. Definitely not the safest solution, but less work than implementing a new IO lib, even though MIO looks great.
extern "system" {
fn read(fd: c_int, buffer: *mut c_void, count: size_t) -> ssize_t;
}
pub fn new(user: User, stream: TcpStream) -> Socket {
// First we need to setup the socket as Non-blocking on POSIX
let fd = stream.as_raw_fd();
unsafe {
let ret_value = libc::fcntl(fd,
libc::consts::os::posix01::F_SETFL,
libc::consts::os::extra::O_NONBLOCK);
// Ensure we didnt get an error code
if ret_value < 0 {
panic!("Unable to set fd as non-blocking")
}
}
Socket {
user: user,
stream: stream
}
}
pub fn read(&mut self) {
let count = 512 as size_t;
let mut buffer = [0u8; 512];
let fd = self.stream.as_raw_fd();
let mut num_read = 0 as ssize_t;
unsafe {
let buf_ptr = buffer.as_mut_ptr();
let void_buf_ptr: *mut c_void = mem::transmute(buf_ptr);
num_read = read(fd, void_buf_ptr, count);
if num_read > 0 {
println!("Read: {}", num_read);
}
println!("test");
}
}

Related

Actix Rust deadpool_postgres: DB connections not released

I've got an Actix-web server that connects to a Postgres DB.
I've noticed that after a 1000 requests my Postgres DB's RAM usage has spiked.
When I stop actix-web, the RAM held by the db is cleared. This leads me to believe that my code is not releasing the connection.
I cannot find an example of connections actually being released. It looks like it's inferred in everyone else's code.
Here's mine:
async fn hellow_world(a : f32, b : f32, pool: &Pool) -> Result<Value, PoolError> {
let client: Client = pool.get().await?;
let sql = format!("select \"json\" from public.table_a WHERE a={} and b={}", a, b);
let stmt = client.prepare(&sql).await?;
let row = client.query_one(&stmt, &[]).await?;
let result : Value = row.get(0);
Ok(result)
}
#[derive(Deserialize)]
pub struct MyRequest {
a: f32,
b: f32
}
#[get("/hello")]
async fn sv_hellow_world(info: web::Query<MyRequest>, db_pool: web::Data<Pool>) -> Result<HttpResponse, Error> {
let response : Value = hellow_world(info.a, info.b, &db_pool).await?;
Ok(HttpResponse::Ok().json(response))
}
#[actix_rt::main]
async fn main() -> std::io::Result<()> {
dotenv().ok();
let config = Config::from_env().unwrap();
let pool = config.pg.create_pool(tokio_postgres::NoTls).unwrap();
env_logger::from_env(Env::default().default_filter_or("info")).init();
let server = HttpServer::new(move || App::new().wrap(Logger::default()).wrap(Logger::new("%a %{User-Agent}i")).data(pool.clone()).service(sv_hellow_world))
.bind("0.0.0.0:3000")?
.run();
server.await
}
Based on further testing, #Werner determined that the code was piling up server-side prepared statements.
It is not clear whether these statements can be closed using this library.
Either of two approaches can be used to avoid this problem:
Use a single, shared prepared statement
Use the direct query form instead of the prepared statement
I recommend the first approach on principle as it is more efficient and protects against SQL Injection. It should look something like this:
async fn hellow_world(a : f32, b : f32, pool: &Pool) -> Result<Value, PoolError> {
let client: Client = pool.get().await?;
let stmt = client.prepare("select \"json\" from public.table_a WHERE a=$1::numeric and b=$2::numeric").await?;
let row = client.query_one(&stmt, &[&a, &b]).await?;
let result : Value = row.get(0);
Ok(result)
}
Using this code, only one prepared statement should be created on each of the pool's connections.

Rewriting looping blocking code to SwiftNIO style non-blocking code

I'm working on a driver that will read data from the network. It doesn't know how much is in a response, other than that when it tries to read and gets 0 bytes back, it is done. So my blocking Swift code looks naively like this:
func readAllBlocking() -> [Byte] {
var buffer: [Byte] = []
var fullBuffer: [Byte] = []
repeat {
buffer = read() // synchronous, blocking
fullBuffer.append(buffer)
} while buffer.count > 0
return fullBuffer
}
How can I rewrite this as a promise that will keep on running until the entire result is read? After trying to wrap my brain around it, I'm still stuck here:
func readAllNonBlocking() -> EventLoopFuture<[Byte]> {
///...?
}
I should add that I can rewrite read() to instead of returning a [Byte] return an EventLoopFuture<[Byte]>
Generally, loops in synchronous programming are turned into recursion to get the same effect with asynchronous programming that uses futures (and also in functional programming).
So your function could look like this:
func readAllNonBlocking(on eventLoop: EventLoop) -> EventLoopFuture<[Byte]> {
// The accumulated chunks
var accumulatedChunks: [Byte] = []
// The promise that will hold the overall result
let promise = eventLoop.makePromise(of: [Byte].self)
// We turn the loop into recursion:
func loop() {
// First, we call `read` to read in the next chunk and hop
// over to `eventLoop` so we can safely write to `accumulatedChunks`
// without a lock.
read().hop(to: eventLoop).map { nextChunk in
// Next, we just append the chunk to the accumulation
accumulatedChunks.append(contentsOf: nextChunk)
guard nextChunk.count > 0 else {
promise.succeed(accumulatedChunks)
return
}
// and if it wasn't empty, we loop again.
loop()
}.cascadeFailure(to: promise) // if anything goes wrong, we fail the whole thing.
}
loop() // Let's kick everything off.
return promise.futureResult
}
I would like to add two things however:
First, what you're implementing above is to simply read in everything until you see EOF, if that piece of software is exposed to the internet, you should definitely add a limit on how many bytes to hold in memory maximally.
Secondly, SwiftNIO is an event driven system so if you were to read these bytes with SwiftNIO, the program would actually look slightly differently. If you're interested what it looks like to simply accumulate all bytes until EOF in SwiftNIO, it's this:
struct AccumulateUntilEOF: ByteToMessageDecoder {
typealias InboundOut = ByteBuffer
func decode(context: ChannelHandlerContext, buffer: inout ByteBuffer) throws -> DecodingState {
// `decode` will be called if new data is coming in.
// We simply return `.needMoreData` because always need more data because our message end is EOF.
// ByteToMessageHandler will automatically accumulate everything for us because we tell it that we need more
// data to decode a message.
return .needMoreData
}
func decodeLast(context: ChannelHandlerContext, buffer: inout ByteBuffer, seenEOF: Bool) throws -> DecodingState {
// `decodeLast` will be called if NIO knows that this is the _last_ time a decode function is called. Usually,
// this is because of EOF or an error.
if seenEOF {
// This is what we've been waiting for, `buffer` should contain all bytes, let's fire them through
// the pipeline.
context.fireChannelRead(self.wrapInboundOut(buffer))
} else {
// Odd, something else happened, probably an error or we were just removed from the pipeline. `buffer`
// will now contain what we received so far but maybe we should just drop it on the floor.
}
buffer.clear()
return .needMoreData
}
}
If you wanted to make a whole program out of this with SwiftNIO, here's an example that is a server which accepts all data until it sees EOF and then literally just writes back the number of received bytes :). Of course, in the real world you would never hold on to all the received bytes to count them (you could just add each individual piece) but I guess it serves as an example.
import NIO
let group = MultiThreadedEventLoopGroup(numberOfThreads: 1)
defer {
try! group.syncShutdownGracefully()
}
struct AccumulateUntilEOF: ByteToMessageDecoder {
typealias InboundOut = ByteBuffer
func decode(context: ChannelHandlerContext, buffer: inout ByteBuffer) throws -> DecodingState {
// `decode` will be called if new data is coming in.
// We simply return `.needMoreData` because always need more data because our message end is EOF.
// ByteToMessageHandler will automatically accumulate everything for us because we tell it that we need more
// data to decode a message.
return .needMoreData
}
func decodeLast(context: ChannelHandlerContext, buffer: inout ByteBuffer, seenEOF: Bool) throws -> DecodingState {
// `decodeLast` will be called if NIO knows that this is the _last_ time a decode function is called. Usually,
// this is because of EOF or an error.
if seenEOF {
// This is what we've been waiting for, `buffer` should contain all bytes, let's fire them through
// the pipeline.
context.fireChannelRead(self.wrapInboundOut(buffer))
} else {
// Odd, something else happened, probably an error or we were just removed from the pipeline. `buffer`
// will now contain what we received so far but maybe we should just drop it on the floor.
}
buffer.clear()
return .needMoreData
}
}
// Just an example "business logic" handler. It will wait for one message
// and just write back the length.
final class SendBackLengthOfFirstInput: ChannelInboundHandler {
typealias InboundIn = ByteBuffer
typealias OutboundOut = ByteBuffer
func channelRead(context: ChannelHandlerContext, data: NIOAny) {
// Once we receive the message, we allocate a response buffer and just write the length of the received
// message in there. We then also close the channel.
let allData = self.unwrapInboundIn(data)
var response = context.channel.allocator.buffer(capacity: 10)
response.writeString("\(allData.readableBytes)\n")
context.writeAndFlush(self.wrapOutboundOut(response)).flatMap {
context.close(mode: .output)
}.whenSuccess {
context.close(promise: nil)
}
}
func errorCaught(context: ChannelHandlerContext, error: Error) {
print("ERROR: \(error)")
context.channel.close(promise: nil)
}
}
let server = try ServerBootstrap(group: group)
// Allow us to reuse the port after the process quits.
.serverChannelOption(ChannelOptions.socket(.init(SOL_SOCKET), .init(SO_REUSEADDR)), value: 1)
// We should allow half-closure because we want to write back after having received an EOF on the input
.childChannelOption(ChannelOptions.allowRemoteHalfClosure, value: true)
// Our program consists of two parts:
.childChannelInitializer { channel in
channel.pipeline.addHandlers([
// 1: The accumulate everything until EOF handler
ByteToMessageHandler(AccumulateUntilEOF(),
// We want 1 MB of buffering max. If you remove this parameter, it'll also
// buffer indefinitely.
maximumBufferSize: 1024 * 1024),
// 2: Our "business logic"
SendBackLengthOfFirstInput()
])
}
// Let's bind port 9999
.bind(to: SocketAddress(ipAddress: "127.0.0.1", port: 9999))
.wait()
// This will never return.
try server.closeFuture.wait()
Demo:
$ echo -n "hello world" | nc localhost 9999
11

Using a callback when handling TCP connections with Tokio

I am trying to have a struct that starts an event loop, listens for TCP connections and calls a callback for each connection.
(The callback will be handed some prepossessed data from the socket. In my example below I just hand it the IP address of the connection but in my real code I will parse the contents that I receive with serde into a struct and pass that into the callback. I hope that doesn't invalidate the following "not working example").
My Cargo.toml:
[package]
name = "lifetime-problem"
version = "0.1.0"
edition = "2018"
[dependencies]
tokio-tcp = "0.1.3"
tokio = "0.1.14"
[[bin]]
name = "lifetime-problem"
path = "main.rs"
and main.rs:
use tokio::prelude::*;
struct Test {
printer: Option<Box<Fn(std::net::SocketAddr) + Sync>>,
}
impl Test {
pub fn start(&mut self) -> Result<(), Box<std::error::Error>> {
let addr = "127.0.0.1:4242".parse::<std::net::SocketAddr>()?;
let listener = tokio::net::TcpListener::bind(&addr)?;
let server = listener
.incoming()
.map_err(|e| eprintln!("failed to accept socket; error = {:?}", e))
.for_each(move |socket: tokio::net::TcpStream| {
let address = socket.peer_addr().expect("");
match self.printer {
Some(callback) => { callback(address); }
None => { println!("{}", address); }
}
Ok(())
});
tokio::run(server);
Ok(())
}
}
fn main() {
let mut x = Test{ printer: None };
x.start();
}
I have tried several things starting from this code (which is adopted directly from the Tokio example).
If I use the code like posted above I get:
error[E0277]: (dyn std::ops::Fn(std::net::SocketAddr) + std::marker::Sync + 'static) cannot be sent between threads safely
for the line 24 (tokio::run(server)).
If I add the Send trait on the Fn in the printer field XOR if I remove the move in the closure in the for_each call I get another error instead:
error[E0495]: cannot infer an appropriate lifetime due to conflicting requirements
which points me to the closure that apparently cannot outlive the start method where it is defined but tokio::run seems to have conflicting requirements for it.
Do you know if I am addressing the callback pattern in totally the wrong way or if there is just some minor error in my code?
First things first:
Compiler will translate Box<Fn(std::net::SocketAddr) + Sync> to Box<Fn(std::net::SocketAddr) + Sync + 'static> unless the lifetime is explicitly specified.
Let's have a look at the errors:
error[E0277]: (dyn std::ops::Fn(std::net::SocketAddr) + std::marker::Sync + 'static) cannot be sent between threads safely
This is self-explanatory. You are trying to move &mut T to another thread, but cannot, because T here is not Send. To send &mut T to another thread T too needs to be of type Send.
Here is the minimal code that will give the same error:
use std::fmt::Debug;
fn func<T> (i:&'static mut T) where T: Debug {
std::thread::spawn(move || {
println!("{:?}", i);
});
}
If I make T above to also be of type Send, the error goes away.
But in your case when you add the Send trait, it gives lifetime error. Why?
&mut self has some lifetime greater than the function start() set by the caller, but there's no guarantee that its 'static. You move this reference into the closure which is passed to the thread and can potentially outlive the scope it is closing over, leading to a dangling reference.
Here's a minimal version, that would give the same error.
use std::fmt::Debug;
fn func<'a, T:'a> (i:&'a mut T) where T: Debug + Sync + Send {
std::thread::spawn(move || {
println!("{:?}", i);
});
}
Sync is not really required here as it is &mut T. Changing &mut T to &T (retaining Sync), will also result into the same error. The onus here is on references and not mutability. So you see, there is some lifetime 'a and it is moved into a closure (given to a thread), which means the closure now contains a reference (disjoint from the main context). So now, what is 'a and how long will it live from the closure's perspective that is invoked from another thread? Not inferable! As a result, the compiler complains saying cannot infer an appropriate lifetime due to conflicting requirements.
If we tweak the code a bit to;
impl Test {
pub fn start(&'static mut self) -> Result<(), Box<std::error::Error>> {
let addr = "127.0.0.1:4242".parse::<std::net::SocketAddr>()?;
let listener = tokio::net::TcpListener::bind(&addr)?;
let server = listener
.incoming()
.map_err(|e| eprintln!("failed to accept socket; error = {:?}", e))
.for_each(move |socket: tokio::net::TcpStream| {
let address = socket.peer_addr().expect("");
match &self.printer {
Some(callback) => { callback(address); }
None => { println!("{}", address); }
}
Ok(())
});
tokio::run(server);
Ok(())
}
}
it will compile fine. There's a guarantee there that self has a 'static lifetime. Please note that in the match statement we need to pass &self.printer, as you cannot move out of a borrowed context.
However, this expects Test to be declared static and that too a mutable one, which is generally not the best way, if you have other options.
Another way is; if it's ok for you to pass Test by value to start() and then further move it into for_each(), the code would look like this:
use tokio::prelude::*;
struct Test {
printer: Option<Box<Fn(std::net::SocketAddr) + Send>>,
}
impl Test {
pub fn start(mut self) -> Result<(), Box<std::error::Error>> {
let addr = "127.0.0.1:4242".parse::<std::net::SocketAddr>()?;
let listener = tokio::net::TcpListener::bind(&addr)?;
let server = listener
.incoming()
.map_err(|e| eprintln!("failed to accept socket; error = {:?}", e))
.for_each(move |socket: tokio::net::TcpStream| {
let address = socket.peer_addr().expect("");
match &self.printer {
Some(callback) => {
callback(address);
}
None => {
println!("{}", address);
}
}
Ok(())
});
tokio::run(server);
Ok(())
}
}
fn main() {
let mut x = Test { printer: None };
x.start();
}

Reading Data from Socket Golang

I'm trying to read data from a telnet session in golang. I wrote the following functions in an attempt to accomplish this.
Initially I was having an issue where I was reading from a socket with no data so it would lock and never return. BufferSocketData is my attempt to work around this issue as I can't know if there is data to read. The idea is it will wait 1 second before determining there is not data in the socket and return an empty string.
GetData seems to work the first time there is new data in the buffer, but beyond that it gets no new data. I'm sure this has something to do with my use of goroutines and channels, I'm new to go and I'm sure I'm not using them correctly.
Any ideas as to why my subsequent reads return no data?
/*
ReadDataFromSocket - Attempts to read any data in the socket.
*/
func ReadDataFromSocket(sock io.Reader, c chan string) {
var recvData = make([]byte, 1024)
var numBytes, _ = sock.Read(recvData)
c <- string(recvData[:numBytes])
}
/*
BufferSocketData - Read information from the socket and store it in the buffer.
*/
func (tn *TelnetLib) BufferSocketData(inp chan string, out chan string) {
var data string
var timeout int64 = 1000 // 1 second timeout.
var start = utils.GetTimestamp()
for utils.GetTimestamp()-start < timeout {
select {
case data = <-inp:
default:
}
if data != "" {
break
}
}
out <- data
}
/*
GetData - Start goroutines to get and buffer data.
*/
func (tn *TelnetLib) GetData() {
var sockCh = make(chan string)
var buffCh = make(chan string)
go ReadDataFromSocket(tn.Conn, sockCh)
go tn.BufferSocketData(sockCh, buffCh)
var data = <-buffCh
if data != "" {
tn.Buffer += data
}
}
Please let me know if you need any additional information.
Use SetReadDeadline to read data with a time limit:
func (tn *TelnetLib) GetData() {
tn.Conn.SetReadDeadline(time.Second)
recvData := make([]byte, 1024)
n, err := tn.Conn.Read(recvData)
if n > 0 {
// do something with recvData[:n]
}
if e, ok := err.(interface{ Timeout() bool }); ok && e.Timeout() {
// handle timeout
} else if err != nil {
// handle error
}
}
Note that a single call Read may not read all data sent by the peer. You may want to accumulate data by calling Read in a loop or call io.ReadFull.

How do I close a Unix socket in Rust?

I have a test that opens and listens to a Unix Domain Socket. The socket is opened and reads data without issues, but it doesn't shutdown gracefully.
This is the error I get when I try to run the test a second time:
thread 'test_1' panicked at 'called Result::unwrap() on an Err
value: Error { repr: Os { code: 48, message: "Address already in use"
} }', ../src/libcore/result.rs:799 note: Run with RUST_BACKTRACE=1
for a backtrace.
The code is available at the Rust playground and there's a Github Gist for it.
use std::io::prelude::*;
use std::thread;
use std::net::Shutdown;
use std::os::unix::net::{UnixStream, UnixListener};
Test Case:
#[test]
fn test_1() {
driver();
assert_eq!("1", "2");
}
Main entry point function
fn driver() {
let listener = UnixListener::bind("/tmp/my_socket.sock").unwrap();
thread::spawn(|| socket_server(listener));
// send a message
busy_work(3);
// try to disconnect the socket
let drop_stream = UnixStream::connect("/tmp/my_socket.sock").unwrap();
let _ = drop_stream.shutdown(Shutdown::Both);
}
Function to send data in intervals
#[allow(unused_variables)]
fn busy_work(threads: i32) {
// Make a vector to hold the children which are spawned.
let mut children = vec![];
for i in 0..threads {
// Spin up another thread
children.push(thread::spawn(|| socket_client()));
}
for child in children {
// Wait for the thread to finish. Returns a result.
let _ = child.join();
}
}
fn socket_client() {
let mut stream = UnixStream::connect("/tmp/my_socket.sock").unwrap();
stream.write_all(b"hello world").unwrap();
}
Function to handle data
fn handle_client(mut stream: UnixStream) {
let mut response = String::new();
stream.read_to_string(&mut response).unwrap();
println!("got response: {:?}", response);
}
Server socket that listens to incoming messages
#[allow(unused_variables)]
fn socket_server(listener: UnixListener) {
// accept connections and process them, spawning a new thread for each one
for stream in listener.incoming() {
match stream {
Ok(mut stream) => {
/* connection succeeded */
let mut response = String::new();
stream.read_to_string(&mut response).unwrap();
if response.is_empty() {
break;
} else {
thread::spawn(|| handle_client(stream));
}
}
Err(err) => {
/* connection failed */
break;
}
}
}
println!("Breaking out of socket_server()");
drop(listener);
}
Please learn to create a minimal reproducible example and then take the time to do so. In this case, there's no need for threads or functions or testing frameworks; running this entire program twice reproduces the error:
use std::os::unix::net::UnixListener;
fn main() {
UnixListener::bind("/tmp/my_socket.sock").unwrap();
}
If you look at the filesystem before and after the test, you will see that the file /tmp/my_socket.sock is not present before the first run and it is present before the second run. Deleting the file allows the program to run to completion again (at which point it recreates the file).
This issue is not unique to Rust:
Note that, once created, this socket file will continue to exist, even after the server exits. If the server subsequently restarts, the file prevents re-binding:
[...]
So, servers should unlink the socket pathname prior to binding it.
You could choose to add some wrapper around the socket that would automatically delete it when it is dropped or create a temporary directory that is cleaned when it is dropped, but I'm not sure how well that would work. You could also create a wrapper function that deletes the file before it opens the socket.
Unlinking the socket when it's dropped
use std::path::{Path, PathBuf};
struct DeleteOnDrop {
path: PathBuf,
listener: UnixListener,
}
impl DeleteOnDrop {
fn bind(path: impl AsRef<Path>) -> std::io::Result<Self> {
let path = path.as_ref().to_owned();
UnixListener::bind(&path).map(|listener| DeleteOnDrop { path, listener })
}
}
impl Drop for DeleteOnDrop {
fn drop(&mut self) {
// There's no way to return a useful error here
let _ = std::fs::remove_file(&self.path).unwrap();
}
}
You may also want to consider implementing Deref / DerefMut to make this into a smart pointer for sockets:
impl std::ops::Deref for DeleteOnDrop {
type Target = UnixListener;
fn deref(&self) -> &Self::Target {
&self.listener
}
}
impl std::ops::DerefMut for DeleteOnDrop {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.listener
}
}
Unlinking the socket before it's opened
This is much simpler:
use std::path::Path;
fn bind(path: impl AsRef<Path>) -> std::io::Result<UnixListener> {
let path = path.as_ref();
std::fs::remove_file(path)?;
UnixListener::bind(path)
}
Note that you can combine the two solutions, such that the socket is deleted before creation and when it's dropped.
I think that deleting during creation is a less-optimal solution: if you ever start a second server, you'll prevent the first server from receiving any more connections. It's probably better to error and tell the user instead.