fixes in async_semaphore for concurrency ID calculation

src/rust/engine/async_semaphore/src/lib.rs

@@ -42,6 +42,7 @@ struct Waiter {
 
 struct Inner {
   waiters: VecDeque<Waiter>,
+  available_ids: VecDeque<usize>,
   available_permits: usize,
   // Used as the source of id in Waiters's because
   // it is monotonically increasing, and only incremented under the mutex lock.
@@ -55,9 +56,15 @@ pub struct AsyncSemaphore {
 
 impl AsyncSemaphore {
   pub fn new(permits: usize) -> AsyncSemaphore {
+    let mut available_ids = VecDeque::new();
+    for id in 1..=permits {
+      available_ids.push_back(id);
+    }
+
     AsyncSemaphore {
       inner: Arc::new(Mutex::new(Inner {
         waiters: VecDeque::new(),
+        available_ids,
         available_permits: permits,
         next_waiter_id: 0,
       })),
@@ -99,10 +106,11 @@ pub struct Permit {
 impl Drop for Permit {
   fn drop(&mut self) {
     let mut inner = self.inner.lock();
-    inner.available_permits += 1;
-    if let Some(waiter) = inner.waiters.front() {
+    inner.available_ids.push_back(self.id);
+    if let Some(waiter) = inner.waiters.get(inner.available_permits) {
       waiter.waker.wake_by_ref()
     }
+    inner.available_permits += 1;
   }
 }
 
@@ -158,7 +166,8 @@ impl Future for PermitFuture {
               // queue, so we don't have to waste time searching for it in the Drop
               // handler.
               self.waiter_id = None;
-              Some(inner.available_permits)
+              let id = inner.available_ids.pop_front().unwrap_or(0); // The unwrap_or case should never happen.
+              Some(id)
             } else {
               // Don't issue a permit to this task if it isn't at the head of the line,
               // we added it as a waiter above.

src/rust/engine/async_semaphore/src/tests.rs

@@ -56,13 +56,91 @@ async fn correct_semaphor_slot_ids() {
   let id3 = rx3.await.unwrap();
   let id4 = rx4.await.unwrap();
 
-  // Process 1 should get ID 2, then process 2 should run with id 1 and complete, then process 3
-  // should run in the same "slot" as process 2 and get the same id. Process 4 is scheduled
-  // later and gets put into "slot" 2.
-  assert_eq!(id1, 2);
-  assert_eq!(id2, 1);
-  assert_eq!(id3, 1);
-  assert_eq!(id4, 2);
+  // Process 1 should get ID 1, then process 2 should run with id 2 and complete, then process 3
+  // should run in the same "slot" as process 2 and get the same id (2). Process 4 is scheduled
+  // later and gets put into "slot" 1.
+  assert_eq!(id1, 1);
+  assert_eq!(id2, 2);
+  assert_eq!(id3, 2);
+  assert_eq!(id4, 1);
+}
+
+#[tokio::test]
+async fn correct_semaphor_slot_ids_2() {
+  let sema = AsyncSemaphore::new(4);
+  let (tx1, rx1) = oneshot::channel();
+  let (tx2, rx2) = oneshot::channel();
+  let (tx3, rx3) = oneshot::channel();
+  let (tx4, rx4) = oneshot::channel();
+  let (tx5, rx5) = oneshot::channel();
+  let (tx6, rx6) = oneshot::channel();
+  let (tx7, rx7) = oneshot::channel();
+
+  println!("Spawning process 1");
+  tokio::spawn(sema.clone().with_acquired(move |id| async move {
+    println!("Exec process 1");
+    tx1.send(id).unwrap();
+    delay_for(Duration::from_millis(20)).await;
+    future::ready(())
+  }));
+  println!("Spawning process 2");
+  tokio::spawn(sema.clone().with_acquired(move |id| async move {
+    println!("Exec process 2");
+    tx2.send(id).unwrap();
+    delay_for(Duration::from_millis(20)).await;
+    future::ready(())
+  }));
+  println!("Spawning process 3");
+  tokio::spawn(sema.clone().with_acquired(move |id| async move {
+    println!("Exec process 3");
+    tx3.send(id).unwrap();
+    delay_for(Duration::from_millis(20)).await;
+    future::ready(())
+  }));
+  println!("Spawning process 4");
+  tokio::spawn(sema.clone().with_acquired(move |id| async move {
+    println!("Exec process 4");
+    tx4.send(id).unwrap();
+    delay_for(Duration::from_millis(20)).await;
+    future::ready(())
+  }));
+  println!("Spawning process 5");
+  tokio::spawn(sema.clone().with_acquired(move |id| async move {
+    println!("Exec process 5");
+    tx5.send(id).unwrap();
+    delay_for(Duration::from_millis(20)).await;
+    future::ready(())
+  }));
+  println!("Spawning process 6");
+  tokio::spawn(sema.clone().with_acquired(move |id| async move {
+    println!("Exec process 6");
+    tx6.send(id).unwrap();
+    delay_for(Duration::from_millis(20)).await;
+    future::ready(())
+  }));
+  println!("Spawning process 7");
+  tokio::spawn(sema.clone().with_acquired(move |id| async move {
+    println!("Exec process 7");
+    tx7.send(id).unwrap();
+    delay_for(Duration::from_millis(20)).await;
+    future::ready(())
+  }));
+
+  let id1 = rx1.await.unwrap();
+  let id2 = rx2.await.unwrap();
+  let id3 = rx3.await.unwrap();
+  let id4 = rx4.await.unwrap();
+  let id5 = rx5.await.unwrap();
+  let id6 = rx6.await.unwrap();
+  let id7 = rx7.await.unwrap();
+
+  assert_eq!(id1, 1);
+  assert_eq!(id2, 2);
+  assert_eq!(id3, 3);
+  assert_eq!(id4, 4);
+  assert_eq!(id5, 1);
+  assert_eq!(id6, 2);
+  assert_eq!(id7, 3);
 }
 
 #[tokio::test]

src/rust/engine/process_execution/src/lib.rs

@@ -531,6 +531,11 @@ impl CommandRunner for BoundedCommandRunner {
       let name = format!("{}-running", req.workunit_name());
 
       semaphore.with_acquired(move |concurrency_id| {
+        log::debug!(
+          "Running {} under semaphor with concurrency id: {}",
+          desc,
+          concurrency_id
+        );
         let mut metadata = WorkunitMetadata::with_level(Level::Info);
         metadata.desc = Some(desc);