take things off the batch in a more safe way (#6171)

Signed-off-by: Jess Frazelle <github@jessfraz.com>
This commit is contained in:
Jess Frazelle
2025-04-05 20:54:32 -07:00
committed by GitHub
parent 48e1a8ed02
commit 38446b5b2a

View File

@ -93,6 +93,16 @@ pub trait EngineManager: std::fmt::Debug + Send + Sync + 'static {
/// Get the artifact commands that have accumulated so far.
fn artifact_commands(&self) -> Arc<RwLock<Vec<ArtifactCommand>>>;
/// Take the batch of commands that have accumulated so far and clear them.
async fn take_batch(&self) -> Vec<(WebSocketRequest, SourceRange)> {
std::mem::take(&mut *self.batch().write().await)
}
/// Take the batch of end commands that have accumulated so far and clear them.
async fn take_batch_end(&self) -> IndexMap<Uuid, (WebSocketRequest, SourceRange)> {
std::mem::take(&mut *self.batch_end().write().await)
}
/// Clear all artifact commands that have accumulated so far.
async fn clear_artifact_commands(&self) {
self.artifact_commands().write().await.clear();
@ -370,11 +380,11 @@ pub trait EngineManager: std::fmt::Debug + Send + Sync + 'static {
source_range: SourceRange,
) -> Result<OkWebSocketResponseData, crate::errors::KclError> {
let all_requests = if batch_end {
let mut requests = self.batch().read().await.clone();
requests.extend(self.batch_end().read().await.values().cloned());
let mut requests = self.take_batch().await.clone();
requests.extend(self.take_batch_end().await.values().cloned());
requests
} else {
self.batch().read().await.clone()
self.take_batch().await.clone()
};
// Return early if we have no commands to send.
@ -442,11 +452,6 @@ pub trait EngineManager: std::fmt::Debug + Send + Sync + 'static {
}
}
// Throw away the old batch queue.
self.batch().write().await.clear();
if batch_end {
self.batch_end().write().await.clear();
}
self.stats().batches_sent.fetch_add(1, Ordering::Relaxed);
// We pop off the responses to cleanup our mappings.