undo batcher changes

foundations/src/batcher/mod.rs

@@ -5,6 +5,7 @@ use std::sync::atomic::{AtomicU64, AtomicUsize};
 use std::sync::Arc;
 
 use tokio::sync::OnceCell;
+use tracing::Instrument;
 
 pub mod dataloader;
 
@@ -224,15 +225,14 @@ impl Drop for CancelOnDrop {
 }
 
 struct BatcherInner<T: BatchOperation> {
- semaphore: Arc<tokio::sync::Semaphore>,
+ semaphore: tokio::sync::Semaphore,
  notify: tokio::sync::Notify,
  sleep_duration: AtomicU64,
  batch_id: AtomicU64,
  max_batch_size: AtomicUsize,
  operation: T,
  name: String,
  active_batch: tokio::sync::RwLock<Option<Batch<T>>>,
- queued_batches: tokio::sync::mpsc::Sender<Batch<T>>,
 }
 
 struct Batch<T: BatchOperation> {
@@ -288,14 +288,18 @@ impl<E: std::error::Error> From<E> for BatcherError<E> {
 
 impl<T: BatchOperation + 'static + Send + Sync> Batch<T> {
  #[tracing::instrument(skip_all, fields(name = %inner.name))]
- async fn run(self, inner: Arc<BatcherInner<T>>, ticket: tokio::sync::OwnedSemaphorePermit) {
+ async fn run(self, inner: Arc<BatcherInner<T>>) {
  self.results
  .get_or_init(|| async move {
- inner.operation.process(self.ops).await.map_err(BatcherError::Batch)
+ let _ticket = inner
+ .semaphore
+ .acquire()
+ .instrument(tracing::debug_span!("Semaphore"))
+ .await
+ .map_err(|_| BatcherError::AcquireSemaphore)?;
+ Ok(inner.operation.process(self.ops).await.map_err(BatcherError::Batch)?)
  })
  .await;
-
- drop(ticket);
  }
 }
 
@@ -308,8 +312,8 @@ pub struct BatcherConfig {
 }
 
 impl<T: BatchOperation + 'static + Send + Sync> BatcherInner<T> {
- fn spawn_batch(self: &Arc<Self>, batch: Batch<T>, ticket: tokio::sync::OwnedSemaphorePermit) {
- tokio::spawn(batch.run(self.clone(), ticket));
+ fn spawn_batch(self: &Arc<Self>, batch: Batch<T>) {
+ tokio::spawn(batch.run(self.clone()));
  }
 
  fn new_batch(&self) -> Batch<T> {
@@ -330,7 +334,6 @@ impl<T: BatchOperation + 'static + Send + Sync> BatcherInner<T> {
  let mut waiters = vec![];
  let mut batch = self.active_batch.write().await;
  let max_documents = self.max_batch_size.load(std::sync::atomic::Ordering::Relaxed);
- let mut batches = vec![];
 
  for document in T::Mode::filter_item_iter(documents) {
  if batch
@@ -339,7 +342,7 @@ impl<T: BatchOperation + 'static + Send + Sync> BatcherInner<T> {
  .unwrap_or(true)
  {
  if let Some(b) = batch.take() {
- batches.push(b);
+ self.spawn_batch(b);
  }
 
  *batch = Some(self.new_batch());
@@ -363,10 +366,6 @@ impl<T: BatchOperation + 'static + Send + Sync> BatcherInner<T> {
  T::Mode::input_add(&mut b.ops, tracker, document);
  }
 
- for batch in batches {
- self.queued_batches.send(batch).await.ok();
- }
-
  waiters
  }
 }
@@ -375,11 +374,8 @@ impl<T: BatchOperation + 'static + Send + Sync> Batcher<T> {
  pub fn new(operation: T) -> Self {
  let config = operation.config();
 
- let (tx, mut rx) = tokio::sync::mpsc::channel(64);
-
  let inner = Arc::new(BatcherInner {
- semaphore: Arc::new(tokio::sync::Semaphore::new(config.concurrency)),
- queued_batches: tx.clone(),
+ semaphore: tokio::sync::Semaphore::new(config.concurrency),
  notify: tokio::sync::Notify::new(),
  batch_id: AtomicU64::new(0),
  active_batch: tokio::sync::RwLock::new(None),
@@ -393,43 +389,21 @@ impl<T: BatchOperation + 'static + Send + Sync> Batcher<T> {
  inner: inner.clone(),
  _auto_loader_abort: CancelOnDrop(
  tokio::task::spawn(async move {
- let mut next_wakeup = None;
  loop {
- tokio::select! {
- Some(batch) = rx.recv() => {
- let ticket = inner.semaphore.clone().acquire_owned().await.unwrap();
- inner.spawn_batch(batch, ticket);
- },
- _ = async {
- if let Some(expires_at) = next_wakeup {
- tokio::time::sleep_until(expires_at).await;
- } else {
- inner.notify.notified().await;
- }
- } => {},
- _ = inner.notify.notified() => {},
- }
-
+ inner.notify.notified().await;
  let Some((id, expires_at)) = inner.active_batch.read().await.as_ref().map(|b| (b.id, b.expires_at))
  else {
  continue;
  };
 
  if expires_at > tokio::time::Instant::now() {
- next_wakeup = Some(expires_at);
- continue;
- } else {
- next_wakeup = None;
+ tokio::time::sleep_until(expires_at).await;
  }
-
- let mut batch = inner.active_batch.write().await;
- let batch = if batch.as_ref().is_some_and(|b| b.id == id) {
- batch.take().unwrap()
- } else {
- continue;
- };
 
- tx.send(batch).await.ok();
+ let mut batch = inner.active_batch.write().await;
+ if batch.as_ref().is_some_and(|b| b.id == id) {
+ inner.spawn_batch(batch.take().unwrap());
+ }
  }
  })
  .abort_handle(),