Add SystemInfo for WT and indecies

Cargo.toml

@@ -25,12 +25,12 @@ lockfree = { version = "0.5.1" }
 worktable_codegen = { path = "codegen", version = "0.5.1" }
 futures = "0.3.30"
 uuid = { version = "1.10.0", features = ["v4"] }
-#data_bucket = "0.2.1"
-data_bucket = { git = "https://github.com/pathscale/DataBucket", branch = "main" }
-# data_bucket = { path = "../DataBucket", version = "0.2.1" }
+data_bucket = { git = "https://github.com/pathscale/DataBucket", rev = "3eb4fc2" }
+#data_bucket = { path = "../DataBucket", version = "0.2.1" }
 performance_measurement_codegen = { path = "performance_measurement/codegen", version = "0.1.0", optional = true }
 performance_measurement = { path = "performance_measurement", version = "0.1.0", optional = true }
-indexset = { version = "0.11.2", features = ["concurrent", "cdc", "multimap"] }
+indexset = { version = "0.11.3", features = ["concurrent", "cdc", "multimap"] }
 convert_case = "0.6.0"
 ordered-float = "5.0.0"
 serde = { version = "1.0.215", features = ["derive"] }
+prettytable-rs = "^0.10"

codegen/src/lib.rs

@@ -1,3 +1,4 @@
+mod mem_stat;
 mod name_generator;
 mod persist_index;
 mod persist_table;
@@ -26,3 +27,10 @@ pub fn persist_table(input: TokenStream) -> TokenStream {
         .unwrap_or_else(|e| e.to_compile_error())
         .into()
 }
+
+#[proc_macro_derive(MemStat)]
+pub fn mem_stat(input: TokenStream) -> TokenStream {
+    mem_stat::expand(input.into())
+        .unwrap_or_else(|e| e.to_compile_error())
+        .into()
+}

codegen/src/mem_stat/mod.rs

@@ -0,0 +1,148 @@
+use proc_macro2::TokenStream;
+use quote::quote;
+use syn::{Data, DeriveInput, Fields, Result, Type};
+
+fn gen_heap_size_body(data: &Data) -> Result<TokenStream> {
+    gen_mem_fn_body(
+        data,
+        quote! { heap_size() },
+        quote! { std::mem::size_of::<Self>() },
+    )
+}
+
+fn gen_used_size_body(data: &Data) -> Result<TokenStream> {
+    gen_mem_fn_body(
+        data,
+        quote! { used_size() },
+        quote! { std::mem::size_of::<Self>() },
+    )
+}
+
+fn gen_mem_fn_body(
+    data: &Data,
+    method: TokenStream,
+    default_for_copy: TokenStream,
+) -> Result<TokenStream> {
+    match data {
+        Data::Struct(data_struct) => {
+            let fields = match &data_struct.fields {
+                Fields::Named(named) => named.named.iter().collect::<Vec<_>>(),
+                Fields::Unnamed(unnamed) => unnamed.unnamed.iter().collect::<Vec<_>>(),
+                Fields::Unit => vec![],
+            };
+
+            if fields.is_empty() {
+                Ok(quote! { 0 })
+            } else if fields.iter().all(|f| is_copy_primitive(&f.ty)) {
+                Ok(default_for_copy)
+            } else {
+                let field_sizes = fields.iter().enumerate().map(|(i, f)| {
+                    let accessor = match &f.ident {
+                        Some(ident) => quote! { self.#ident },
+                        None => {
+                            let index = syn::Index::from(i);
+                            quote! { self.#index }
+                        }
+                    };
+                    quote! { size += #accessor.#method; }
+                });
+
+                Ok(quote! {
+                    let mut size = 0;
+                    #(#field_sizes)*
+                    size
+                })
+            }
+        }
+
+        Data::Enum(enum_data) => {
+            let arms = enum_data.variants.iter().map(|variant| {
+                let name = &variant.ident;
+                match &variant.fields {
+                    Fields::Unit => {
+                        quote! {
+                            Self::#name => 0,
+                        }
+                    }
+                    Fields::Unnamed(fields) => {
+                        let bindings: Vec<_> = (0..fields.unnamed.len())
+                            .map(|i| syn::Ident::new(&format!("f{}", i), variant.ident.span()))
+                            .collect();
+
+                        let calls = bindings
+                            .iter()
+                            .map(|b| quote! { #b.#method })
+                            .collect::<Vec<_>>();
+                        quote! {
+                            Self::#name(#(#bindings),*) => {
+                                0 #(+ #calls)*
+                            },
+                        }
+                    }
+                    Fields::Named(fields) => {
+                        let bindings: Vec<_> = fields
+                            .named
+                            .iter()
+                            .map(|f| f.ident.as_ref().unwrap())
+                            .collect();
+
+                        let calls = bindings
+                            .iter()
+                            .map(|b| quote! { #b.#method })
+                            .collect::<Vec<_>>();
+                        quote! {
+                            Self::#name { #(#bindings),* } => {
+                                0 #(+ #calls)*
+                            },
+                        }
+                    }
+                }
+            });
+
+            Ok(quote! {
+                match self {
+                    #(#arms)*
+                }
+            })
+        }
+
+        _ => Err(syn::Error::new_spanned(
+            method,
+            "#[derive(MemStat)] only supports structs and enums",
+        )),
+    }
+}
+
+pub fn expand(input: proc_macro2::TokenStream) -> Result<TokenStream> {
+    let input: DeriveInput = syn::parse2(input)?;
+    let name = &input.ident;
+
+    let heap = gen_heap_size_body(&input.data)?;
+    let used = gen_used_size_body(&input.data)?;
+
+    Ok(quote! {
+        impl MemStat for #name {
+            fn heap_size(&self) -> usize {
+                #heap
+            }
+            fn used_size(&self) -> usize {
+                #used
+            }
+        }
+    })
+}
+
+fn is_copy_primitive(ty: &Type) -> bool {
+    matches!(
+        ty,
+        Type::Path(type_path)
+            if type_path.qself.is_none() &&
+               type_path.path.segments.len() == 1 &&
+               matches!(
+                   type_path.path.segments[0].ident.to_string().as_str(),
+                   "u8" | "u16" | "u32" | "u64" | "usize" |
+                   "i8" | "i16" | "i32" | "i64" | "isize" |
+                   "bool" | "char" | "f64" | "f32"
+               )
+    )
+}

codegen/src/worktable/generator/index.rs

@@ -25,7 +25,7 @@ impl Generator {
 
     /// Generates table's secondary index struct definition. It has fields with index names and types varying on index
     /// uniqueness. For unique index it's `TreeIndex<T, Link`, for non-unique `TreeIndex<T, Arc<LockFreeSet<Link>>>`.
-    /// Index also derives `PersistIndex` macro.
+    /// Index also derives `PersistIndex` and `MemStat` macro.
     fn gen_type_def(&mut self) -> TokenStream {
         let name_generator = WorktableNameGenerator::from_table_name(self.name.to_string());
         let ident = name_generator.get_index_type_ident();
@@ -47,11 +47,11 @@ impl Generator {
 
         let derive = if self.is_persist {
             quote! {
-                #[derive(Debug, Default, PersistIndex)]
+                #[derive(Debug, MemStat, Default, PersistIndex)]
             }
         } else {
             quote! {
-                #[derive(Debug, Default)]
+                #[derive(Debug, MemStat, Default)]
             }
         };
 
@@ -73,12 +73,14 @@ impl Generator {
         let save_row_fn = self.gen_save_row_index_fn();
         let delete_row_fn = self.gen_delete_row_index_fn();
         let process_difference_fn = self.gen_process_difference_index_fn();
+        let info_fn = self.gen_index_info_fn();
 
         quote! {
             impl TableSecondaryIndex<#row_type_ident, #avt_type_ident> for #index_type_ident {
                 #save_row_fn
                 #delete_row_fn
                 #process_difference_fn
+                #info_fn
             }
         }
     }
@@ -357,6 +359,50 @@ impl Generator {
             }
         }
     }
+
+    fn gen_index_info_fn(&self) -> TokenStream {
+        let rows = self.columns.indexes.values().map(|idx| {
+            let index_field_name = &idx.name;
+            let index_name_str = index_field_name.to_string();
+
+            if idx.is_unique {
+                quote! {
+
+                    info.push(IndexInfo {
+                        name: #index_name_str.to_string(),
+                        index_type: IndexKind::Unique,
+                        key_count: self.#index_field_name.len(),
+                        capacity: self.#index_field_name.capacity(),
+                        heap_size: self.#index_field_name.heap_size(),
+                        used_size: self.#index_field_name.used_size(),
+                        node_count: self.#index_field_name.node_count(),
+
+
+                    });
+                }
+            } else {
+                quote! {
+                    info.push(IndexInfo {
+                        name: #index_name_str.to_string(),
+                        index_type: IndexKind::NonUnique,
+                        key_count: self.#index_field_name.len(),
+                        capacity: self.#index_field_name.capacity(),
+                        heap_size: self.#index_field_name.heap_size(),
+                        used_size: self.#index_field_name.used_size(),
+                        node_count: self.#index_field_name.node_count(),
+                    });
+                }
+            }
+        });
+
+        quote! {
+            fn index_info(&self) -> Vec<IndexInfo> {
+                let mut info = Vec::new();
+                #(#rows)*
+                info
+            }
+        }
+    }
 }
 
 // TODO: tests...

codegen/src/worktable/generator/row.rs

@@ -96,7 +96,7 @@ impl Generator {
             .collect();
 
         quote! {
-            #[derive(rkyv::Archive, Debug, rkyv::Deserialize, Clone, rkyv::Serialize, PartialEq)]
+            #[derive(rkyv::Archive, Debug, rkyv::Deserialize, Clone, rkyv::Serialize, PartialEq, MemStat)]
             #[rkyv(derive(Debug))]
             #[repr(C)]
             pub struct #ident {

codegen/src/worktable/generator/table/impls.rs

@@ -19,6 +19,7 @@ impl Generator {
         let iter_with_fn = self.gen_table_iter_with_fn();
         let iter_with_async_fn = self.gen_table_iter_with_async_fn();
         let count_fn = self.gen_table_count_fn();
+        let system_info_fn = self.gen_system_info_fn();
 
         quote! {
             impl #ident {
@@ -31,6 +32,7 @@ impl Generator {
                 #get_next_fn
                 #iter_with_fn
                 #iter_with_async_fn
+                #system_info_fn
             }
         }
     }
@@ -222,4 +224,12 @@ impl Generator {
             }
         }
     }
+
+    fn gen_system_info_fn(&self) -> TokenStream {
+        quote! {
+            pub fn system_info(&self) -> SystemInfo {
+                self.0.system_info()
+            }
+        }
+    }
 }

examples/Cargo.toml

@@ -13,4 +13,7 @@ eyre = "0.6.12"
 futures = "0.3.30"
 async-std = "1.10"
 either = "1.15.0"
+ordered-float = "5.0.0"
+indexset = { version = "0.11.3", features = ["concurrent", "cdc", "multimap"] }
+tokio = { version = "1", features = ["full"] }