Adding a negative-test for coalesce

The has_job method exposed a weakness in WrappingHashset which was that
two different iterators would interfere with each other. So the library
was updated to 0.5, and we should only use that version and higher in
the future.

Author: SpamapS

Cargo.lock

@@ -117,7 +117,7 @@ impl ConnHandler {
     }
 
     fn handle_noop(&self) -> Result<Packet, io::Error> {
-        Ok(new_req(GRAB_JOB, Bytes::new()))
+        Ok(new_req(GRAB_JOB_UNIQ, Bytes::new()))
     }
 
     fn handle_no_job(&self) -> Result<Packet, io::Error> {

rustygear/src/conn.rs

@@ -17,7 +17,7 @@ tokio = { version = "1.37.0", features = ["full"] }
 tokio-util = { version = "0.7.10", features = ["codec"] }
 tower-service = "0.3"
 futures = "0.3"
-wrappinghashset = ">=0.4.1"
+wrappinghashset = ">=0.5.0"
 clap = { version = "4.5" , features = ["cargo"] }
 futures-sink = "0.3"
 uuid = { version = "1.8.0", features = ["v4"] }

rustygeard/Cargo.toml

@@ -23,6 +23,7 @@ pub trait HandleJobStorage {
     fn new_job_storage() -> SharedJobStorage;
     fn coalesce_unique(&mut self, unique: &Bytes, remote: Option<usize>) -> Option<Bytes>;
     fn add_job(&mut self, job: Arc<Job>, priority: JobQueuePriority, remote: Option<usize>);
+    fn has_job(&self, worker: &mut Worker) -> bool;
     fn get_job(&mut self, worker: &mut Worker) -> Option<Arc<Job>>;
 }
 
@@ -154,6 +155,28 @@ impl HandleJobStorage for SharedJobStorage {
         );
     }
 
+    fn has_job(&self, worker: &mut Worker) -> bool {
+        let storage = self.lock().unwrap();
+        for func in worker.iter() {
+            if let Some(prios) = storage.queues.get(&func) {
+                let mut prio = 0;
+                for q in prios {
+                    prio = prio + 1;
+                    for job in q.iter() {
+                        if let Some(_job) = job.upgrade() {
+                            debug!(
+                                "Worker {:?} has a job in func {:?} prio {}",
+                                worker, func, prio
+                            );
+                            return true;
+                        }
+                    }
+                }
+            }
+        }
+        return false;
+    }
+
     fn get_job(&mut self, worker: &mut Worker) -> Option<Arc<Job>> {
         let mut storage = self.lock().unwrap();
         let mut job: Option<Arc<Job>> = None;

rustygeard/src/queues.rs

@@ -225,8 +225,14 @@ impl GearmanService {
     fn handle_pre_sleep(&self) -> Result<Packet, io::Error> {
         let worker = self.worker.clone();
         let ref mut w = worker.lock().unwrap();
-        self.workers.clone().sleep(w, self.conn_id);
-        Ok(no_response())
+        let queues = self.queues.clone();
+        if queues.has_job(w) {
+            // We said NO_JOB, but now we have a job, so we need to wake it back up immediately.
+            Ok(new_noop())
+        } else {
+            self.workers.clone().sleep(w, self.conn_id);
+            Ok(no_response())
+        }
     }
 
     fn handle_submit_job(

rustygeard/src/service.rs

@@ -14,7 +14,6 @@ fn test_server_starts() {
 
 #[tokio::test]
 async fn test_server_coalesces_uniqs() {
-    std::env::set_var("RUST_LOG", "debug");
     let server = start_test_server().expect("Starting test server");
     let mut client1 = Client::new()
         .add_server(&server.addr().to_string())
@@ -28,7 +27,7 @@ async fn test_server_coalesces_uniqs() {
         .connect()
         .await
         .expect("Connecting client2");
-    let (tx, rx) = tokio::sync::mpsc::channel(1);
+    let (tx, rx) = tokio::sync::mpsc::channel(2);
     let rx = Arc::new(Mutex::new(rx));
     let server_addr = server.addr().to_string().clone();
     thread::spawn(move || {
@@ -54,7 +53,7 @@ async fn test_server_coalesces_uniqs() {
                 })
                 .await
                 .expect("Sending CAN_DO and setting up worker function")
-                .do_one_job()
+                .work()
                 .await
                 .expect("Doing one job");
         });
@@ -89,4 +88,34 @@ async fn test_server_coalesces_uniqs() {
     } else {
         panic!("Response 1 was not WORK_COMPLETE: {:?}", response1);
     }
+    // And now, make sure different uniqids do not coalesce
+    let mut job1b = client1
+        .submit_unique("uniqfunc", b"uniqid1b", b"")
+        .await
+        .expect("submitting uniqid1b job");
+    let mut job2b = client2
+        .submit_unique("uniqfunc", b"uniqid2b", b"")
+        .await
+        .expect("submitting uniqid2b job");
+    tx.send(()).await.expect("Sending to let the worker finish");
+    tx.send(()).await.expect("Sending to let the worker finish");
+    let response1b = job1b.response().await.expect("Getting response from job1b");
+    let response2b = job2b.response().await.expect("Getting response from job2b");
+    if let WorkUpdate::Complete {
+        handle: handle1b,
+        payload: _,
+    } = response1b
+    {
+        if let WorkUpdate::Complete {
+            handle: handle2b,
+            payload: _,
+        } = response2b
+        {
+            assert_ne!(handle1b, handle2b);
+        } else {
+            panic!("job2b not a WORK_COMPLETE: {:?}", response2b);
+        }
+    } else {
+        panic!("job1b not a WORK_COMPLETE: {:?}", response1b);
+    }
 }