1use crate::{OperationId, WorkspaceId};
10use serde::{Deserialize, Serialize};
11use std::collections::{BTreeMap, BTreeSet};
12use std::future::Future;
13use std::path::{Component, Path, PathBuf};
14use std::sync::Mutex;
15use thiserror::Error;
16
17const WORKSPACE_CONTEXT_VERSION: u32 = 1;
18
19#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize, Deserialize)]
21#[serde(transparent)]
22pub struct WorkspaceContextId(OperationId);
23
24impl WorkspaceContextId {
25 #[must_use]
27 pub fn new() -> Self {
28 Self(OperationId::new())
29 }
30
31 #[must_use]
33 pub const fn from_bytes(bytes: [u8; 16]) -> Self {
34 Self(OperationId::from_bytes(bytes))
35 }
36
37 #[must_use]
39 pub const fn into_bytes(self) -> [u8; 16] {
40 self.0.into_bytes()
41 }
42}
43
44impl Default for WorkspaceContextId {
45 fn default() -> Self {
46 Self::new()
47 }
48}
49
50#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize, Deserialize)]
52#[serde(transparent)]
53pub struct WorkspaceRootId(OperationId);
54
55impl WorkspaceRootId {
56 #[must_use]
58 pub fn new() -> Self {
59 Self(OperationId::new())
60 }
61
62 #[must_use]
64 pub const fn from_bytes(bytes: [u8; 16]) -> Self {
65 Self(OperationId::from_bytes(bytes))
66 }
67
68 #[must_use]
70 pub const fn into_bytes(self) -> [u8; 16] {
71 self.0.into_bytes()
72 }
73}
74
75impl Default for WorkspaceRootId {
76 fn default() -> Self {
77 Self::new()
78 }
79}
80
81#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
83#[serde(rename_all = "snake_case")]
84pub enum WorkspaceContextState {
85 Active,
87 Frozen,
89 Discarded,
91}
92
93#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
95pub struct WorkspaceContextRoot {
96 pub root_id: WorkspaceRootId,
98 pub source_path: PathBuf,
100 pub workspace_id: WorkspaceId,
102 pub workspace_name: String,
104 pub parent_workspace_id: Option<WorkspaceId>,
106 pub mount_path: Option<PathBuf>,
108}
109
110#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
112pub struct WorkspaceContext {
113 pub version: u32,
115 pub revision: u64,
117 pub context_id: WorkspaceContextId,
119 pub parent_context_id: Option<WorkspaceContextId>,
121 pub roots: BTreeMap<WorkspaceRootId, WorkspaceContextRoot>,
123 pub state: WorkspaceContextState,
125}
126
127#[derive(Clone, Copy, Debug, Eq, PartialEq)]
129pub enum WorkspaceRouteKind {
130 Mount,
132 Source,
134}
135
136#[derive(Clone, Debug, Eq, PartialEq)]
138pub struct WorkspaceRoute {
139 pub context_id: WorkspaceContextId,
141 pub root_id: WorkspaceRootId,
143 pub kind: WorkspaceRouteKind,
145 pub prefix: PathBuf,
147 pub relative_path: PathBuf,
149}
150
151pub trait WorkspaceContextStore: Send + Sync {
153 type Error: std::error::Error + Send + Sync + 'static;
155
156 fn load(
158 &self,
159 context_id: WorkspaceContextId,
160 ) -> impl Future<Output = Result<Option<WorkspaceContext>, Self::Error>> + Send;
161
162 fn list(&self) -> impl Future<Output = Result<Vec<WorkspaceContext>, Self::Error>> + Send;
165
166 fn compare_and_swap(
168 &self,
169 context_id: WorkspaceContextId,
170 expected_revision: u64,
171 replacement: WorkspaceContext,
172 ) -> impl Future<Output = Result<bool, Self::Error>> + Send;
173
174 fn compare_and_swap_many(
177 &self,
178 expected_revisions: BTreeMap<WorkspaceContextId, u64>,
179 replacements: Vec<WorkspaceContext>,
180 require_exact_set: bool,
181 ) -> impl Future<Output = Result<bool, Self::Error>> + Send;
182
183 fn discard_subtree(
185 &self,
186 parent: WorkspaceContextId,
187 child: WorkspaceContextId,
188 maximum: u32,
189 ) -> impl Future<Output = Result<WorkspaceContextDiscardOutcome, Self::Error>> + Send;
190}
191
192#[derive(Clone, Debug, Eq, PartialEq)]
194pub enum WorkspaceContextDiscardOutcome {
195 Discarded(Vec<WorkspaceContextId>),
197 IncompatibleState,
199 UnauthorizedParent,
201 AlreadyDiscarded,
203 TraversalLimit,
205}
206
207pub struct WorkspaceContextRegistry<S> {
209 store: S,
210}
211
212impl<S> WorkspaceContextRegistry<S> {
213 #[must_use]
215 pub const fn new(store: S) -> Self {
216 Self { store }
217 }
218
219 #[must_use]
221 pub const fn store(&self) -> &S {
222 &self.store
223 }
224}
225
226#[derive(Debug, Error)]
228pub enum WorkspaceContextError<E: std::error::Error + 'static> {
229 #[error("workspace-context store failed: {0}")]
231 Store(E),
232 #[error("workspace-context state is incompatible")]
234 IncompatibleState,
235 #[error("workspace context conflicts with its durable registration")]
237 ConflictingRegistration,
238 #[error("workspace context roots do not match their direct parent")]
240 InvalidRoots,
241 #[error("workspace context path is not canonical")]
243 NonCanonicalPath,
244 #[error("only the direct parent may perform this workspace-context operation")]
246 UnauthorizedParent,
247 #[error("workspace-context traversal exceeded its safe bound")]
249 TraversalLimit,
250 #[error("workspace context is discarded")]
252 Discarded,
253 #[error("workspace root is still bound by a live descendant")]
255 RootInUse,
256}
257
258impl<S: WorkspaceContextStore> WorkspaceContextRegistry<S> {
259 pub async fn register_root(
261 &self,
262 context_id: WorkspaceContextId,
263 roots: impl IntoIterator<Item = WorkspaceContextRoot>,
264 ) -> Result<WorkspaceContext, WorkspaceContextError<S::Error>> {
265 self.register(context_id, None, roots).await
266 }
267
268 pub async fn register_child(
272 &self,
273 context_id: WorkspaceContextId,
274 parent_context_id: WorkspaceContextId,
275 roots: impl IntoIterator<Item = WorkspaceContextRoot>,
276 ) -> Result<WorkspaceContext, WorkspaceContextError<S::Error>> {
277 self.register(context_id, Some(parent_context_id), roots)
278 .await
279 }
280
281 pub async fn adopt_root(
289 &self,
290 context_id: WorkspaceContextId,
291 root: WorkspaceContextRoot,
292 ) -> Result<WorkspaceContext, WorkspaceContextError<S::Error>> {
293 collect_roots([root.clone()])?;
294 loop {
295 let mut context = self.resolve(context_id).await?;
296 if context.state == WorkspaceContextState::Discarded {
297 return Err(WorkspaceContextError::Discarded);
298 }
299 if let Some(existing) = context.roots.get(&root.root_id) {
300 return if existing == &root {
301 Ok(context)
302 } else {
303 Err(WorkspaceContextError::ConflictingRegistration)
304 };
305 }
306 if context
307 .roots
308 .values()
309 .any(|existing| existing.source_path == root.source_path)
310 {
311 return Err(WorkspaceContextError::ConflictingRegistration);
312 }
313
314 let mut expected = BTreeMap::from([(context_id, context.revision)]);
315 match context.parent_context_id {
316 None if root.parent_workspace_id.is_some() => {
317 return Err(WorkspaceContextError::InvalidRoots);
318 }
319 None => {}
320 Some(parent_context_id) => {
321 let parent = self.resolve(parent_context_id).await?;
322 let valid_parent = parent.state != WorkspaceContextState::Discarded
323 && parent.roots.get(&root.root_id).is_some_and(|parent_root| {
324 parent_root.source_path == root.source_path
325 && root.parent_workspace_id == Some(parent_root.workspace_id)
326 });
327 if !valid_parent {
328 return Err(WorkspaceContextError::InvalidRoots);
329 }
330 expected.insert(parent_context_id, parent.revision);
331 }
332 }
333
334 context.roots.insert(root.root_id, root.clone());
335 context.revision = context
336 .revision
337 .checked_add(1)
338 .ok_or(WorkspaceContextError::IncompatibleState)?;
339 if self
340 .store
341 .compare_and_swap_many(expected, vec![context.clone()], false)
342 .await
343 .map_err(WorkspaceContextError::Store)?
344 {
345 return Ok(context);
346 }
347 }
348 }
349
350 pub async fn remove_root(
357 &self,
358 context_id: WorkspaceContextId,
359 root_id: WorkspaceRootId,
360 ) -> Result<WorkspaceContext, WorkspaceContextError<S::Error>> {
361 loop {
362 let contexts = self
363 .store
364 .list()
365 .await
366 .map_err(WorkspaceContextError::Store)?;
367 let by_id = contexts
368 .into_iter()
369 .map(|context| (context.context_id, context))
370 .collect::<BTreeMap<_, _>>();
371 let mut context = by_id
372 .get(&context_id)
373 .cloned()
374 .ok_or(WorkspaceContextError::IncompatibleState)?;
375 validate_record(&context, context_id)?;
376 if context.state == WorkspaceContextState::Discarded {
377 return Err(WorkspaceContextError::Discarded);
378 }
379 if !context.roots.contains_key(&root_id) || context.roots.len() == 1 {
380 return Err(WorkspaceContextError::InvalidRoots);
381 }
382 for candidate in by_id.values() {
383 if candidate.context_id == context_id
384 || candidate.state == WorkspaceContextState::Discarded
385 || !candidate.roots.contains_key(&root_id)
386 {
387 continue;
388 }
389 let mut parent = candidate.parent_context_id;
390 let mut visited = BTreeSet::new();
391 while let Some(id) = parent {
392 if id == context_id {
393 return Err(WorkspaceContextError::RootInUse);
394 }
395 if !visited.insert(id) {
396 return Err(WorkspaceContextError::IncompatibleState);
397 }
398 parent = by_id
399 .get(&id)
400 .ok_or(WorkspaceContextError::IncompatibleState)?
401 .parent_context_id;
402 }
403 }
404 context.roots.remove(&root_id);
405 context.revision = context
406 .revision
407 .checked_add(1)
408 .ok_or(WorkspaceContextError::IncompatibleState)?;
409 let expected = by_id
410 .iter()
411 .map(|(id, record)| (*id, record.revision))
412 .collect();
413 if self
414 .store
415 .compare_and_swap_many(expected, vec![context.clone()], true)
416 .await
417 .map_err(WorkspaceContextError::Store)?
418 {
419 return Ok(context);
420 }
421 }
422 }
423
424 pub async fn resolve(
426 &self,
427 context_id: WorkspaceContextId,
428 ) -> Result<WorkspaceContext, WorkspaceContextError<S::Error>> {
429 let context = self
430 .store
431 .load(context_id)
432 .await
433 .map_err(WorkspaceContextError::Store)?
434 .ok_or(WorkspaceContextError::IncompatibleState)?;
435 validate_record(&context, context_id)?;
436 Ok(context)
437 }
438
439 pub async fn authorize_parent(
441 &self,
442 child: WorkspaceContextId,
443 parent: WorkspaceContextId,
444 ) -> Result<WorkspaceContext, WorkspaceContextError<S::Error>> {
445 let context = self.resolve(child).await?;
446 if context.state == WorkspaceContextState::Discarded {
447 return Err(WorkspaceContextError::Discarded);
448 }
449 if context.parent_context_id != Some(parent) {
450 return Err(WorkspaceContextError::UnauthorizedParent);
451 }
452 Ok(context)
453 }
454
455 pub async fn set_mount(
458 &self,
459 context_id: WorkspaceContextId,
460 root_id: WorkspaceRootId,
461 mount_path: Option<PathBuf>,
462 ) -> Result<WorkspaceContext, WorkspaceContextError<S::Error>> {
463 if mount_path
464 .as_deref()
465 .is_some_and(|path| !is_canonical(path))
466 {
467 return Err(WorkspaceContextError::NonCanonicalPath);
468 }
469 self.update(context_id, |context| {
470 if context.state == WorkspaceContextState::Discarded {
471 return Err(WorkspaceContextError::Discarded);
472 }
473 let root = context
474 .roots
475 .get_mut(&root_id)
476 .ok_or(WorkspaceContextError::InvalidRoots)?;
477 root.mount_path = mount_path.clone();
478 Ok(())
479 })
480 .await
481 }
482
483 pub async fn set_workspace(
488 &self,
489 context_id: WorkspaceContextId,
490 root_id: WorkspaceRootId,
491 workspace_id: WorkspaceId,
492 workspace_name: String,
493 parent_workspace_id: Option<WorkspaceId>,
494 ) -> Result<WorkspaceContext, WorkspaceContextError<S::Error>> {
495 if workspace_name.is_empty() {
496 return Err(WorkspaceContextError::InvalidRoots);
497 }
498 loop {
499 let mut context = self.resolve(context_id).await?;
500 if context.state == WorkspaceContextState::Discarded {
501 return Err(WorkspaceContextError::Discarded);
502 }
503 let mut expected = BTreeMap::from([(context_id, context.revision)]);
504 match context.parent_context_id {
505 None if parent_workspace_id.is_some() => {
506 return Err(WorkspaceContextError::InvalidRoots);
507 }
508 None => {}
509 Some(parent_context_id) => {
510 let parent = self.resolve(parent_context_id).await?;
511 if parent.state == WorkspaceContextState::Discarded
512 || parent
513 .roots
514 .get(&root_id)
515 .is_none_or(|root| Some(root.workspace_id) != parent_workspace_id)
516 {
517 return Err(WorkspaceContextError::InvalidRoots);
518 }
519 expected.insert(parent_context_id, parent.revision);
520 }
521 }
522 let root = context
523 .roots
524 .get_mut(&root_id)
525 .ok_or(WorkspaceContextError::InvalidRoots)?;
526 root.workspace_id = workspace_id;
527 root.workspace_name.clone_from(&workspace_name);
528 root.parent_workspace_id = parent_workspace_id;
529 context.revision = context
530 .revision
531 .checked_add(1)
532 .ok_or(WorkspaceContextError::IncompatibleState)?;
533 if self
534 .store
535 .compare_and_swap_many(expected, vec![context.clone()], false)
536 .await
537 .map_err(WorkspaceContextError::Store)?
538 {
539 return Ok(context);
540 }
541 }
542 }
543
544 pub async fn set_active(
546 &self,
547 context_id: WorkspaceContextId,
548 active: bool,
549 ) -> Result<WorkspaceContext, WorkspaceContextError<S::Error>> {
550 self.update(context_id, |context| {
551 if context.state == WorkspaceContextState::Discarded {
552 return Err(WorkspaceContextError::Discarded);
553 }
554 context.state = if active {
555 WorkspaceContextState::Active
556 } else {
557 WorkspaceContextState::Frozen
558 };
559 Ok(())
560 })
561 .await
562 }
563
564 pub async fn discard_subtree(
567 &self,
568 parent: WorkspaceContextId,
569 child: WorkspaceContextId,
570 maximum: u32,
571 ) -> Result<Vec<WorkspaceContextId>, WorkspaceContextError<S::Error>> {
572 match self
573 .store
574 .discard_subtree(parent, child, maximum)
575 .await
576 .map_err(WorkspaceContextError::Store)?
577 {
578 WorkspaceContextDiscardOutcome::Discarded(contexts) => Ok(contexts),
579 WorkspaceContextDiscardOutcome::IncompatibleState => {
580 Err(WorkspaceContextError::IncompatibleState)
581 }
582 WorkspaceContextDiscardOutcome::UnauthorizedParent => {
583 Err(WorkspaceContextError::UnauthorizedParent)
584 }
585 WorkspaceContextDiscardOutcome::AlreadyDiscarded => {
586 Err(WorkspaceContextError::Discarded)
587 }
588 WorkspaceContextDiscardOutcome::TraversalLimit => {
589 Err(WorkspaceContextError::TraversalLimit)
590 }
591 }
592 }
593
594 pub async fn route(
597 &self,
598 context_id: WorkspaceContextId,
599 path: &Path,
600 ) -> Result<Option<WorkspaceRoute>, WorkspaceContextError<S::Error>> {
601 if !is_canonical(path) {
602 return Err(WorkspaceContextError::NonCanonicalPath);
603 }
604 let context = self.resolve(context_id).await?;
605 if context.state == WorkspaceContextState::Discarded {
606 return Err(WorkspaceContextError::Discarded);
607 }
608 let mounted = context
609 .roots
610 .values()
611 .filter_map(|root| root.mount_path.as_ref().map(|prefix| (root, prefix)))
612 .filter(|(_, prefix)| path.starts_with(prefix))
613 .max_by_key(|(_, prefix)| prefix.components().count());
614 let selected = mounted
615 .map(|(root, prefix)| (root, prefix, WorkspaceRouteKind::Mount))
616 .or_else(|| {
617 context
618 .roots
619 .values()
620 .filter(|root| path.starts_with(&root.source_path))
621 .max_by_key(|root| root.source_path.components().count())
622 .map(|root| (root, &root.source_path, WorkspaceRouteKind::Source))
623 });
624 Ok(selected.map(|(root, prefix, kind)| WorkspaceRoute {
625 context_id,
626 root_id: root.root_id,
627 kind,
628 prefix: prefix.clone(),
629 relative_path: path
630 .strip_prefix(prefix)
631 .unwrap_or_else(|_| Path::new(""))
632 .to_path_buf(),
633 }))
634 }
635
636 async fn register(
637 &self,
638 context_id: WorkspaceContextId,
639 parent_context_id: Option<WorkspaceContextId>,
640 roots: impl IntoIterator<Item = WorkspaceContextRoot>,
641 ) -> Result<WorkspaceContext, WorkspaceContextError<S::Error>> {
642 let roots = collect_roots(roots)?;
643 if parent_context_id.is_none()
644 && roots
645 .values()
646 .any(|root| root.parent_workspace_id.is_some())
647 {
648 return Err(WorkspaceContextError::InvalidRoots);
649 }
650 loop {
651 let mut expected = BTreeMap::from([(context_id, 0)]);
652 if let Some(parent_id) = parent_context_id {
653 let parent = self.resolve(parent_id).await?;
654 if parent.state == WorkspaceContextState::Discarded
655 || roots.len() != parent.roots.len()
656 || roots.iter().any(|(root_id, root)| {
657 parent.roots.get(root_id).is_none_or(|parent_root| {
658 root.source_path != parent_root.source_path
659 || root.parent_workspace_id != Some(parent_root.workspace_id)
660 })
661 })
662 {
663 return Err(WorkspaceContextError::InvalidRoots);
664 }
665 expected.insert(parent_id, parent.revision);
666 }
667 let record = WorkspaceContext {
668 version: WORKSPACE_CONTEXT_VERSION,
669 revision: 1,
670 context_id,
671 parent_context_id,
672 roots: roots.clone(),
673 state: WorkspaceContextState::Active,
674 };
675 if self
676 .store
677 .compare_and_swap_many(expected, vec![record.clone()], false)
678 .await
679 .map_err(WorkspaceContextError::Store)?
680 {
681 return Ok(record);
682 }
683 if let Some(existing) = self
684 .store
685 .load(context_id)
686 .await
687 .map_err(WorkspaceContextError::Store)?
688 {
689 validate_record(&existing, context_id)?;
690 return if same_registration(&existing, &record) {
691 Ok(existing)
692 } else {
693 Err(WorkspaceContextError::ConflictingRegistration)
694 };
695 }
696 }
697 }
698
699 async fn update(
700 &self,
701 context_id: WorkspaceContextId,
702 mut action: impl FnMut(&mut WorkspaceContext) -> Result<(), WorkspaceContextError<S::Error>>,
703 ) -> Result<WorkspaceContext, WorkspaceContextError<S::Error>> {
704 loop {
705 let mut context = self.resolve(context_id).await?;
706 let expected = context.revision;
707 action(&mut context)?;
708 context.revision = expected
709 .checked_add(1)
710 .ok_or(WorkspaceContextError::IncompatibleState)?;
711 if self
712 .store
713 .compare_and_swap(context_id, expected, context.clone())
714 .await
715 .map_err(WorkspaceContextError::Store)?
716 {
717 return Ok(context);
718 }
719 }
720 }
721}
722
723fn collect_roots<E: std::error::Error + 'static>(
724 roots: impl IntoIterator<Item = WorkspaceContextRoot>,
725) -> Result<BTreeMap<WorkspaceRootId, WorkspaceContextRoot>, WorkspaceContextError<E>> {
726 let mut collected = BTreeMap::new();
727 for root in roots {
728 if root.workspace_name.is_empty()
729 || !is_canonical(&root.source_path)
730 || root
731 .mount_path
732 .as_deref()
733 .is_some_and(|path| !is_canonical(path))
734 || collected.insert(root.root_id, root).is_some()
735 {
736 return Err(WorkspaceContextError::InvalidRoots);
737 }
738 }
739 if collected.is_empty() {
740 return Err(WorkspaceContextError::InvalidRoots);
741 }
742 Ok(collected)
743}
744
745fn validate_record<E: std::error::Error + 'static>(
746 context: &WorkspaceContext,
747 expected: WorkspaceContextId,
748) -> Result<(), WorkspaceContextError<E>> {
749 if !record_is_valid(context, expected) {
750 return Err(WorkspaceContextError::IncompatibleState);
751 }
752 Ok(())
753}
754
755fn record_is_valid(context: &WorkspaceContext, expected: WorkspaceContextId) -> bool {
756 !(context.version != WORKSPACE_CONTEXT_VERSION
757 || context.context_id != expected
758 || context.revision == 0
759 || context.roots.is_empty()
760 || context.roots.iter().any(|(root_id, root)| {
761 root_id != &root.root_id
762 || root.workspace_name.is_empty()
763 || !is_canonical(&root.source_path)
764 || root
765 .mount_path
766 .as_deref()
767 .is_some_and(|path| !is_canonical(path))
768 }))
769}
770
771fn same_registration(left: &WorkspaceContext, right: &WorkspaceContext) -> bool {
772 left.version == right.version
773 && left.context_id == right.context_id
774 && left.parent_context_id == right.parent_context_id
775 && left.roots == right.roots
776}
777
778fn is_canonical(path: &Path) -> bool {
779 path.is_absolute()
780 && path
781 .components()
782 .all(|component| !matches!(component, Component::CurDir | Component::ParentDir))
783}
784
785pub(crate) type WorkspaceContextChildren =
786 BTreeMap<WorkspaceContextId, BTreeSet<WorkspaceContextId>>;
787
788pub(crate) fn update_context_children(
789 children: &mut WorkspaceContextChildren,
790 before: Option<&WorkspaceContext>,
791 after: Option<&WorkspaceContext>,
792) {
793 if let Some(before) = before
794 && before.parent_context_id != after.and_then(|record| record.parent_context_id)
795 && let Some(parent) = before.parent_context_id
796 && let Some(siblings) = children.get_mut(&parent)
797 {
798 siblings.remove(&before.context_id);
799 if siblings.is_empty() {
800 children.remove(&parent);
801 }
802 }
803 if let Some(after) = after
804 && before.and_then(|record| record.parent_context_id) != after.parent_context_id
805 && let Some(parent) = after.parent_context_id
806 {
807 children.entry(parent).or_default().insert(after.context_id);
808 }
809}
810
811pub(crate) fn context_children(
812 records: impl IntoIterator<Item = WorkspaceContext>,
813) -> WorkspaceContextChildren {
814 let mut children = WorkspaceContextChildren::new();
815 for record in records {
816 update_context_children(&mut children, None, Some(&record));
817 }
818 children
819}
820
821pub(crate) fn plan_context_subtree_discard(
822 root: Option<&WorkspaceContext>,
823 children: &WorkspaceContextChildren,
824 parent: WorkspaceContextId,
825 child: WorkspaceContextId,
826 maximum: u32,
827) -> Result<Vec<WorkspaceContextId>, WorkspaceContextDiscardOutcome> {
828 if maximum == 0 {
829 return Err(WorkspaceContextDiscardOutcome::TraversalLimit);
830 }
831 let Some(root) = root else {
832 return Err(WorkspaceContextDiscardOutcome::IncompatibleState);
833 };
834 if root.state == WorkspaceContextState::Discarded {
835 return Err(WorkspaceContextDiscardOutcome::AlreadyDiscarded);
836 }
837 if root.parent_context_id != Some(parent) {
838 return Err(WorkspaceContextDiscardOutcome::UnauthorizedParent);
839 }
840 let maximum = maximum as usize;
841 let mut pending = vec![child];
842 let mut seen = BTreeSet::new();
843 while let Some(context_id) = pending.pop() {
844 if !seen.insert(context_id) || seen.len() > maximum {
845 return Err(WorkspaceContextDiscardOutcome::TraversalLimit);
846 }
847 if let Some(descendants) = children.get(&context_id) {
848 if seen.len() + pending.len() + descendants.len() > maximum {
849 return Err(WorkspaceContextDiscardOutcome::TraversalLimit);
850 }
851 pending.extend(descendants.iter().copied());
852 }
853 }
854 Ok(seen.into_iter().collect())
855}
856
857pub(crate) fn discard_context_subtree(
858 records: &mut BTreeMap<WorkspaceContextId, WorkspaceContext>,
859 children: &WorkspaceContextChildren,
860 parent: WorkspaceContextId,
861 child: WorkspaceContextId,
862 maximum: u32,
863) -> WorkspaceContextDiscardOutcome {
864 let discarded =
865 match plan_context_subtree_discard(records.get(&child), children, parent, child, maximum) {
866 Ok(discarded) => discarded,
867 Err(outcome) => return outcome,
868 };
869 let mut replacements = Vec::with_capacity(discarded.len());
870 for context_id in &discarded {
871 let Some(mut context) = records.get(context_id).cloned() else {
872 return WorkspaceContextDiscardOutcome::IncompatibleState;
873 };
874 if !record_is_valid(&context, *context_id) {
875 return WorkspaceContextDiscardOutcome::IncompatibleState;
876 }
877 let Some(revision) = context.revision.checked_add(1) else {
878 return WorkspaceContextDiscardOutcome::IncompatibleState;
879 };
880 context.state = WorkspaceContextState::Discarded;
881 for root in context.roots.values_mut() {
882 root.mount_path = None;
883 }
884 context.revision = revision;
885 replacements.push(context);
886 }
887 for replacement in replacements {
888 records.insert(replacement.context_id, replacement);
889 }
890 WorkspaceContextDiscardOutcome::Discarded(discarded)
891}
892
893#[derive(Default)]
895pub struct MemoryWorkspaceContextStore {
896 state: Mutex<MemoryWorkspaceContextState>,
897}
898
899#[derive(Default)]
900struct MemoryWorkspaceContextState {
901 records: BTreeMap<WorkspaceContextId, WorkspaceContext>,
902 children: WorkspaceContextChildren,
903}
904
905impl MemoryWorkspaceContextStore {
906 #[must_use]
908 pub fn new() -> Self {
909 Self::default()
910 }
911}
912
913#[derive(Debug, Error)]
915#[error("workspace-context memory store is unavailable")]
916pub struct MemoryWorkspaceContextStoreError;
917
918impl WorkspaceContextStore for MemoryWorkspaceContextStore {
919 type Error = MemoryWorkspaceContextStoreError;
920
921 async fn load(
922 &self,
923 context_id: WorkspaceContextId,
924 ) -> Result<Option<WorkspaceContext>, Self::Error> {
925 self.state
926 .lock()
927 .map_err(|_| MemoryWorkspaceContextStoreError)
928 .map(|state| state.records.get(&context_id).cloned())
929 }
930
931 async fn list(&self) -> Result<Vec<WorkspaceContext>, Self::Error> {
932 self.state
933 .lock()
934 .map_err(|_| MemoryWorkspaceContextStoreError)
935 .map(|state| state.records.values().cloned().collect())
936 }
937
938 async fn compare_and_swap(
939 &self,
940 context_id: WorkspaceContextId,
941 expected_revision: u64,
942 replacement: WorkspaceContext,
943 ) -> Result<bool, Self::Error> {
944 let mut state = self
945 .state
946 .lock()
947 .map_err(|_| MemoryWorkspaceContextStoreError)?;
948 let revision = state
949 .records
950 .get(&context_id)
951 .map_or(0, |record| record.revision);
952 if revision != expected_revision {
953 return Ok(false);
954 }
955 let before = state.records.get(&context_id).cloned();
956 update_context_children(&mut state.children, before.as_ref(), Some(&replacement));
957 state.records.insert(context_id, replacement);
958 Ok(true)
959 }
960
961 async fn compare_and_swap_many(
962 &self,
963 expected_revisions: BTreeMap<WorkspaceContextId, u64>,
964 replacements: Vec<WorkspaceContext>,
965 require_exact_set: bool,
966 ) -> Result<bool, Self::Error> {
967 let mut state = self
968 .state
969 .lock()
970 .map_err(|_| MemoryWorkspaceContextStoreError)?;
971 if require_exact_set && state.records.len() != expected_revisions.len() {
972 return Ok(false);
973 }
974 if expected_revisions.iter().any(|(context_id, expected)| {
975 state
976 .records
977 .get(context_id)
978 .map_or(0, |record| record.revision)
979 != *expected
980 }) || replacements
981 .iter()
982 .any(|replacement| !expected_revisions.contains_key(&replacement.context_id))
983 {
984 return Ok(false);
985 }
986 for replacement in replacements {
987 let before = state.records.get(&replacement.context_id).cloned();
988 update_context_children(&mut state.children, before.as_ref(), Some(&replacement));
989 state.records.insert(replacement.context_id, replacement);
990 }
991 Ok(true)
992 }
993
994 async fn discard_subtree(
995 &self,
996 parent: WorkspaceContextId,
997 child: WorkspaceContextId,
998 maximum: u32,
999 ) -> Result<WorkspaceContextDiscardOutcome, Self::Error> {
1000 let mut state = self
1001 .state
1002 .lock()
1003 .map_err(|_| MemoryWorkspaceContextStoreError)?;
1004 let MemoryWorkspaceContextState { records, children } = &mut *state;
1005 Ok(discard_context_subtree(
1006 records, children, parent, child, maximum,
1007 ))
1008 }
1009}
1010
1011#[cfg(test)]
1012#[allow(clippy::expect_used)]
1013mod tests {
1014 use super::*;
1015 use std::sync::{
1016 Arc,
1017 atomic::{AtomicUsize, Ordering},
1018 };
1019
1020 #[derive(Default)]
1021 struct CountingContextStore {
1022 inner: MemoryWorkspaceContextStore,
1023 loads: Arc<AtomicUsize>,
1024 lists: Arc<AtomicUsize>,
1025 swaps: Arc<AtomicUsize>,
1026 }
1027
1028 impl CountingContextStore {
1029 fn reset(&self) {
1030 self.loads.store(0, Ordering::Relaxed);
1031 self.lists.store(0, Ordering::Relaxed);
1032 self.swaps.store(0, Ordering::Relaxed);
1033 }
1034
1035 fn counts(&self) -> (usize, usize, usize) {
1036 (
1037 self.loads.load(Ordering::Relaxed),
1038 self.lists.load(Ordering::Relaxed),
1039 self.swaps.load(Ordering::Relaxed),
1040 )
1041 }
1042 }
1043
1044 impl WorkspaceContextStore for CountingContextStore {
1045 type Error = MemoryWorkspaceContextStoreError;
1046
1047 async fn load(
1048 &self,
1049 context_id: WorkspaceContextId,
1050 ) -> Result<Option<WorkspaceContext>, Self::Error> {
1051 self.loads.fetch_add(1, Ordering::Relaxed);
1052 self.inner.load(context_id).await
1053 }
1054
1055 async fn list(&self) -> Result<Vec<WorkspaceContext>, Self::Error> {
1056 self.lists.fetch_add(1, Ordering::Relaxed);
1057 self.inner.list().await
1058 }
1059
1060 async fn compare_and_swap(
1061 &self,
1062 context_id: WorkspaceContextId,
1063 expected_revision: u64,
1064 replacement: WorkspaceContext,
1065 ) -> Result<bool, Self::Error> {
1066 self.swaps.fetch_add(1, Ordering::Relaxed);
1067 self.inner
1068 .compare_and_swap(context_id, expected_revision, replacement)
1069 .await
1070 }
1071
1072 async fn compare_and_swap_many(
1073 &self,
1074 expected_revisions: BTreeMap<WorkspaceContextId, u64>,
1075 replacements: Vec<WorkspaceContext>,
1076 require_exact_set: bool,
1077 ) -> Result<bool, Self::Error> {
1078 self.swaps.fetch_add(1, Ordering::Relaxed);
1079 self.inner
1080 .compare_and_swap_many(expected_revisions, replacements, require_exact_set)
1081 .await
1082 }
1083
1084 async fn discard_subtree(
1085 &self,
1086 parent: WorkspaceContextId,
1087 child: WorkspaceContextId,
1088 maximum: u32,
1089 ) -> Result<WorkspaceContextDiscardOutcome, Self::Error> {
1090 self.swaps.fetch_add(1, Ordering::Relaxed);
1091 self.inner.discard_subtree(parent, child, maximum).await
1092 }
1093 }
1094
1095 fn context(byte: u8) -> WorkspaceContextId {
1096 WorkspaceContextId::from_bytes([byte; 16])
1097 }
1098
1099 fn root_id(byte: u8) -> WorkspaceRootId {
1100 WorkspaceRootId::from_bytes([byte; 16])
1101 }
1102
1103 fn workspace(byte: u8) -> WorkspaceId {
1104 WorkspaceId::from_bytes([byte; 16])
1105 }
1106
1107 fn numbered_context(number: u128) -> WorkspaceContextId {
1108 WorkspaceContextId::from_bytes(number.to_le_bytes())
1109 }
1110
1111 fn numbered_root(number: u128) -> WorkspaceRootId {
1112 WorkspaceRootId::from_bytes(number.to_le_bytes())
1113 }
1114
1115 fn numbered_workspace(number: u128) -> WorkspaceId {
1116 WorkspaceId::from_bytes(number.to_le_bytes())
1117 }
1118
1119 fn scaled_roots(generation: u128, count: u128) -> Vec<WorkspaceContextRoot> {
1120 (0..count)
1121 .map(|number| WorkspaceContextRoot {
1122 root_id: numbered_root(number + 1),
1123 source_path: PathBuf::from(if cfg!(windows) {
1124 format!("C:\\scale-root-{number}")
1125 } else {
1126 format!("/scale-root-{number}")
1127 }),
1128 workspace_id: numbered_workspace(generation * count + number + 1),
1129 workspace_name: format!("scale-{generation}-{number}"),
1130 parent_workspace_id: (generation > 0)
1131 .then(|| numbered_workspace((generation - 1) * count + number + 1)),
1132 mount_path: None,
1133 })
1134 .collect()
1135 }
1136
1137 fn root(root_byte: u8, workspace_byte: u8, parent: Option<u8>) -> WorkspaceContextRoot {
1138 WorkspaceContextRoot {
1139 root_id: root_id(root_byte),
1140 source_path: PathBuf::from(if cfg!(windows) { "C:\\work" } else { "/work" }),
1141 workspace_id: workspace(workspace_byte),
1142 workspace_name: format!("workspace-{workspace_byte}"),
1143 parent_workspace_id: parent.map(workspace),
1144 mount_path: None,
1145 }
1146 }
1147
1148 #[tokio::test]
1149 async fn recursive_contexts_require_exact_direct_parent_roots() {
1150 let registry = WorkspaceContextRegistry::new(MemoryWorkspaceContextStore::new());
1151 registry
1152 .register_root(context(1), [root(9, 1, None)])
1153 .await
1154 .expect("root");
1155 registry
1156 .register_child(context(2), context(1), [root(9, 2, Some(1))])
1157 .await
1158 .expect("child");
1159 registry
1160 .register_child(context(3), context(2), [root(9, 3, Some(2))])
1161 .await
1162 .expect("grandchild");
1163
1164 registry
1165 .authorize_parent(context(3), context(2))
1166 .await
1167 .expect("direct parent");
1168 assert!(matches!(
1169 registry.authorize_parent(context(3), context(1)).await,
1170 Err(WorkspaceContextError::UnauthorizedParent)
1171 ));
1172 }
1173
1174 #[tokio::test]
1175 async fn root_removal_requires_descendants_to_release_the_binding_first() {
1176 let registry = WorkspaceContextRegistry::new(MemoryWorkspaceContextStore::new());
1177 registry
1178 .register_root(context(41), scaled_roots(0, 2))
1179 .await
1180 .expect("two root bindings");
1181 registry
1182 .register_child(context(42), context(41), scaled_roots(1, 2))
1183 .await
1184 .expect("child bindings");
1185 assert!(matches!(
1186 registry.remove_root(context(41), numbered_root(2)).await,
1187 Err(WorkspaceContextError::RootInUse)
1188 ));
1189 let child = registry
1190 .remove_root(context(42), numbered_root(2))
1191 .await
1192 .expect("child releases its root");
1193 assert!(!child.roots.contains_key(&numbered_root(2)));
1194 let parent = registry
1195 .remove_root(context(41), numbered_root(2))
1196 .await
1197 .expect("parent releases unused root");
1198 assert!(!parent.roots.contains_key(&numbered_root(2)));
1199 assert!(matches!(
1200 registry.remove_root(context(41), numbered_root(1)).await,
1201 Err(WorkspaceContextError::InvalidRoots)
1202 ));
1203 }
1204
1205 #[tokio::test]
1206 async fn roots_are_adopted_lazily_through_the_exact_parent() {
1207 let registry = WorkspaceContextRegistry::new(MemoryWorkspaceContextStore::new());
1208 registry
1209 .register_root(context(1), [root(9, 1, None)])
1210 .await
1211 .expect("root");
1212 registry
1213 .register_child(context(2), context(1), [root(9, 2, Some(1))])
1214 .await
1215 .expect("child");
1216
1217 let source_path = PathBuf::from(if cfg!(windows) {
1218 "C:\\second-work"
1219 } else {
1220 "/second-work"
1221 });
1222 let parent_root = WorkspaceContextRoot {
1223 root_id: root_id(10),
1224 source_path: source_path.clone(),
1225 workspace_id: workspace(3),
1226 workspace_name: "workspace-3".to_owned(),
1227 parent_workspace_id: None,
1228 mount_path: None,
1229 };
1230 registry
1231 .adopt_root(context(1), parent_root.clone())
1232 .await
1233 .expect("parent adopts root");
1234 let child_root = WorkspaceContextRoot {
1235 root_id: root_id(10),
1236 source_path,
1237 workspace_id: workspace(4),
1238 workspace_name: "workspace-4".to_owned(),
1239 parent_workspace_id: Some(workspace(3)),
1240 mount_path: None,
1241 };
1242 let adopted = registry
1243 .adopt_root(context(2), child_root.clone())
1244 .await
1245 .expect("child adopts parent root");
1246 assert_eq!(adopted.roots.get(&root_id(10)), Some(&child_root));
1247 assert_eq!(
1248 registry
1249 .adopt_root(context(2), child_root)
1250 .await
1251 .expect("exact retry"),
1252 adopted
1253 );
1254 }
1255
1256 #[tokio::test]
1257 async fn child_cannot_adopt_a_root_outside_its_direct_parent() {
1258 let registry = WorkspaceContextRegistry::new(MemoryWorkspaceContextStore::new());
1259 registry
1260 .register_root(context(1), [root(9, 1, None)])
1261 .await
1262 .expect("root");
1263 registry
1264 .register_child(context(2), context(1), [root(9, 2, Some(1))])
1265 .await
1266 .expect("child");
1267
1268 let unauthorized = WorkspaceContextRoot {
1269 root_id: root_id(10),
1270 source_path: PathBuf::from(if cfg!(windows) {
1271 "C:\\outside"
1272 } else {
1273 "/outside"
1274 }),
1275 workspace_id: workspace(4),
1276 workspace_name: "workspace-4".to_owned(),
1277 parent_workspace_id: Some(workspace(3)),
1278 mount_path: None,
1279 };
1280 assert!(matches!(
1281 registry.adopt_root(context(2), unauthorized).await,
1282 Err(WorkspaceContextError::InvalidRoots)
1283 ));
1284 }
1285
1286 #[tokio::test]
1287 async fn mount_routing_wins_and_discard_is_recursive() {
1288 let registry = WorkspaceContextRegistry::new(MemoryWorkspaceContextStore::new());
1289 registry
1290 .register_root(context(1), [root(9, 1, None)])
1291 .await
1292 .expect("root");
1293 registry
1294 .register_child(context(2), context(1), [root(9, 2, Some(1))])
1295 .await
1296 .expect("child");
1297 registry
1298 .register_child(context(3), context(2), [root(9, 3, Some(2))])
1299 .await
1300 .expect("grandchild");
1301 let mount = PathBuf::from(if cfg!(windows) {
1302 "C:\\mounts\\child"
1303 } else {
1304 "/mounts/child"
1305 });
1306 registry
1307 .set_mount(context(2), root_id(9), Some(mount.clone()))
1308 .await
1309 .expect("set mount");
1310 let route = registry
1311 .route(context(2), &mount.join("src/lib.rs"))
1312 .await
1313 .expect("route")
1314 .expect("mapped");
1315 assert_eq!(route.kind, WorkspaceRouteKind::Mount);
1316 assert_eq!(route.relative_path, PathBuf::from("src/lib.rs"));
1317
1318 let discarded = registry
1319 .discard_subtree(context(1), context(2), 8)
1320 .await
1321 .expect("discard");
1322 assert_eq!(discarded, vec![context(2), context(3)]);
1323 assert_eq!(
1324 registry
1325 .resolve(context(3))
1326 .await
1327 .expect("grandchild")
1328 .state,
1329 WorkspaceContextState::Discarded
1330 );
1331 assert!(matches!(
1332 registry.route(context(2), &mount).await,
1333 Err(WorkspaceContextError::Discarded)
1334 ));
1335 assert!(matches!(
1336 registry
1337 .set_mount(context(2), root_id(9), Some(mount.clone()))
1338 .await,
1339 Err(WorkspaceContextError::Discarded)
1340 ));
1341 assert!(matches!(
1342 registry.authorize_parent(context(2), context(1)).await,
1343 Err(WorkspaceContextError::Discarded)
1344 ));
1345 assert!(matches!(
1346 registry.discard_subtree(context(1), context(2), 8).await,
1347 Err(WorkspaceContextError::Discarded)
1348 ));
1349 }
1350
1351 #[tokio::test]
1352 async fn workspace_rebinding_preserves_direct_parent_lineage() {
1353 let registry = WorkspaceContextRegistry::new(MemoryWorkspaceContextStore::new());
1354 registry
1355 .register_root(context(1), [root(9, 1, None)])
1356 .await
1357 .expect("root");
1358 registry
1359 .register_child(context(2), context(1), [root(9, 2, Some(1))])
1360 .await
1361 .expect("child");
1362
1363 assert!(matches!(
1364 registry
1365 .set_workspace(
1366 context(2),
1367 root_id(9),
1368 workspace(3),
1369 "topic".to_owned(),
1370 Some(workspace(99)),
1371 )
1372 .await,
1373 Err(WorkspaceContextError::InvalidRoots)
1374 ));
1375 assert!(matches!(
1376 registry
1377 .set_workspace(
1378 context(1),
1379 root_id(9),
1380 workspace(4),
1381 "root-topic".to_owned(),
1382 Some(workspace(1)),
1383 )
1384 .await,
1385 Err(WorkspaceContextError::InvalidRoots)
1386 ));
1387 let child = registry.resolve(context(2)).await.expect("child");
1388 assert_eq!(child.roots[&root_id(9)].workspace_id, workspace(2));
1389 }
1390
1391 #[tokio::test]
1392 async fn many_roots_and_deep_lineage_have_constant_store_work_per_step() {
1393 const ROOTS: u128 = 64;
1394 const DEPTH: u128 = 64;
1395
1396 let registry = WorkspaceContextRegistry::new(CountingContextStore::default());
1397 registry
1398 .register_root(numbered_context(1), scaled_roots(0, ROOTS))
1399 .await
1400 .expect("large root context");
1401
1402 for generation in 1..DEPTH {
1403 registry.store().reset();
1404 registry
1405 .register_child(
1406 numbered_context(generation + 1),
1407 numbered_context(generation),
1408 scaled_roots(generation, ROOTS),
1409 )
1410 .await
1411 .expect("deep child context");
1412 assert_eq!(
1413 registry.store().counts(),
1414 (1, 0, 1),
1415 "child registration must inspect only its direct parent"
1416 );
1417 }
1418
1419 registry.store().reset();
1420 let route = registry
1421 .route(
1422 numbered_context(DEPTH),
1423 Path::new(if cfg!(windows) {
1424 "C:\\scale-root-63\\nested\\file"
1425 } else {
1426 "/scale-root-63/nested/file"
1427 }),
1428 )
1429 .await
1430 .expect("route")
1431 .expect("matching root");
1432 assert_eq!(route.root_id, numbered_root(ROOTS));
1433 assert_eq!(registry.store().counts(), (1, 0, 0));
1434 }
1435
1436 #[tokio::test]
1437 async fn ten_thousand_durable_contexts_and_128_active_children_remain_direct() {
1438 const CHILDREN: u128 = 10_000;
1439 const ACTIVE_CHILDREN: u128 = 128;
1440
1441 let registry = WorkspaceContextRegistry::new(CountingContextStore::default());
1442 let parent_context = numbered_context(1);
1443 let root = scaled_roots(0, 1).pop().expect("root binding");
1444 let parent_workspace = root.workspace_id;
1445 registry
1446 .register_root(parent_context, [root.clone()])
1447 .await
1448 .expect("root context");
1449
1450 for number in 0..CHILDREN {
1451 registry.store().reset();
1452 registry
1453 .register_child(
1454 numbered_context(number + 2),
1455 parent_context,
1456 [WorkspaceContextRoot {
1457 workspace_id: numbered_workspace(number + 2),
1458 workspace_name: format!("sibling-{number}"),
1459 parent_workspace_id: Some(parent_workspace),
1460 ..root.clone()
1461 }],
1462 )
1463 .await
1464 .expect("sibling child context");
1465 assert_eq!(
1466 registry.store().counts(),
1467 (1, 0, 1),
1468 "registration must inspect only the direct parent"
1469 );
1470 }
1471
1472 for number in ACTIVE_CHILDREN..CHILDREN {
1473 registry.store().reset();
1474 registry
1475 .set_active(numbered_context(number + 2), false)
1476 .await
1477 .expect("freeze retained sibling");
1478 assert_eq!(registry.store().counts(), (1, 0, 1));
1479 }
1480
1481 registry.store().reset();
1482 let records = registry.store().list().await.expect("durable contexts");
1483 assert_eq!(registry.store().counts(), (0, 1, 0));
1484 assert_eq!(
1485 records.len(),
1486 usize::try_from(CHILDREN + 1).expect("context count fits usize")
1487 );
1488 assert_eq!(
1489 records
1490 .iter()
1491 .filter(|record| {
1492 record.parent_context_id == Some(parent_context)
1493 && record.state == WorkspaceContextState::Active
1494 })
1495 .count(),
1496 usize::try_from(ACTIVE_CHILDREN).expect("active child count fits usize")
1497 );
1498
1499 registry.store().reset();
1500 let route = registry
1501 .route(
1502 numbered_context(ACTIVE_CHILDREN + 1),
1503 &root.source_path.join("nested/file"),
1504 )
1505 .await
1506 .expect("route")
1507 .expect("matching root");
1508 assert_eq!(route.root_id, root.root_id);
1509 assert_eq!(registry.store().counts(), (1, 0, 0));
1510
1511 registry.store().reset();
1512 assert!(matches!(
1513 registry
1514 .discard_subtree(parent_context, numbered_context(2), 0)
1515 .await,
1516 Err(WorkspaceContextError::TraversalLimit)
1517 ));
1518 assert_eq!(
1519 registry.store().counts(),
1520 (0, 0, 1),
1521 "the registry must delegate one atomic bounded discard without listing all contexts"
1522 );
1523 assert_eq!(
1524 registry
1525 .resolve(numbered_context(2))
1526 .await
1527 .expect("bounded rejection changes nothing")
1528 .state,
1529 WorkspaceContextState::Active
1530 );
1531 }
1532}