5
\$\begingroup\$

Background

I've been working on an algorithmic trading platform in Rust and utilizing the newly released futures-rs library to keep everything asynchronous and non-blocking.

One issue I've come across with this is that the majority of database interface libraries currently available for Rust are blocking, which poses a large issue for my application since it's very heavily focused on speed and efficiency.

Description

To overcome this, I devised a system that my application can use to fire off database queries asynchronously. It works by starting up multiple threads on which Postgres connections are created. The system distributes new queries to available connections, handles to which are stored in a VecDeque as Receiver objects.

If all the connections are busy, the query is stored in a different queue and popped out by the worker threads as they complete their previous queries. Both the connection handle queue and the query queue are held in Arc<Mutex<VecDeque>> objects to allow them to be accessed by different threads.


I've tested the system to ensure that it's indeed asynchronous, but I wanted to ask around and see if there was anything I could do to make it better either syntactically or performance-wise.

use std::collections::VecDeque;
use std::thread;
use std::sync::{Arc, Mutex};

use postgres;
use futures::stream::{Stream, channel, Sender, Receiver};
use futures::{Future, oneshot, Complete};

use transport::postgres::get_client;

// helper types to keep function declarations clean
type QueryError = postgres::error::Error;
type SenderQueue = Arc<Mutex<VecDeque<Sender<(String, Complete<()>), ()>>>>;
type QueryQueue = Arc<Mutex<VecDeque<String>>>;

pub struct QueryServer {
    conn_count: usize, // how many connections to open
    query_queue: QueryQueue, // internal query queue
    conn_queue: SenderQueue, // senders for idle query threads
}

// locks the QueryQueue and returns a queued query, if there are any.
fn try_get_new_query(query_queue: QueryQueue) -> Option<String> {
    let mut qq_inner = query_queue.lock().unwrap();
    // there is a queued query
    if !qq_inner.is_empty() {
        return Some(qq_inner.pop_front().unwrap())
    }else{
        // No new queries
        return None
    }
}

// executes the query and blocks the calling thread until it completes
#[allow(unused_must_use)]
fn execute_query(query: String, client: &postgres::Connection) {
    client.execute(query.as_str(), &[])
        /*.map_err(|err| println!("Error saving tick: {:?}", err) )*/;
}

// Creates a query processor that awaits requests
fn init_query_processor(rx: Receiver<(String, Complete<()>), ()>, query_queue: QueryQueue){
    // get a connection to the postgres database
    let client = get_client().expect("Couldn't create postgres connection.");
    // Handler for new queries from main thread
    // This blocks the worker thread until a new message is received
    // .wait() consumes the stream immediately, so the main thread has to wait
    // for the worker to push a message saying it's done before sending more messages
    for tup in rx.wait() {
        let (query, done_tx) = tup.unwrap();
        execute_query(query, &client);
        // keep trying to get queued queries to exeucte until the queue is empty
        while let Some(new_query) = try_get_new_query(query_queue.clone()) {
            execute_query(new_query, &client);
        }
        // Let the main thread know it's safe to use the sender again
        // This essentially indicates that the worker thread is idle
        done_tx.complete(());
    }
}

impl QueryServer {
    pub fn new(conn_count: usize) -> QueryServer {
        let mut conn_queue = VecDeque::with_capacity(conn_count);
        let query_queue = Arc::new(Mutex::new(VecDeque::new()));
        for _ in 0..conn_count {
            // channel for getting the Sender back from the worker thread
            let (tx, rx) = channel::<(String, Complete<()>), ()>();
            let qq_copy = query_queue.clone();
            thread::spawn(move || { init_query_processor(rx, qq_copy) });
            // store the sender which can be used to send queries
            // to the worker in the connection queue
            conn_queue.push_back(tx);
        }

        QueryServer {
            conn_count: conn_count,
            query_queue: query_queue,
            conn_queue: Arc::new(Mutex::new(conn_queue))
        }
    }

    // queues up a query to execute that doesn't return a result
    pub fn execute(&mut self, query: String) {
        // no connections available
        let temp_lock_res = self.conn_queue.lock().unwrap().is_empty();
        // Force the guard locking conn_queue to go out of scope
        // this prevents the lock from being held through the entire if/else
        let copy_res = temp_lock_res.clone();
        if copy_res {
            // push query to the query queue
            self.query_queue.lock().unwrap().push_back(query);
        }else{
            let tx = self.conn_queue.lock().unwrap().pop_front().unwrap();
            let cq_clone = self.conn_queue.clone();
            // future for notifying main thread when query is done and worker is idle
            let (c, o) = oneshot::<()>();
            tx.send(Ok((query, c))).and_then(|new_tx| {
                // Wait until the worker thread signals that it is idle
                o.and_then(move |_| {
                    // Put the Sender for the newly idle
                    // worker into the connection queue
                    cq_clone.lock().unwrap().push_back(new_tx);
                    Ok(())
                }).forget();
                Ok(())
            }).forget();
        }
    }
}
\$\endgroup\$
1
  • 1
    \$\begingroup\$ Welcome to Code Review! I hope you get some great answers. \$\endgroup\$ – Phrancis Aug 28 '16 at 1:10
3
\$\begingroup\$

Here are a few shallow comments on the code.

This code

let mut qq_inner = query_queue.lock().unwrap();
// there is a queued query
if !qq_inner.is_empty() {
    return Some(qq_inner.pop_front().unwrap())
}else{
    // No new queries
    return None
}

looks like a verbose way of writing

query_queue.lock().unwrap().pop_front()

The signature

fn try_get_new_query(query_queue: QueryQueue) -> Option<String>

means that it takes an Arc, but the code has no interest in ownership so should really be

fn try_get_new_query(query_queue: &Mutex<VecDeque<String>>) -> Option<String>

instead. The inner type is better labelled QueryQueue and the outer types are Arc<QueryQueue> and &QueryQueue. This gives your API more flexibility, and since Arc dereferences to a normal &-reference, this does not prevent any prior use-cases.

This code

#[allow(unused_must_use)]
fn execute_query(query: String, client: &postgres::Connection) {
    client.execute(query.as_str(), &[])
        /*.map_err(|err| println!("Error saving tick: {:?}", err) )*/;
}

firstly has a very strange comment, since map_err should almost never involve println!. I could understand logging, but not if it throws away the error.

Secondly, it'd be nicer to write as

fn execute_query(query: String, client: &postgres::Connection) {
    let _ = client.execute(query.as_str(), &[]);
}

I'd also write it as

fn execute_query(query: &str, client: &postgres::Connection) {
    let _ = client.execute(query, &[]);
}

since it also doesn't really make calling any harder and increases flexibility.

This is a totally optional style point, but IMO,

thread::spawn(move || { init_query_processor(rx, qq_copy) });

is nicer as

thread::spawn(move || init_query_processor(rx, qq_copy));

In execute, there's another minor style point which is that }else{ should be } else {.

Personally, given that let (c, o) = oneshot::<()>(); only uses c and o once, there's no great need to shorten them so much. YMMV.

\$\endgroup\$
1
  • \$\begingroup\$ Really great information; I appreciate the amazing reply! Now that I look through my code, there were a lot of relics from debugging and previous versions that would have made sense to be changed. Thanks again for taking the time to review it for me. \$\endgroup\$ – Ameo Aug 28 '16 at 18:27

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.