Skip to content

Commit

Permalink
Add low-level packet sending to WorkerJob
Browse files Browse the repository at this point in the history
This will allow workers to construct any Gearman packet they need to
send, which can be a little dangerous but may be useful for
compatibility with other clients. We specifically use it here to send a
broken packet to test our client's handling of it.

With this we re-work some error handling to make the client more
resilient, and test that it is, in fact, more resilient.
  • Loading branch information
SpamapS committed Jun 18, 2024
1 parent 7ba4803 commit 5b77bfc
Show file tree
Hide file tree
Showing 4 changed files with 122 additions and 63 deletions.
69 changes: 36 additions & 33 deletions rustygear/src/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -227,21 +227,15 @@ impl WorkerJob {
payload.put_u8(b'\0');
payload.extend(denominator.as_bytes());
let packet = new_res(WORK_STATUS, payload.freeze());
self.send_packet(packet).await.map_err(|e| {
if e.is::<std::io::Error>() {
*e.downcast::<std::io::Error>().expect("downcast after is")
} else {
std::io::Error::new(io::ErrorKind::Other, e.to_string())
}
})
self.send_packet(packet).await
}

async fn send_packet(&mut self, packet: Packet) -> Result<(), Box<dyn Error>> {
pub async fn send_packet(&mut self, packet: Packet) -> Result<(), io::Error> {
match self.sink_tx.send(packet).await {
Err(_) => Err(Box::new(io::Error::new(
Err(_) => Err(io::Error::new(
io::ErrorKind::NotConnected,
"Connection closed",
))),
)),
Ok(_) => Ok(()),
}
}
Expand All @@ -250,7 +244,7 @@ impl WorkerJob {
///
/// This method is typically called by the [Client::work] method upon return
/// of an error from the assigned closure.
pub async fn work_fail(&mut self) -> Result<(), Box<dyn Error>> {
pub async fn work_fail(&mut self) -> Result<(), io::Error> {
let packet = new_res(WORK_FAIL, self.handle.clone());
self.send_packet(packet).await
}
Expand All @@ -259,7 +253,7 @@ impl WorkerJob {
///
/// This method is typically called by the [Client::work] method upon return of
/// the assigned closure.
pub async fn work_complete(&mut self, response: Vec<u8>) -> Result<(), Box<dyn Error>> {
pub async fn work_complete(&mut self, response: Vec<u8>) -> Result<(), io::Error> {
let mut payload = BytesMut::with_capacity(self.handle.len() + 1 + self.payload.len());
payload.extend(self.handle.clone());
payload.put_u8(b'\0');
Expand Down Expand Up @@ -449,29 +443,38 @@ impl Client {
let tx = tx.clone();
while let Some(frame) = stream.next().await {
trace!("Frame read: {:?}", frame);
let response = match frame {
Err(e) => Err(e.to_string()),
Ok(frame) => {
let handler = handler.clone();
debug!("Locking handler");
let mut handler = handler;
debug!("Locked handler");
handler
.call(frame)
.map_err(|e| e.to_string())
}
};
match response {
Err(e) => {
error!("conn dropped?: {}", e);
break;
}
Ok(response) => {
if let Err(_) = tx.send(response).await
{
error!("receiver dropped")
// This lexical scope is needed because the compiler can't figure out
// that response's error is dropped before the await.
// See: https://github.com/rust-lang/rust/pull/107421 for the fix
// which is only in nightly as of this writing.
let packet = {
let response = match frame {
Err(e) => {
Err(Box::new(e) as Box<dyn Error>)
}
Ok(frame) => {
let handler = handler.clone();
debug!("Locking handler");
let mut handler = handler;
debug!("Locked handler");
handler.call(frame)
} //.map_err(|e| e)
// Ugh this map_err
};
match response {
Err(e) => {
if e.is::<io::Error>() {
error!("conn dropped?: {}", e);
break;
}
error!("There was a non-fatal error while processing a packet: {}", e);
continue;
}
Ok(packet) => packet,
}
};
if let Err(_) = tx.send(packet).await {
warn!("receiver dropped")
}
}
reader_conns
Expand Down
30 changes: 30 additions & 0 deletions rustygeard/src/testutil.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@ use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::mpsc::{sync_channel, Receiver, SyncSender};
use std::time::Duration;

use rustygear::client::Client;

use crate::server::GearmanServer;

pub struct ServerGuard {
Expand Down Expand Up @@ -66,3 +68,31 @@ pub fn start_test_server() -> Option<ServerGuard> {
}
return None;
}

pub async fn connect(addr: &SocketAddr) -> Client {
connect_with_client_id(addr, "tests").await
}

pub async fn connect_with_client_id(addr: &SocketAddr, client_id: &'static str) -> Client {
let client = Client::new().add_server(&addr.to_string());
client
.set_client_id(client_id)
.connect()
.await
.expect("Failed to connect to server")
}

pub async fn worker(addr: &SocketAddr) -> Client {
connect(addr)
.await
.can_do("testfunc", |workerjob| {
Ok(format!(
"worker saw {} with unique [{}]",
String::from_utf8_lossy(workerjob.payload()),
String::from_utf8_lossy(workerjob.unique())
)
.into_bytes())
})
.await
.expect("CAN_DO should work")
}
32 changes: 2 additions & 30 deletions rustygeard/tests/client.rs
Original file line number Diff line number Diff line change
@@ -1,37 +1,9 @@
use std::{io::ErrorKind, net::SocketAddr, time::Duration};
use std::{io::ErrorKind, time::Duration};

use rustygear::client::{Client, WorkUpdate, WorkerJob};
use rustygeard::testutil::start_test_server;
use rustygeard::testutil::{connect, connect_with_client_id, start_test_server, worker};
use tokio::time::{sleep, timeout};

async fn connect(addr: &SocketAddr) -> Client {
connect_with_client_id(addr, "tests").await
}

async fn connect_with_client_id(addr: &SocketAddr, client_id: &'static str) -> Client {
let client = Client::new().add_server(&addr.to_string());
client
.set_client_id(client_id)
.connect()
.await
.expect("Failed to connect to server")
}

async fn worker(addr: &SocketAddr) -> Client {
connect(addr)
.await
.can_do("testfunc", |workerjob| {
Ok(format!(
"worker saw {} with unique [{}]",
String::from_utf8_lossy(workerjob.payload()),
String::from_utf8_lossy(workerjob.unique())
)
.into_bytes())
})
.await
.expect("CAN_DO should work")
}

#[tokio::test]
async fn test_client_connects() {
let server = start_test_server().unwrap();
Expand Down
54 changes: 54 additions & 0 deletions rustygeard/tests/errors.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
use std::time::Duration;

use bytes::BytesMut;
use rustygear::{
client::{Client, WorkUpdate, WorkerJob},
constants::WORK_STATUS,
util::new_req,
};
use rustygeard::testutil::{connect, connect_with_client_id, start_test_server};
use tokio::time::timeout;

#[tokio::test]
async fn test_worker_sends_bad_work_status() {
let server = start_test_server().unwrap();
let worker = connect_with_client_id(server.addr(), "status-worker").await;
fn sends_status(work: &mut WorkerJob) -> Result<Vec<u8>, std::io::Error> {
let rt = tokio::runtime::Builder::new_current_thread()
.build()
.unwrap();
let mut data = BytesMut::new();
data.extend(work.handle());
data.extend(b"\0notnumbers\0notnumdenom");
let packet = new_req(WORK_STATUS, data.freeze());
rt.block_on(work.send_packet(packet))?;
Ok("Done".into())
}
let mut worker = worker
.can_do("statusfunc", sends_status)
.await
.expect("CAN_DO should succeed");
let mut client: Client = connect(server.addr()).await;
let mut job = client
.submit("statusfunc", b"statuspayload")
.await
.expect("Submit should succeed");
worker
.do_one_job()
.await
.expect("One job should be completed");
// We'll ignore the broken status packet and still get the WorkComplete
// The timeout is here to protect the test suite because response can
// Easily get disconnected from things if errors aren't handled right.
let response = timeout(Duration::from_millis(500), job.response())
.await
.expect("Response happens within 500ms")
.expect("Response to non-background job should not error");
assert!(matches!(
response,
WorkUpdate::Complete {
handle: _,
payload: _
}
));
}

0 comments on commit 5b77bfc

Please sign in to comment.