cleanups

Author: tonyalaribe

Cargo.lock

@@ -68,6 +68,7 @@ ahash = "0.8"
 lru = "0.16.1"
 serde_bytes = "0.11.19"
 dashmap = "6.1"
+parking_lot = "0.12"
 envy = "0.4"
 tdigests = "1.0"
 bincode = { version = "2.0", features = ["serde"] }

Cargo.toml

@@ -138,10 +138,14 @@ impl BufferedWriteLayer {
                 return Ok(estimated_size);
             }
 
-            // Exponential backoff: spin_loop for first few attempts, then yield
+            // Exponential backoff: spin_loop for first few attempts, then brief sleep.
+            // Note: Using std::thread::sleep in this sync function called from async context.
+            // This is acceptable because: (1) max sleep is ~1ms, (2) only under high contention,
+            // (3) converting to async would require spawn_blocking which adds more overhead.
             if attempt < 5 {
                 std::hint::spin_loop();
             } else {
+                // Max backoff = 1μs << 10 = 1024μs ≈ 1ms
                 let backoff_micros = CAS_BACKOFF_BASE_MICROS << attempt.min(CAS_BACKOFF_MAX_EXPONENT);
                 std::thread::sleep(std::time::Duration::from_micros(backoff_micros));
             }

src/buffered_write_layer.rs

@@ -36,16 +36,26 @@ use deltalake::operations::create::CreateBuilder;
 use deltalake::{DeltaTable, DeltaTableBuilder};
 use futures::StreamExt;
 use instrumented_object_store::instrument_object_store;
+use std::sync::Mutex;
 use serde::{Deserialize, Serialize};
 use sqlx::{PgPool, postgres::PgPoolOptions};
 use std::fmt;
+use std::sync::OnceLock;
 use std::{any::Any, collections::HashMap, sync::Arc};
 use tokio::sync::RwLock;
 use tokio_util::sync::CancellationToken;
 use tracing::field::Empty;
 use tracing::{Instrument, debug, error, info, instrument, warn};
 use url::Url;
 
+/// Mutex to serialize access to environment variable modifications.
+/// Required because delta-rs uses std::env::var() for AWS credential resolution,
+/// and std::env::set_var is unsafe in multi-threaded contexts.
+static ENV_MUTEX: OnceLock<Mutex<()>> = OnceLock::new();
+fn env_mutex() -> &'static Mutex<()> {
+    ENV_MUTEX.get_or_init(|| Mutex::new(()))
+}
+
 // Changed to support multiple tables per project: (project_id, table_name) -> DeltaTable
 pub type ProjectConfigs = Arc<RwLock<HashMap<(String, String), Arc<RwLock<DeltaTable>>>>>;
 
@@ -1136,17 +1146,18 @@ impl Database {
     async fn create_or_load_delta_table(
         &self, storage_uri: &str, storage_options: HashMap<String, String>, cached_store: Arc<dyn object_store::ObjectStore>,
     ) -> Result<DeltaTable> {
-        // SAFETY: delta-rs internally uses std::env::var() for AWS credential resolution.
-        // While set_var is unsafe in multi-threaded contexts (potential data races with concurrent
-        // env reads), this is acceptable here because:
-        // 1. We only set AWS_* vars which are read by the AWS SDK during client initialization
-        // 2. The values are consistent across calls (same credentials for same storage_options)
-        // 3. Delta table creation happens early in request processing, before parallel query execution
-        // 4. The alternative (forking processes or thread-local storage) adds significant complexity
-        for (key, value) in &storage_options {
-            if key.starts_with("AWS_") {
-                unsafe {
-                    std::env::set_var(key, value);
+        // delta-rs uses std::env::var() for AWS credential resolution.
+        // We serialize access with ENV_MUTEX to prevent data races from concurrent set_var calls.
+        {
+            let _guard = env_mutex().lock();
+            for (key, value) in &storage_options {
+                if key.starts_with("AWS_") {
+                    // SAFETY: Protected by ENV_MUTEX. set_var is only unsafe due to potential
+                    // concurrent reads, which we prevent by holding the mutex during the entire
+                    // block. The mutex ensures only one thread modifies env vars at a time.
+                    unsafe {
+                        std::env::set_var(key, value);
+                    }
                 }
             }
         }
@@ -1194,9 +1205,8 @@ impl Database {
 
         // Fallback to legacy batch queue if configured
         let enable_queue = self.config.core.enable_batch_queue;
-        if !skip_queue && enable_queue && self.batch_queue.is_some() {
+        if !skip_queue && enable_queue && let Some(ref queue) = self.batch_queue {
             span.record("use_queue", true);
-            let queue = self.batch_queue.as_ref().unwrap();
             for batch in batches {
                 if let Err(e) = queue.queue(batch) {
                     return Err(anyhow::anyhow!("Queue error: {}", e));
@@ -1724,12 +1734,19 @@ impl ProjectRoutingTable {
         // delta table provider expects indices based on its own schema.
         let delta_schema = provider.schema();
         let translated_projection = projection.map(|proj| {
-            proj.iter()
-                .filter_map(|&idx| {
-                    let col_name = self.schema.field(idx).name();
-                    delta_schema.fields().iter().position(|f| f.name() == col_name)
-                })
-                .collect::<Vec<_>>()
+            let mut translated = Vec::with_capacity(proj.len());
+            for &idx in proj {
+                let col_name = self.schema.field(idx).name();
+                if let Some(delta_idx) = delta_schema.fields().iter().position(|f| f.name() == col_name) {
+                    translated.push(delta_idx);
+                } else {
+                    warn!(
+                        "Column '{}' requested in projection but not found in Delta schema for table '{}'",
+                        col_name, self.table_name
+                    );
+                }
+            }
+            translated
         });
 
         let delta_plan = provider.scan(state, translated_projection.as_ref(), filters, limit).await?;

src/database.rs

@@ -12,7 +12,7 @@ use datafusion::sql::sqlparser::dialect::GenericDialect;
 use datafusion::sql::sqlparser::parser::Parser as SqlParser;
 use std::sync::atomic::{AtomicI64, AtomicUsize, Ordering};
 use std::sync::{Arc, RwLock};
-use tracing::{debug, instrument, warn};
+use tracing::{debug, info, instrument, warn};
 
 // 10-minute buckets balance flush granularity vs overhead. Shorter = more flushes,
 // longer = larger Delta files. Matches default flush interval for aligned boundaries.
@@ -46,7 +46,7 @@ fn schemas_compatible(existing: &SchemaRef, incoming: &SchemaRef) -> bool {
         }
     }
     if new_fields > 0 {
-        debug!("Schema evolution: {} new nullable field(s) added", new_fields);
+        info!("Schema evolution: {} new nullable field(s) added", new_fields);
     }
     true
 }

src/mem_buffer.rs

@@ -71,6 +71,8 @@ pub struct LoggingSimpleQueryHandler {
 }
 
 impl LoggingSimpleQueryHandler {
+    /// Create a new LoggingSimpleQueryHandler.
+    /// Note: auth_manager is unused since datafusion-postgres 0.14.0 moved auth to server level.
     pub fn new(session_context: Arc<SessionContext>, _auth_manager: Arc<AuthManager>) -> Self {
         Self {
             inner: DfSessionService::new(session_context),
@@ -144,6 +146,8 @@ pub struct LoggingExtendedQueryHandler {
 }
 
 impl LoggingExtendedQueryHandler {
+    /// Create a new LoggingExtendedQueryHandler.
+    /// Note: auth_manager is unused since datafusion-postgres 0.14.0 moved auth to server level.
     pub fn new(session_context: Arc<SessionContext>, _auth_manager: Arc<AuthManager>) -> Self {
         Self {
             inner: DfSessionService::new(session_context),

src/pgwire_handlers.rs

@@ -118,7 +118,7 @@ impl CompactColumn {
         Self {
             null_bitmap: data.nulls().map(|n| n.buffer().as_slice().to_vec()),
             buffers: data.buffers().iter().map(|b| b.as_slice().to_vec()).collect(),
-            children: data.child_data().iter().map(|c| Self::from_array_data(c)).collect(),
+            children: data.child_data().iter().map(Self::from_array_data).collect(),
             null_count: data.null_count(),
             child_lens: data.child_data().iter().map(|c| c.len()).collect(),
         }
@@ -128,7 +128,7 @@ impl CompactColumn {
         Self {
             null_bitmap: data.nulls().map(|n| n.buffer().as_slice().to_vec()),
             buffers: data.buffers().iter().map(|b| b.as_slice().to_vec()).collect(),
-            children: data.child_data().iter().map(|c| Self::from_array_data(c)).collect(),
+            children: data.child_data().iter().map(Self::from_array_data).collect(),
             null_count: data.null_count(),
             child_lens: data.child_data().iter().map(|c| c.len()).collect(),
         }
@@ -454,9 +454,11 @@ fn deserialize_wal_entry(data: &[u8]) -> Result<WalEntry, WalError> {
     }
 
     if data[0..4] == WAL_MAGIC {
-        // v1+ format: data[4] is version byte (>= 1), data[5] is operation
-        // v0 format: data[4] is operation (0-2), no version byte
-        // Distinguish: if data[4] > 2, it must be a version byte
+        // WAL format detection based on byte 4:
+        // - v0 (legacy): data[4] is operation byte (0=Insert, 1=Delete, 2=Update)
+        // - v1+ (current): data[4] is version byte (>=128), data[5] is operation
+        // Since WalOperation values are 0-2 and WAL_VERSION is 128, we can safely
+        // distinguish formats: if data[4] > 2, it must be a version byte, not an operation.
         if data[4] > 2 {
             if data.len() < 6 {
                 return Err(WalError::TooShort { len: data.len() });