Remove num_tokens_used from VQueueMeta

Author: AhmedSoliman

crates/storage-api/src/vqueue_table/metadata.rs

@@ -74,21 +74,18 @@ pub struct VQueueMeta {
     /// being added to the vqueue. The capacity configuration will limit this value.
     #[bilrost(tag(2))]
     pub(crate) length: u32,
-    /// Number of concurrency tokens being used
-    #[bilrost(tag(3))]
-    pub(crate) num_tokens_used: u32,
     /// The number of entries waiting to be dequeued. The vector index implies the priority
-    #[bilrost(tag(4), encoding(packed))]
+    #[bilrost(tag(3), encoding(packed))]
     pub(crate) num_waiting: [u32; EffectivePriority::NUM_PRIORITIES],
-    #[bilrost(tag(5))]
+    #[bilrost(tag(4))]
     pub(crate) num_running: u32,
-    #[bilrost(tag(6))]
+    #[bilrost(tag(5))]
     pub(crate) stats: VQueueStatistics,
-    #[bilrost(tag(7))]
+    #[bilrost(tag(6))]
     pub(crate) scope: Option<Scope>,
-    #[bilrost(tag(8))]
+    #[bilrost(tag(7))]
     pub(crate) limit_key: LimitKey<ReString>,
-    #[bilrost(tag(9))]
+    #[bilrost(tag(8))]
     lock_name: Option<LockName>,
 }
 
@@ -101,7 +98,6 @@ impl VQueueMeta {
         Self {
             is_paused: false,
             length: 0,
-            num_tokens_used: 0,
             num_waiting: [0; EffectivePriority::NUM_PRIORITIES],
             num_running: 0,
             stats: VQueueStatistics::new(WallClock::recent_ms()),
@@ -131,10 +127,6 @@ impl VQueueMeta {
         &self.limit_key
     }
 
-    pub fn tokens_used(&self) -> u32 {
-        self.num_tokens_used
-    }
-
     pub fn len(&self) -> u32 {
         self.length
     }
@@ -199,14 +191,6 @@ impl VQueueMeta {
         self.num_waiting[priority as usize] -= 1;
     }
 
-    fn acquire_token(&mut self) {
-        self.num_tokens_used += 1;
-    }
-
-    fn release_token(&mut self) {
-        self.num_tokens_used = self.num_tokens_used.saturating_sub(1);
-    }
-
     pub fn apply_update(&mut self, update: &Update) -> anyhow::Result<()> {
         debug_assert!(self.length >= self.total_waiting());
         let now = update.ts;
@@ -244,13 +228,8 @@ impl VQueueMeta {
                     let latency_ms = now.to_unix_millis().saturating_sub_ms(visible_since);
                     self.stats.update_avg_queue_duration(latency_ms);
                 }
-
-                if !priority.token_held() {
-                    self.acquire_token();
-                }
             }
             Action::Park {
-                should_release_concurrency_token,
                 priority,
                 previous_stage,
             } => {
@@ -269,12 +248,6 @@ impl VQueueMeta {
                         // do nothing.
                     }
                 }
-
-                if should_release_concurrency_token && priority.token_held() {
-                    // Release the token immediately on park if this entry doesn't require
-                    // holding while parked.
-                    self.release_token();
-                }
             }
             Action::WakeUp { priority } => {
                 debug_assert!(self.length > 0);
@@ -306,10 +279,6 @@ impl VQueueMeta {
                         // do nothing.
                     }
                 }
-
-                if priority.token_held() {
-                    self.release_token();
-                }
             }
         }
         Ok(())
@@ -379,7 +348,6 @@ pub enum Action {
     #[bilrost(tag(4), message)]
     Park {
         priority: EffectivePriority,
-        should_release_concurrency_token: bool,
         previous_stage: Stage,
     },
     // Wake up after pause or suspend.

crates/vqueues/src/lib.rs

@@ -433,15 +433,12 @@ where
 
     /// Park an entry
     ///
-    /// If `should_release_concurrency_token` is true, the parked entry will release its token
-    ///
     /// Returns `true` if the entry was found in inbox and parked correctly, `false` otherwise.
     pub fn park(
         &mut self,
         at: UniqueTimestamp,
         card: &EntryCard,
         previous_stage: Stage,
-        should_release_concurrency_token: bool,
     ) -> Result<bool, StorageError> {
         let meta = self.cache.get_mut(self.cache_key).unwrap();
 
@@ -457,7 +454,6 @@ where
         let update = metadata::Update::new(
             at,
             metadata::Action::Park {
-                should_release_concurrency_token,
                 priority: card.priority,
                 previous_stage,
             },
@@ -485,16 +481,10 @@ where
             unreachable!("Cannot remove an item from a dormant vqueue");
         }
 
-        let mut modified_card = card.clone();
-        if should_release_concurrency_token && card.priority.token_held() {
-            // adjust the priority to reflect releasing the token
-            modified_card.priority = EffectivePriority::Started;
-        }
-
         self.storage
-            .put_inbox_entry(meta.vqueue_id(), Stage::Park, &modified_card);
+            .put_inbox_entry(meta.vqueue_id(), Stage::Park, card);
         self.storage
-            .put_vqueue_entry_state(meta.vqueue_id(), &modified_card, Stage::Park, ());
+            .put_vqueue_entry_state(meta.vqueue_id(), card, Stage::Park, ());
 
         if let Some(collector) = self.action_collector.as_deref_mut() {
             collector.push(A::from(event));

crates/vqueues/src/scheduler.rs

@@ -54,8 +54,6 @@ pub struct VQueueSchedulerStatus {
     pub remaining_running: u32,
     /// Number of items waiting in the inbox stage.
     pub waiting_inbox: u32,
-    /// Number of concurrency tokens currently in use by this vqueue.
-    pub tokens_used: u32,
     /// The current scheduling status of this vqueue.
     pub status: SchedulingStatus,
 }

crates/vqueues/src/scheduler/drr.rs

@@ -352,7 +352,6 @@ impl<S: VQueueStore> DRRScheduler<S> {
                 wait_stats: qstate.get_head_wait_stats(),
                 remaining_running: qstate.num_remaining_in_running_stage(),
                 waiting_inbox: qstate.num_waiting_inbox(),
-                tokens_used: qstate.num_tokens_used(),
                 status: self.eligible.get_status(qstate),
             };
 
@@ -370,7 +369,6 @@ impl<S: VQueueStore> DRRScheduler<S> {
             wait_stats: qstate.get_head_wait_stats(),
             remaining_running: qstate.num_remaining_in_running_stage(),
             waiting_inbox: qstate.num_waiting_inbox(),
-            tokens_used: qstate.num_tokens_used(),
             status: self.eligible.get_status(qstate),
         })
     }
@@ -895,7 +893,6 @@ mod tests {
         // 5 left in inbox from scheduler's perspective
         assert_eq!(status.waiting_inbox, 5);
         assert_eq!(status.remaining_running, 0);
-        assert_eq!(status.tokens_used, 2);
 
         // Now enqueue high priority item
         let mut txn = rocksdb.transaction();
@@ -948,7 +945,6 @@ mod tests {
         // we added one, and took 2 (5 + 1 - 2 = 4)
         assert_eq!(status.waiting_inbox, 4);
         assert_eq!(status.remaining_running, 0);
-        assert_eq!(status.tokens_used, 4);
         // let's confirm all the items
         events.clear();
         let mut txn = rocksdb.transaction();
@@ -1144,14 +1140,12 @@ mod tests {
         );
         assert_eq!(status.waiting_inbox, 1);
         assert_eq!(status.remaining_running, 0);
-        assert_eq!(status.tokens_used, 2);
 
         // qid2 is exhausted
         let status = scheduler.get_status(&qid2).unwrap();
         assert_eq!(status.status, SchedulingStatus::Empty);
         assert_eq!(status.waiting_inbox, 0);
         assert_eq!(status.remaining_running, 0);
-        assert_eq!(status.tokens_used, 1);
 
         // Pop the permit from the scheduler to release the concurrency token.
         // Only inbox items acquire permits, running items yield without permits.

crates/vqueues/src/scheduler/vqueue_state.rs

@@ -417,8 +417,4 @@ impl<S: VQueueStore> VQueueState<S> {
             .total_waiting()
             .saturating_sub(self.unconfirmed_assignments.len() as u32)
     }
-
-    pub fn num_tokens_used(&self) -> u32 {
-        self.meta.tokens_used() + self.unconfirmed_assignments.len() as u32
-    }
 }

crates/worker/src/partition/state_machine/lifecycle/paused.rs

@@ -10,7 +10,7 @@
 
 use crate::debug_if_leader;
 use crate::partition::state_machine::lifecycle::event::ApplyEventCommand;
-use crate::partition::state_machine::{CommandHandler, Error, ParkCause, StateMachineApplyContext};
+use crate::partition::state_machine::{CommandHandler, Error, StateMachineApplyContext};
 use restate_storage_api::invocation_status_table::{
     InvocationStatus, ReadInvocationStatusTable, WriteInvocationStatusTable,
 };
@@ -58,12 +58,8 @@ where
         debug_if_leader!(ctx.is_leader, "Paused the invocation");
 
         if Configuration::pinned().common.experimental_enable_vqueues {
-            ctx.vqueue_park_invocation(
-                &self.invocation_id,
-                &invoked_meta.invocation_target,
-                ParkCause::Pause,
-            )
-            .await?;
+            ctx.vqueue_park_invocation(&self.invocation_id, &invoked_meta.invocation_target)
+                .await?;
         }
 
         let mut invocation_status = InvocationStatus::Paused(invoked_meta);

crates/worker/src/partition/state_machine/lifecycle/suspend.rs

@@ -8,7 +8,7 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
-use crate::partition::state_machine::{CommandHandler, Error, ParkCause, StateMachineApplyContext};
+use crate::partition::state_machine::{CommandHandler, Error, StateMachineApplyContext};
 use restate_storage_api::invocation_status_table::{InvocationStatus, WriteInvocationStatusTable};
 use restate_storage_api::journal_table_v2::ReadJournalTable;
 use restate_storage_api::lock_table::WriteLockTable;
@@ -88,7 +88,6 @@ where
                 ctx.vqueue_park_invocation(
                     &self.invocation_id,
                     &in_flight_invocation_metadata.invocation_target,
-                    ParkCause::Suspend,
                 )
                 .await?;
             }

crates/worker/src/partition/state_machine/mod.rs

@@ -84,9 +84,8 @@ use restate_types::invocation::{
     AttachInvocationRequest, IngressInvocationResponseSink, InvocationMutationResponseSink,
     InvocationQuery, InvocationResponse, InvocationTarget, InvocationTargetType,
     InvocationTermination, JournalCompletionTarget, NotifySignalRequest, ResponseResult,
-    ServiceInvocation, ServiceInvocationResponseSink, ServiceInvocationSpanContext, ServiceType,
-    Source, SubmitNotificationSink, TerminationFlavor, VirtualObjectHandlerType,
-    WorkflowHandlerType,
+    ServiceInvocation, ServiceInvocationResponseSink, ServiceInvocationSpanContext, Source,
+    SubmitNotificationSink, TerminationFlavor, VirtualObjectHandlerType, WorkflowHandlerType,
 };
 use restate_types::invocation::{InvocationInput, SpanRelation};
 use restate_types::journal::Completion;
@@ -4475,12 +4474,8 @@ impl<S> StateMachineApplyContext<'_, S> {
         metadata.timestamps.update(self.record_created_at);
 
         if Configuration::pinned().common.experimental_enable_vqueues {
-            self.vqueue_park_invocation(
-                &invocation_id,
-                &metadata.invocation_target,
-                ParkCause::Suspend,
-            )
-            .await?;
+            self.vqueue_park_invocation(&invocation_id, &metadata.invocation_target)
+                .await?;
         }
 
         self.storage
@@ -5115,7 +5110,6 @@ impl<S> StateMachineApplyContext<'_, S> {
         &mut self,
         invocation_id: &InvocationId,
         invocation_target: &InvocationTarget,
-        cause: ParkCause,
     ) -> Result<(), Error>
     where
         S: WriteVQueueTable + WriteLockTable + ReadVQueueTable,
@@ -5150,34 +5144,10 @@ impl<S> StateMachineApplyContext<'_, S> {
 
         let now = UniqueTimestamp::from_unix_millis_unchecked(self.record_created_at);
 
-        let should_release_concurrency_token = match cause {
-            ParkCause::Suspend => {
-                // Always hold on to your concurrency token until the invocation is completed if
-                // we are suspending for all types (services, VOs, workflows). This has the benefit
-                // of an easy to reason about concurrency model for our users. The downside is that
-                // callers might deadlock if they call a limited service which has no more
-                // concurrency tokens left and there is a cyclic dependency (e.g. a limited service
-                // calling itself).
-                false
-            }
-            ParkCause::Pause => {
-                // We release the concurrency token in case we are pausing a service because
-                // unpausing requires human intervention, and we don't want to block other service
-                // invocations.
-                //
-                // Note that we don't do this for paused VOs and workflows because they need to
-                // ensure that no other instance can run while they hold their lock. Technically,
-                // we still have the service_status_table which stores the locking information, and
-                // we need to keep things in sync until we decide what to do with this table.
-                matches!(invocation_target.service_ty(), ServiceType::Service)
-            }
-        };
-
         vqueue.park(
             now,
             &entry_state_header.current_entry_card(),
             entry_state_header.stage(),
-            should_release_concurrency_token,
         )?;
 
         Ok(())
@@ -5378,15 +5348,6 @@ impl<S> StateMachineApplyContext<'_, S> {
     }
 }
 
-/// Cause for parking an invocation
-#[derive(Debug)]
-enum ParkCause {
-    /// The invocation suspends to await completion or signals
-    Suspend,
-    /// The invocation pauses because it depleted it retries or was manually paused
-    Pause,
-}
-
 // To write completions in the effects log
 struct CompletionResultFmt<'a>(&'a CompletionResult);