osv 0.3.1 → 0.3.2

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 9a6579b30fb8761382666c3fcfa434d9234c3ebf508cad7b3045b9f4ebc2f8b7
-  data.tar.gz: afb793b1df201e876e955ed75904c729aa353bdee3ae374085902efdd4601603
+  metadata.gz: 687f5a3426558ea4d4d342f153dae1f11da93807daddcc6944888d92cda7502b
+  data.tar.gz: e9f766b94e4c7f806e81eec859a0146b6a4758c5a69c759b084f64e01fd7a888
 SHA512:
-  metadata.gz: 4318353ec8b32f026d3a3c8675a26b6a6545faba213bba28ada9b0ab7d978ea8d7dd6a8e202aaefb765040cc0d54990fd94f8263bbf4ecd6a4150ab9f48e7cc6
-  data.tar.gz: d0c590105f070773f82a2e00c0c069439c55055cfe958e7eb3bcf443a83fb1d3dd1f099b93f2835a30797024451666e42f18aeaa361432824d18f77cfaa4748f
+  metadata.gz: bb14e59bca79613064a07cbcfa6c2531f08c68eb131e512b9822c7295bcacfd051de8dcd49100d12c8377d90a438d528dc5b20bcf6a1aa7c4e44f8e2aaae7111
+  data.tar.gz: a3a83f9653ef5bc41c908de14c96b756c719fd2750c8739b4a2165ff91ed82d34fdfde184c853185a8ff516e324d01816c65764f05adf58db6b2159ae66b0c74

data/Cargo.lock

@@ -43,6 +43,12 @@ version = "2.6.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b048fb63fd8b5923fc5aa7b340d8e156aec7ec02f0c78fa8a6ddc2613f6f71de"
 
+[[package]]
+name = "bumpalo"
+version = "3.16.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "79296716171880943b8470b5f8d03aa55eb2e645a4874bdbb28adb49162e012c"
+
 [[package]]
 name = "cexpr"
 version = "0.6.0"
@@ -232,6 +238,7 @@ dependencies = [
 name = "osv"
 version = "0.1.0"
 dependencies = [
+ "bumpalo",
  "csv",
  "kanal",
  "magnus 0.7.1",
@@ -260,18 +267,18 @@ dependencies = [
 
 [[package]]
 name = "rb-sys"
-version = "0.9.103"
+version = "0.9.104"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "91dbe37ab6ac2fba187480fb6544b92445e41e5c6f553bf0c33743f3c450a1df"
+checksum = "e2e26425f064a90404ed5e33fee2137b02a9c6d1c83e19394f4d8a476b9d76a2"
 dependencies = [
  "rb-sys-build",
 ]
 
 [[package]]
 name = "rb-sys-build"
-version = "0.9.103"
+version = "0.9.104"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c4d56a49dcb646b70b758789c0d16c055a386a4f2a3346333abb69850fa860ce"
+checksum = "c9802c9003c5648ee0a067e9aa8960d402d5f764f682f93c1ed49eec72f6d7fc"
 dependencies = [
  "bindgen",
  "lazy_static",

data/ext/osv/Cargo.toml

@@ -7,9 +7,10 @@ edition = "2021"
 crate-type = ["cdylib"]
 
 [dependencies]
-csv = "1.3.1"
+csv = "^1.3"
 magnus = { version = "0.7", features = ["rb-sys"] }
-rb-sys = "0.9"
+rb-sys = "^0.9"
 serde = { version = "1.0", features = ["derive"] }
 serde_magnus = "0.8.1"
 kanal = "0.1.0-pre8"
+bumpalo = { version = "3.16.0", features = ["collections"] }

data/ext/osv/src/csv/builder.rs

@@ -1,4 +1,5 @@
 use super::{
+    header_cache::StringCache,
     parser::RecordParser,
     reader::{ReadImpl, RecordReader},
 };
@@ -89,11 +90,19 @@ impl<'a, T: RecordParser + Send + 'static> RecordReaderBuilder<'a, T> {
         let headers = RecordReader::<T>::get_headers(self.ruby, &mut reader, self.has_headers)?;
         let null_string = self.null_string;
 
+        let static_headers = StringCache::intern_many(&headers).map_err(|e| {
+            Error::new(
+                self.ruby.exception_runtime_error(),
+                format!("Failed to intern headers: {e}"),
+            )
+        })?;
+        let headers_for_cleanup = static_headers.clone();
+
         let (sender, receiver) = kanal::bounded(self.buffer);
         let handle = thread::spawn(move || {
             let mut record = csv::StringRecord::new();
             while let Ok(true) = reader.read_record(&mut record) {
-                let row = T::parse(&headers, &record, &null_string);
+                let row = T::parse(&static_headers, &record, &null_string);
                 if sender.send(row).is_err() {
                     break;
                 }
@@ -104,6 +113,7 @@ impl<'a, T: RecordParser + Send + 'static> RecordReaderBuilder<'a, T> {
 
         Ok(RecordReader {
             reader: ReadImpl::MultiThreaded {
+                headers: headers_for_cleanup,
                 receiver,
                 handle: Some(handle),
             },

data/ext/osv/src/csv/header_cache.rs

@@ -0,0 +1,63 @@
+use std::{
+    collections::HashMap,
+    sync::{atomic::AtomicU32, LazyLock, Mutex},
+};
+
+static STRING_CACHE: LazyLock<Mutex<HashMap<&'static str, AtomicU32>>> =
+    LazyLock::new(|| Mutex::new(HashMap::with_capacity(100)));
+
+pub struct StringCache {}
+
+impl StringCache {
+    #[allow(dead_code)]
+    pub fn intern(string: String) -> Result<&'static str, String> {
+        let mut cache = STRING_CACHE.lock().map_err(|e| e.to_string())?;
+
+        if let Some((&existing, count)) = cache.get_key_value(string.as_str()) {
+            count.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
+            Ok(existing)
+        } else {
+            let leaked = Box::leak(string.into_boxed_str());
+            cache.insert(leaked, AtomicU32::new(1));
+            Ok(leaked)
+        }
+    }
+
+    pub fn intern_many(strings: &[String]) -> Result<Vec<&'static str>, String> {
+        let mut cache = STRING_CACHE.lock().map_err(|e| e.to_string())?;
+        let mut result = Vec::with_capacity(strings.len());
+
+        for string in strings {
+            let static_str: &'static str =
+                if let Some((&existing, count)) = cache.get_key_value(string.as_str()) {
+                    count.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
+                    existing
+                } else {
+                    let leaked = Box::leak(string.clone().into_boxed_str());
+                    cache.insert(leaked, AtomicU32::new(1));
+                    leaked
+                };
+            result.push(static_str);
+        }
+
+        Ok(result)
+    }
+
+    pub fn clear(headers: &[&'static str]) -> Result<(), String> {
+        let cache = STRING_CACHE.lock().map_err(|e| e.to_string())?;
+
+        for header in headers {
+            if let Some(count) = cache.get(header) {
+                let remaining = count.fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
+                if remaining == 0 {
+                    let ptr = *header as *const str as *mut str;
+                    unsafe {
+                        let _ = Box::from_raw(ptr);
+                    }
+                }
+            }
+        }
+
+        Ok(())
+    }
+}

data/ext/osv/src/csv/mod.rs

@@ -1,8 +1,8 @@
 mod builder;
+mod header_cache;
 mod parser;
 mod reader;
 mod record;
 
 pub use builder::RecordReaderBuilder;
 pub use record::CsvRecord;
-

data/ext/osv/src/csv/parser.rs

@@ -2,24 +2,24 @@ use std::collections::HashMap;
 
 pub trait RecordParser {
     type Output;
-    fn parse<'a>(
-        headers: &'a [String],
+    fn parse(
+        headers: &[&'static str],
         record: &csv::StringRecord,
         null_string: &str,
     ) -> Self::Output;
 }
 
-impl RecordParser for HashMap<String, Option<String>> {
+impl RecordParser for HashMap<&'static str, Option<String>> {
     type Output = Self;
-    fn parse<'a>(
-        headers: &'a [String],
+    fn parse(
+        headers: &[&'static str],
         record: &csv::StringRecord,
         null_string: &str,
     ) -> Self::Output {
         let mut map = HashMap::with_capacity(headers.len());
         for (header, field) in headers.iter().zip(record.iter()) {
             map.insert(
-                header.clone(),
+                *header,
                 if field == null_string {
                     None
                 } else {
@@ -33,8 +33,8 @@ impl RecordParser for HashMap<String, Option<String>> {
 
 impl RecordParser for Vec<Option<String>> {
     type Output = Self;
-    fn parse<'a>(
-        _headers: &'a [String],
+    fn parse(
+        _headers: &[&'static str],
         record: &csv::StringRecord,
         null_string: &str,
     ) -> Self::Output {

data/ext/osv/src/csv/reader.rs

@@ -1,4 +1,4 @@
-use super::parser::RecordParser;
+use super::{header_cache::StringCache, parser::RecordParser};
 use magnus::{Error, Ruby};
 use std::{io::Read, thread};
 
@@ -6,14 +6,36 @@ pub struct RecordReader<T: RecordParser> {
     pub(crate) reader: ReadImpl<T>,
 }
 
+impl<T: RecordParser> Drop for RecordReader<T> {
+    fn drop(&mut self) {
+        match &mut self.reader {
+            ReadImpl::MultiThreaded {
+                receiver,
+                handle,
+                headers,
+            } => {
+                receiver.close();
+                if let Some(handle) = handle.take() {
+                    let _ = handle.join();
+                }
+                StringCache::clear(headers).unwrap();
+            }
+            ReadImpl::SingleThreaded { headers, .. } => {
+                StringCache::clear(headers).unwrap();
+            }
+        }
+    }
+}
+
 #[allow(dead_code)]
 pub enum ReadImpl<T: RecordParser> {
     SingleThreaded {
         reader: csv::Reader<Box<dyn Read + Send + 'static>>,
-        headers: Vec<String>,
+        headers: Vec<&'static str>,
         null_string: String,
     },
     MultiThreaded {
+        headers: Vec<&'static str>,
         receiver: kanal::Receiver<T::Output>,
         handle: Option<thread::JoinHandle<()>>,
     },
@@ -48,7 +70,9 @@ impl<T: RecordParser> Iterator for RecordReader<T> {
 
     fn next(&mut self) -> Option<Self::Item> {
         match &mut self.reader {
-            ReadImpl::MultiThreaded { receiver, handle } => match receiver.recv() {
+            ReadImpl::MultiThreaded {
+                receiver, handle, ..
+            } => match receiver.recv() {
                 Ok(record) => Some(record),
                 Err(_) => {
                     if let Some(handle) = handle.take() {

data/ext/osv/src/csv/record.rs

@@ -1,17 +1,23 @@
-use magnus::{IntoValue, Ruby, Value};
+use magnus::{IntoValue, RHash, Ruby, Value};
 use std::collections::HashMap;
 
 #[derive(Debug)]
 pub enum CsvRecord {
     Vec(Vec<Option<String>>),
-    Map(HashMap<String, Option<String>>),
+    Map(HashMap<&'static str, Option<String>>),
 }
 
 impl IntoValue for CsvRecord {
     fn into_value_with(self, handle: &Ruby) -> Value {
         match self {
             CsvRecord::Vec(vec) => vec.into_value_with(handle),
-            CsvRecord::Map(map) => map.into_value_with(handle),
+            CsvRecord::Map(map) => {
+                let hash = RHash::new();
+                for (k, v) in map {
+                    hash.aset(k, v).unwrap();
+                }
+                hash.into_value_with(handle)
+            }
         }
     }
 }

data/ext/osv/src/reader.rs

@@ -1,15 +1,15 @@
-use std::collections::HashMap;
-
 use crate::csv::{CsvRecord, RecordReaderBuilder};
 use crate::utils::*;
 use magnus::value::ReprValue;
 use magnus::{block::Yield, Error, KwArgs, RHash, Ruby, Symbol, Value};
+use std::collections::HashMap;
 
 pub fn parse_csv(
-    ruby: &Ruby,
     rb_self: Value,
     args: &[Value],
 ) -> Result<Yield<Box<dyn Iterator<Item = CsvRecord>>>, Error> {
+    let ruby = unsafe { Ruby::get_unchecked() };
+
     let CsvArgs {
         to_read,
         has_headers,
@@ -18,7 +18,7 @@ pub fn parse_csv(
         null_string,
         buffer_size,
         result_type,
-    } = parse_csv_args(ruby, args)?;
+    } = parse_csv_args(&ruby, args)?;
 
     if !ruby.block_given() {
         return create_enumerator(EnumeratorArgs {
@@ -35,7 +35,7 @@ pub fn parse_csv(
 
     let iter: Box<dyn Iterator<Item = CsvRecord>> = match result_type.as_str() {
         "hash" => Box::new(
-            RecordReaderBuilder::<HashMap<String, Option<String>>>::new(ruby, to_read)
+            RecordReaderBuilder::<HashMap<&'static str, Option<String>>>::new(&ruby, to_read)
                 .has_headers(has_headers)
                 .delimiter(delimiter)
                 .quote_char(quote_char)
@@ -45,7 +45,7 @@ pub fn parse_csv(
                 .map(CsvRecord::Map),
         ),
         "array" => Box::new(
-            RecordReaderBuilder::<Vec<Option<String>>>::new(ruby, to_read)
+            RecordReaderBuilder::<Vec<Option<String>>>::new(&ruby, to_read)
                 .has_headers(has_headers)
                 .delimiter(delimiter)
                 .quote_char(quote_char)

data/lib/osv/version.rb

@@ -1,3 +1,3 @@
 module OSV
-  VERSION = "0.3.1"
+  VERSION = "0.3.2"
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: osv
 version: !ruby/object:Gem::Version
-  version: 0.3.1
+  version: 0.3.2
 platform: ruby
 authors:
 - Nathan Jaremko
@@ -56,6 +56,7 @@ files:
 - ext/osv/Cargo.toml
 - ext/osv/extconf.rb
 - ext/osv/src/csv/builder.rs
+- ext/osv/src/csv/header_cache.rs
 - ext/osv/src/csv/mod.rs
 - ext/osv/src/csv/parser.rs
 - ext/osv/src/csv/reader.rs