feat: Deposit token logic and add OnlyIfUsed and GoBackN mode (numaproj#2962)

Signed-off-by: Vaibhav Tiwari <vaibhav.tiwari33@gmail.com>

Author: vaibhavtiwari33

rust/numaflow-core/src/pipeline/isb/jetstream/reader.rs

@@ -385,8 +385,14 @@ impl<C: NumaflowTypeConfig> JetStreamReader<C> {
             .map_err(|e| Error::ISB(format!("Failed to acquire semaphore permit: {e}")))?;
 
         // Apply rate limiting if configured.
-        let effective_batch_size = match &self.rate_limiter {
+        let (effective_batch_size, token_grabbed_epoch) = match &self.rate_limiter {
             Some(rate_limiter) => {
+                // Note the epoch at which we got the estimated batch size from the rate limiter
+                let cur_epoch = std::time::SystemTime::now()
+                    .duration_since(std::time::UNIX_EPOCH)
+                    .expect("Time went backwards beyond unix epoch")
+                    .as_secs();
+
                 let tokens_acquired = rate_limiter
                     .acquire_n(Some(batch_size), Some(Duration::from_secs(1)))
                     .await;
@@ -396,9 +402,9 @@ impl<C: NumaflowTypeConfig> JetStreamReader<C> {
                     return Ok(vec![]);
                 }
 
-                effective_batch_size
+                (effective_batch_size, cur_epoch)
             }
-            None => batch_size,
+            None => (batch_size, 0),
         };
 
         let start = Instant::now();
@@ -436,6 +442,23 @@ impl<C: NumaflowTypeConfig> JetStreamReader<C> {
             }
         };
 
+        // deposit the unused tokens back to the rate limiter
+        // utilize the cur_epoch we calculated when we initially acquired the tokens
+        // to ensure that the tokens are deposited for the correct epoch.
+        match &self.rate_limiter {
+            Some(rate_limiter) => {
+                rate_limiter
+                    .deposit_unused(
+                        effective_batch_size
+                            .checked_sub(jetstream_messages.len())
+                            .unwrap_or(0),
+                        token_grabbed_epoch,
+                    )
+                    .await;
+            }
+            None => {}
+        }
+
         pipeline_metrics()
             .jetstream_isb
             .read_time_total