Back VortexSession with ArcSwap for mutable but less errorprone session management (#8631)

I think immutable session is nice but hard to integrate it with for a
lot of systems

Author: robert3005

Cargo.lock

@@ -12,14 +12,15 @@ or other shared resources.
 
 The session is built on two primitives from the `vortex-session` crate:
 
-- **`VortexSession`** -- a cloneable, thread-safe map from Rust `TypeId` to a boxed value. Any
-  type that is `Send + Sync + Debug + 'static` can be stored as a session variable.
+- **`VortexSession`** -- a cloneable, thread-safe map from Rust `TypeId` to a shared
+  (`Arc`-wrapped) value. Any type that is `Clone + Send + Sync + Debug + 'static` can be stored
+  as a session variable.
 - **`Registry<T>`** -- a concurrent map from string IDs to values of type `T`, used by each
   component to look up registered plugins at runtime.
 
-Because `VortexSession` is backed by an `Arc<DashMap>`, cloning is cheap and all clones share
-the same state. This makes it safe to hand the session to multiple threads, tasks, or I/O
-operations without coordination.
+Because `VortexSession` is backed by an `ArcSwap`, cloning is cheap and all clones share the
+same state, with lock-free reads and copy-on-write writes. This makes it safe to hand the
+session to multiple threads, tasks, or I/O operations without coordination.
 
 ## Component Registries
 
@@ -95,14 +96,13 @@ all built-in components and encodings:
 let session = VortexSession::default();
 ```
 
-For tests or specialized use-cases, sessions can be assembled from individual components using
-the `.with::<T>()` builder:
+For tests or specialized use-cases, sessions can be assembled from individual components by
+starting from an empty session and chaining the `.with::<T>()` helpers:
 
 ```rust
-let session = VortexSession::builder()
+let session = VortexSession::empty()
     .with::<ArraySession>()
     .with::<LayoutSession>()
     .with::<ScalarFnSession>()
-    .with::<RuntimeSession>()
-    .build();
+    .with::<RuntimeSession>();
 ```

docs/developer-guide/internals/session.md

@@ -20,9 +20,8 @@ use crate::RLE;
 use crate::rle::RLEArrayExt;
 
 pub(crate) fn initialize(session: &VortexSession) {
-    session
-        .kernels()
-        .register_execute_parent_kernel(Slice.id(), RLE, SliceExecuteAdaptor(RLE));
+    let kernels = session.kernels();
+    kernels.register_execute_parent_kernel(Slice.id(), RLE, SliceExecuteAdaptor(RLE));
 }
 
 impl SliceKernel for RLE {

encodings/fastlanes/src/rle/kernel.rs

@@ -10,7 +10,6 @@ use vortex_session::VortexSession;
 use crate::ZigZag;
 
 pub(crate) fn initialize(session: &VortexSession) {
-    session
-        .kernels()
-        .register_execute_parent_kernel(Dict.id(), ZigZag, TakeExecuteAdaptor(ZigZag));
+    let kernels = session.kernels();
+    kernels.register_execute_parent_kernel(Dict.id(), ZigZag, TakeExecuteAdaptor(ZigZag));
 }

encodings/zigzag/src/kernel.rs

@@ -56,7 +56,7 @@ mod native_runtime {
         if vortex_cuda::cuda_available() {
             use vortex_cuda::CudaSessionExt;
             session = session.with::<vortex_cuda::CudaSession>();
-            vortex_cuda::initialize_cuda(session.cuda_session());
+            vortex_cuda::initialize_cuda(&session.cuda_session());
         }
         session
     });

fuzz/src/lib.rs

@@ -264,6 +264,7 @@ impl<V: AggregateFnVTable> DynAccumulator for Accumulator<V> {
 mod tests {
     use vortex_buffer::buffer;
     use vortex_error::VortexResult;
+    use vortex_session::SessionExt;
     use vortex_session::VortexSession;
 
     use crate::ArrayRef;

vortex-array/src/aggregate_fn/accumulator.rs

@@ -195,10 +195,9 @@ mod tests {
 
     #[test]
     fn unknown_aggregate_fn_id_allow_unknown() {
-        let session = VortexSession::builder()
+        let session = VortexSession::empty()
             .with::<AggregateFnSession>()
-            .allow_unknown()
-            .build();
+            .allow_unknown();
 
         let proto = pb::AggregateFn {
             id: "vortex.test.foreign_aggregate".to_string(),

vortex-array/src/aggregate_fn/proto.rs

@@ -5,6 +5,7 @@ use std::any::Any;
 use std::sync::Arc;
 
 use vortex_session::SessionExt;
+use vortex_session::SessionGuard;
 use vortex_session::SessionVar;
 
 use crate::aggregate_fn::AggregateFnId;
@@ -221,7 +222,7 @@ impl AggregateFnSession {
 /// Extension trait for accessing aggregate function session data.
 pub trait AggregateFnSessionExt: SessionExt {
     /// Returns the aggregate function session data.
-    fn aggregate_fns(&self) -> &AggregateFnSession {
+    fn aggregate_fns(&self) -> SessionGuard<'_, AggregateFnSession> {
         self.get::<AggregateFnSession>()
     }
 }

vortex-array/src/aggregate_fn/session.rs

@@ -28,9 +28,8 @@ use crate::optimizer::kernels::ArrayKernelsExt;
 use crate::optimizer::rules::ArrayParentReduceRule;
 
 pub(crate) fn initialize(session: &VortexSession) {
-    session
-        .kernels()
-        .register_execute_parent_kernel(Dict.id(), Filter, TakeExecuteAdaptor(Filter));
+    let kernels = session.kernels();
+    kernels.register_execute_parent_kernel(Dict.id(), Filter, TakeExecuteAdaptor(Filter));
 }
 
 pub trait FilterReduce: VTable {

vortex-array/src/arrays/filter/kernel.rs

@@ -119,8 +119,9 @@ pub use varbinview::VarBinViewArray;
 pub mod variant;
 pub use variant::Variant;
 pub use variant::VariantArray;
+use vortex_session::VortexSession;
 
-pub(crate) fn initialize(session: &vortex_session::VortexSession) {
+pub(crate) fn initialize(session: &VortexSession) {
     bool::initialize(session);
     chunked::initialize(session);
     decimal::initialize(session);