WIP: Moving vectors to fallible vectors

ahnlich/Cargo.toml

@@ -41,3 +41,4 @@ deadpool = { version = "0.10", features = ["rt_tokio_1"]}
 opentelemetry =  { version = "0.23.0", features = ["trace"] }
 tracing-opentelemetry = "0.24.0"
 log = "0.4"
+fallible_collections = "0.4.9"

ahnlich/ai/Cargo.toml

@@ -35,6 +35,8 @@ serde_json.workspace = true
 termcolor = "1.4.1"
 strum = { version = "0.26", features = ["derive"] }
 log.workspace = true
+fallible_collections.workspace = true
+rayon.workspace = true
 
 [dev-dependencies]
 db = { path = "../db", version = "*" }

ahnlich/ai/src/engine/store.rs

@@ -9,6 +9,7 @@ use ahnlich_types::keyval::StoreKey;
 use ahnlich_types::keyval::StoreName;
 use ahnlich_types::keyval::StoreValue;
 use ahnlich_types::metadata::MetadataValue;
+use fallible_collections::FallibleVec;
 use flurry::HashMap as ConcurrentHashMap;
 use serde::Deserialize;
 use serde::Serialize;
@@ -162,7 +163,7 @@ impl AIStoreHandler {
     ) -> Result<StoreValidateResponse, AIProxyError> {
         let metadata_key = &*AHNLICH_AI_RESERVED_META_KEY;
         let store = self.get(store_name)?;
-        let mut output = Vec::with_capacity(inputs.len());
+        let mut output: Vec<_> = FallibleVec::try_with_capacity(inputs.len())?;
         let mut delete_hashset = StdHashSet::new();
 
         for (store_input, mut store_value) in inputs {

ahnlich/ai/src/error.rs

@@ -2,6 +2,7 @@ use ahnlich_types::{
     ai::{AIStoreInputType, PreprocessAction},
     keyval::StoreName,
 };
+use fallible_collections::TryReserveError;
 use thiserror::Error;
 use tokio::sync::oneshot::error::RecvError;
 
@@ -67,4 +68,12 @@ pub enum AIProxyError {
         index_model_dim: usize,
         query_model_dim: usize,
     },
+    #[error("allocation error {0:?}")]
+    Allocation(TryReserveError),
+}
+
+impl From<TryReserveError> for AIProxyError {
+    fn from(input: TryReserveError) -> Self {
+        Self::Allocation(input)
+    }
 }

ahnlich/ai/src/manager/mod.rs

@@ -6,6 +6,7 @@ use crate::engine::ai::models::Model;
 use crate::error::AIProxyError;
 use ahnlich_types::ai::{AIModel, AIStoreInputType, ImageAction, PreprocessAction, StringAction};
 use ahnlich_types::keyval::{StoreInput, StoreKey};
+use fallible_collections::FallibleVec;
 use std::collections::HashMap;
 use task_manager::Task;
 use task_manager::TaskManager;
@@ -46,7 +47,7 @@ impl ModelThread {
         process_action: PreprocessAction,
     ) -> ModelThreadResponse {
         let model: Model = (&self.model).into();
-        let mut response = Vec::with_capacity(inputs.len());
+        let mut response: Vec<_> = FallibleVec::try_with_capacity(inputs.len())?;
         // move this from for loop into vec of inputs
         for input in inputs {
             let processed_input = self.preprocess_store_input(process_action, input)?;

ahnlich/ai/src/server/task.rs

@@ -7,6 +7,7 @@ use ahnlich_types::keyval::{StoreInput, StoreValue};
 use ahnlich_types::metadata::MetadataValue;
 use ahnlich_types::predicate::{Predicate, PredicateCondition};
 use ahnlich_types::version::VERSION;
+use fallible_collections::vec::TryFromIterator;
 use std::collections::HashSet;
 use std::net::SocketAddr;
 use std::sync::Arc;
@@ -333,23 +334,22 @@ impl AhnlichProtocol for AIProxyTask {
                         {
                             Ok(res) => {
                                 if let ServerResponse::GetSimN(response) = res {
-                                    // conversion to store input here
-                                    let mut output = Vec::new();
-
-                                    // TODO: Can Parallelize
-                                    for (store_key, store_value, sim) in response.into_iter() {
-                                        let temp =
-                                            self.store_handler.store_key_val_to_store_input_val(
-                                                vec![(store_key, store_value)],
-                                            );
-
-                                        if let Some(valid_result) = temp.first().take() {
-                                            let valid_result = valid_result.to_owned();
-                                            output.push((valid_result.0, valid_result.1, sim))
-                                        }
-                                    }
-
-                                    Ok(AIServerResponse::GetSimN(output))
+                                    TryFromIterator::try_from_iterator(
+                                        response.into_iter().flat_map(
+                                            |(store_key, store_value, sim)| {
+                                                // TODO: Can parallelize
+                                                self.store_handler
+                                                    .store_key_val_to_store_input_val(vec![(
+                                                        store_key,
+                                                        store_value,
+                                                    )])
+                                                    .into_iter()
+                                                    .map(move |v| (v.0, v.1, sim))
+                                            },
+                                        ),
+                                    )
+                                    .map_err(|e| AIProxyError::from(e).to_string())
+                                    .map(AIServerResponse::GetSimN)
                                 } else {
                                     Err(AIProxyError::UnexpectedDBResponse(format!("{:?}", res))
                                         .to_string())

ahnlich/client/Cargo.toml

@@ -22,6 +22,7 @@ bincode.workspace = true
 async-trait.workspace = true
 tokio.workspace = true
 deadpool.workspace = true
+fallible_collections.workspace = true
 [dev-dependencies]
 db = { path = "../db", version = "*" }
 ai = { path = "../ai", version = "*" }

ahnlich/client/src/conn/db.rs

@@ -27,7 +27,7 @@ impl Connection for DBConn {
     }
 
     async fn is_conn_valid(&mut self) -> Result<(), AhnlichError> {
-        let mut queries = Self::ServerQuery::with_capacity(1);
+        let mut queries = Self::ServerQuery::with_capacity(1)?;
         queries.push(DBQuery::Ping);
         let response = self.send_query(queries).await?;
         let mut expected_response = ServerResult::with_capacity(1);

ahnlich/client/src/db.rs

@@ -215,7 +215,7 @@ impl DbClient {
         tracing_id: Option<String>,
     ) -> Result<DbPipeline, AhnlichError> {
         Ok(DbPipeline {
-            queries: ServerDBQuery::with_capacity_and_tracing_id(capacity, tracing_id),
+            queries: ServerDBQuery::with_capacity_and_tracing_id(capacity, tracing_id)?,
             conn: self.pool.get().await?,
         })
     }
@@ -420,7 +420,7 @@ impl DbClient {
         tracing_id: Option<String>,
     ) -> Result<ServerResponse, AhnlichError> {
         let mut conn = self.pool.get().await?;
-        let mut queries = ServerDBQuery::with_capacity_and_tracing_id(1, tracing_id);
+        let mut queries = ServerDBQuery::with_capacity_and_tracing_id(1, tracing_id)?;
         queries.push(query);
         let res = conn
             .send_query(queries)

ahnlich/client/src/error.rs

@@ -1,11 +1,17 @@
+use ahnlich_types::bincode::BincodeSerError;
+use fallible_collections::TryReserveError;
 use thiserror::Error;
 
 #[derive(Error, Debug)]
 pub enum AhnlichError {
     #[error("std io error {0}")]
     Standard(#[from] std::io::Error),
-    #[error("bincode serialize error {0}")]
-    BinCode(#[from] bincode::Error),
+    #[error("{0}")]
+    BinCodeSerAndDeser(#[from] BincodeSerError),
+    #[error("allocation error {0:?}")]
+    Allocation(TryReserveError),
+    #[error("bincode deserialize error {0}")]
+    Bincode(#[from] bincode::Error),
     #[error("db error {0}")]
     DbError(String),
     #[error("empty response")]
@@ -27,3 +33,9 @@ impl From<deadpool::managed::BuildError> for AhnlichError {
         Self::PoolError(format!("{input}"))
     }
 }
+
+impl From<TryReserveError> for AhnlichError {
+    fn from(input: TryReserveError) -> Self {
+        Self::Allocation(input)
+    }
+}

ahnlich/db/Cargo.toml

@@ -37,6 +37,7 @@ serde_json.workspace = true
 async-trait.workspace = true
 rayon.workspace = true
 log.workspace = true
+fallible_collections.workspace = true
 
 
 [dev-dependencies]

ahnlich/db/src/engine/store.rs

@@ -14,6 +14,7 @@ use ahnlich_types::predicate::PredicateCondition;
 use ahnlich_types::similarity::Algorithm;
 use ahnlich_types::similarity::NonLinearAlgorithm;
 use ahnlich_types::similarity::Similarity;
+use fallible_collections::FallibleVec;
 use flurry::HashMap as ConcurrentHashMap;
 use serde::Deserialize;
 use serde::Serialize;
@@ -588,7 +589,7 @@ impl Store {
             .collect();
         let pinned = self.id_to_value.pin();
         let (mut inserted, mut updated) = (0, 0);
-        let mut inserted_keys = Vec::new();
+        let mut inserted_keys: Vec<_> = FallibleVec::try_with_capacity(res.len())?;
         for (key, val) in res {
             if pinned.insert(key, val.clone()).is_some() {
                 updated += 1;

ahnlich/db/src/errors.rs

@@ -1,10 +1,11 @@
 use ahnlich_types::keyval::StoreName;
 use ahnlich_types::metadata::MetadataKey;
 use ahnlich_types::similarity::NonLinearAlgorithm;
+use fallible_collections::TryReserveError;
 use thiserror::Error;
 
 /// TODO: Move to shared rust types so library can deserialize it from the TCP response
-#[derive(Error, Debug, Eq, PartialEq, PartialOrd, Ord)]
+#[derive(Error, Debug, Eq, PartialEq)]
 pub enum ServerError {
     #[error("Predicate {0} not found in store, attempt CREATEPREDINDEX with predicate")]
     PredicateNotFound(MetadataKey),
@@ -21,4 +22,12 @@ pub enum ServerError {
     },
     #[error("Could not deserialize query, error is {0}")]
     QueryDeserializeError(String),
+    #[error("allocation error {0:?}")]
+    Allocation(TryReserveError),
+}
+
+impl From<TryReserveError> for ServerError {
+    fn from(input: TryReserveError) -> Self {
+        Self::Allocation(input)
+    }
 }

ahnlich/typegen/src/tracers/query/db.rs

@@ -88,7 +88,8 @@ pub fn trace_db_query_enum() -> Registry {
     let server_query =
         ServerDBQuery::from_queries(&[deletepred_variant.clone(), set_query.clone()]);
     let trace_id = "00-djf9039023r3-1er".to_string();
-    let server_query_with_trace_id = ServerDBQuery::with_capacity_and_tracing_id(2, Some(trace_id));
+    let server_query_with_trace_id = ServerDBQuery::with_capacity_and_tracing_id(2, Some(trace_id))
+        .expect("Could not create server query");
 
     let _ = tracer
         .trace_value(&mut samples, &create_store)

ahnlich/types/Cargo.toml

@@ -10,3 +10,5 @@ ndarray.workspace = true
 serde.workspace = true
 bincode.workspace = true
 once_cell.workspace = true
+fallible_collections.workspace = true
+thiserror.workspace = true

ahnlich/types/src/bincode.rs

@@ -1,6 +1,7 @@
 use crate::version::VERSION;
 use bincode::config::DefaultOptions;
 use bincode::config::Options;
+use fallible_collections::vec::FallibleVec;
 use serde::de::DeserializeOwned;
 use serde::Serialize;
 
@@ -25,17 +26,18 @@ pub trait BinCodeSerAndDeser
 where
     Self: Serialize + DeserializeOwned + Send,
 {
-    fn serialize(&self) -> Result<Vec<u8>, bincode::Error> {
+    fn serialize(&self) -> Result<Vec<u8>, BincodeSerError> {
         let config = DefaultOptions::new()
             .with_fixint_encoding()
             .with_little_endian();
         let serialized_version_data = config.serialize(&*VERSION)?;
         let serialized_data = config.serialize(self)?;
         let data_length = serialized_data.len() as u64;
         // serialization appends the length buffer to be read first
-        let mut buffer = Vec::with_capacity(
+        let mut buffer: Vec<_> = FallibleVec::try_with_capacity(
             MAGIC_BYTES.len() + VERSION_LENGTH + LENGTH_HEADER_SIZE + serialized_data.len(),
-        );
+        )
+        .map_err(BincodeSerError::Allocation)?;
         buffer.extend(MAGIC_BYTES);
         buffer.extend(serialized_version_data);
         buffer.extend(&data_length.to_le_bytes());
@@ -63,3 +65,11 @@ where
 pub trait BinCodeSerAndDeserResponse: BinCodeSerAndDeser {
     fn from_error(err: String) -> Self;
 }
+
+#[derive(thiserror::Error, Debug)]
+pub enum BincodeSerError {
+    #[error("bincode serialize error {0}")]
+    BinCode(#[from] bincode::Error),
+    #[error("allocation error {0:?}")]
+    Allocation(fallible_collections::TryReserveError),
+}

ahnlich/types/src/db/query.rs

@@ -1,3 +1,5 @@
+use fallible_collections::FallibleVec;
+use fallible_collections::TryReserveError;
 use std::collections::HashSet;
 use std::num::NonZeroUsize;
 
@@ -87,17 +89,20 @@ pub struct ServerQuery {
 }
 
 impl ServerQuery {
-    pub fn with_capacity(len: usize) -> Self {
-        Self {
-            queries: Vec::with_capacity(len),
+    pub fn with_capacity(len: usize) -> Result<Self, TryReserveError> {
+        Ok(Self {
+            queries: FallibleVec::try_with_capacity(len)?,
             trace_id: None,
-        }
+        })
     }
-    pub fn with_capacity_and_tracing_id(len: usize, trace_id: Option<String>) -> Self {
-        Self {
-            queries: Vec::with_capacity(len),
+    pub fn with_capacity_and_tracing_id(
+        len: usize,
+        trace_id: Option<String>,
+    ) -> Result<Self, TryReserveError> {
+        Ok(Self {
+            queries: FallibleVec::try_with_capacity(len)?,
             trace_id,
-        }
+        })
     }
 
     pub fn push(&mut self, entry: Query) {

ahnlich/utils/Cargo.toml

@@ -20,3 +20,4 @@ serde_json.workspace = true
 log.workspace = true
 cap = "0.1.2"
 tokio-util.workspace = true
+fallible_collections.workspace = true