Add databend_main and databend_test proc macros

Cargo.toml

@@ -15,9 +15,9 @@ members = [
     "common/functions",
     "common/infallible",
     "common/io",
+    "common/macros",
     "common/management",
-    "common/mem/mem-allocator",
-    "common/mem/mem-derive",
+    "common/mem-allocator",
     "common/planners",
     "common/meta/api",
     "common/meta/embedded",

common/datavalues/Cargo.toml

@@ -8,11 +8,11 @@ edition = "2021"
 
 [dependencies] # In alphabetical order
 # Workspace dependencies
-common-arrow = {path = "../arrow"}
-common-exception = {path = "../exception"}
-common-io = {path = "../io"}
-common-mem-derive = {path = "../mem/mem-derive"}
-common-mem-allocator = {path = "../mem/mem-allocator"}
+common-arrow = { path = "../arrow" }
+common-exception = { path = "../exception" }
+common-io = { path = "../io" }
+common-macros = { path = "../macros" }
+common-mem-allocator = { path = "../mem-allocator" }
 
 # Github dependencies
 

common/datavalues/src/data_field.rs

@@ -15,7 +15,7 @@
 use std::collections::BTreeMap;
 
 use common_arrow::arrow::datatypes::Field as ArrowField;
-use common_mem_derive::*;
+use common_macros::MallocSizeOf;
 
 use crate::DataType;
 

common/datavalues/src/data_value.rs

@@ -25,7 +25,7 @@ use common_arrow::arrow::datatypes::Field as ArrowField;
 use common_exception::ErrorCode;
 use common_exception::Result;
 use common_io::prelude::*;
-use common_mem_derive::*;
+use common_macros::MallocSizeOf;
 
 use crate::arrays::ListBooleanArrayBuilder;
 use crate::arrays::ListBuilderTrait;

common/datavalues/src/types/data_type.rs

@@ -15,7 +15,7 @@
 use core::fmt;
 
 use common_arrow::arrow::datatypes::DataType as ArrowDataType;
-use common_mem_derive::*;
+use common_macros::MallocSizeOf;
 
 use crate::DataField;
 use crate::PhysicalDataType;

common/mem/mem-derive/Cargo.toml → common/macros/Cargo.toml

@@ -1,5 +1,5 @@
 [package]
-name = "common-mem-derive"
+name = "common-macros"
 version = "0.1.0"
 edition = "2021"
 
@@ -9,6 +9,8 @@ edition = "2021"
 proc-macro = true
 
 [dependencies]
+quote = "1"
 proc-macro2 = "1"
 syn = { version = "1", features = ["full"] }
 synstructure = "0.12"
+common-base = { path = "../base" }

common/macros/src/async_entrypoint.rs

@@ -0,0 +1,128 @@
+// Copyright 2020 Datafuse Labs.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use proc_macro::TokenStream;
+use proc_macro2::Ident;
+use syn::spanned::Spanned;
+
+fn token_stream_with_error(mut tokens: TokenStream, error: syn::Error) -> TokenStream {
+    tokens.extend(TokenStream::from(error.into_compile_error()));
+    tokens
+}
+
+pub fn async_main(_args: TokenStream, item: TokenStream) -> TokenStream {
+    let input: syn::ItemFn = match syn::parse(item.clone()) {
+        Ok(it) => it,
+        Err(e) => return token_stream_with_error(item, e),
+    };
+
+    match input.sig.inputs.len() {
+        0 => parse_knobs(input, false, false),
+        1 => parse_knobs(input, false, true),
+        _ => token_stream_with_error(item, syn::Error::new(input.sig.span(), "")),
+    }
+}
+
+pub fn async_test(_args: TokenStream, item: TokenStream) -> TokenStream {
+    let input: syn::ItemFn = match syn::parse(item.clone()) {
+        Ok(it) => it,
+        Err(e) => return token_stream_with_error(item, e),
+    };
+
+    match input.sig.inputs.len() {
+        0 => parse_knobs(input, true, false),
+        1 => parse_knobs(input, true, true),
+        _ => token_stream_with_error(item, syn::Error::new(input.sig.span(), "")),
+    }
+}
+
+fn parse_knobs(mut input: syn::ItemFn, is_test: bool, has_tracker: bool) -> TokenStream {
+    let mut inner_impl = input.clone();
+    inner_impl.sig.ident = Ident::new("main_impl", inner_impl.sig.ident.span());
+
+    input.sig.asyncness = None;
+    input.sig.inputs.clear();
+    let (last_stmt_start_span, last_stmt_end_span) = {
+        let mut last_stmt = input
+            .block
+            .stmts
+            .last()
+            .map(quote::ToTokens::into_token_stream)
+            .unwrap_or_default()
+            .into_iter();
+
+        let start = last_stmt
+            .next()
+            .map_or_else(proc_macro2::Span::call_site, |t| t.span());
+        let end = last_stmt.last().map_or(start, |t| t.span());
+        (start, end)
+    };
+
+    let rt = quote::quote_spanned! {last_stmt_start_span =>
+        common_base::Runtime::with_default_worker_threads().unwrap()
+    };
+
+    let wait_in_future = match has_tracker {
+        true => quote::quote! {
+            let tracker = runtime.get_tracker();
+            main_impl(tracker.clone())
+        },
+        false => quote::quote! { main_impl() },
+    };
+
+    let header = if is_test {
+        quote::quote! {
+            use common_base::Runtime;
+            #[::core::prelude::v1::test]
+        }
+    } else {
+        quote::quote! {}
+    };
+
+    let body = &input.block;
+    let brace_token = input.block.brace_token;
+    let (tail_return, tail_semicolon) = match body.stmts.last() {
+        Some(syn::Stmt::Semi(expr, _)) => (
+            match expr {
+                syn::Expr::Return(_) => quote::quote! { return },
+                _ => quote::quote! {},
+            },
+            quote::quote! {
+                ;
+            },
+        ),
+        _ => (quote::quote! {}, quote::quote! {}),
+    };
+
+    input.block = syn::parse2(quote::quote_spanned! {last_stmt_end_span=>
+        {
+            #inner_impl
+
+            {
+                use common_base::BlockingWait;
+                let runtime = #rt;
+                #tail_return #wait_in_future.wait_in(&runtime, None).unwrap()#tail_semicolon
+            }
+        }
+    })
+    .expect("Parsing failure");
+    input.block.brace_token = brace_token;
+
+    let result = quote::quote! {
+        #header
+        #input
+    };
+
+    result.into()
+}

common/macros/src/lib.rs

@@ -0,0 +1,34 @@
+// Copyright 2020 Datafuse Labs.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#[allow(dead_code)]
+mod async_entrypoint;
+mod malloc_sizeof;
+
+synstructure::decl_derive!([MallocSizeOf, attributes(ignore_malloc_size_of, conditional_malloc_size_of)] => malloc_sizeof::malloc_size_of_derive);
+
+#[allow(unused)]
+use proc_macro::TokenStream;
+
+#[proc_macro_attribute]
+#[cfg(not(test))] // Work around for rust-lang/rust#62127
+pub fn databend_main(args: TokenStream, item: TokenStream) -> TokenStream {
+    async_entrypoint::async_main(args, item)
+}
+
+#[proc_macro_attribute]
+#[cfg(not(test))] // Work around for rust-lang/rust#62127
+pub fn databend_test(args: TokenStream, item: TokenStream) -> TokenStream {
+    async_entrypoint::async_test(args, item)
+}

common/mem/mem-derive/src/lib.rs → common/macros/src/malloc_sizeof.rs

@@ -21,15 +21,10 @@
 //! A crate for deriving the MallocSizeOf trait.
 //! this is only used for databend
 
-extern crate proc_macro2;
-#[macro_use]
-extern crate syn;
-#[macro_use]
-extern crate synstructure;
+use syn::*;
+use synstructure::*;
 
-decl_derive!([MallocSizeOf, attributes(ignore_malloc_size_of, conditional_malloc_size_of)] => malloc_size_of_derive);
-
-fn malloc_size_of_derive(s: synstructure::Structure) -> proc_macro2::TokenStream {
+pub fn malloc_size_of_derive(s: synstructure::Structure) -> proc_macro2::TokenStream {
     let match_body = s.each(|binding| {
         let mut ignore = false;
         let mut conditional = false;

common/mem/mem-allocator/Cargo.toml → common/mem-allocator/Cargo.toml

@@ -6,8 +6,8 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-common-base = {path = "../../base" }
+common-base = {path = "../base" }
 tikv-jemalloc-sys = "0.4.2"
-common-infallible = { path = "../../infallible" }
+common-infallible = { path = "../infallible" }
 parking_lot = "0.11"
-common-mem-derive = { path = "../mem-derive" }
+common-macros = { path = "../macros" }

common/mem/mem-allocator/tests/derive_test.rs → common/mem-allocator/tests/derive_test.rs

@@ -16,10 +16,10 @@ use std::collections::HashMap;
 use std::sync::Arc;
 
 use common_infallible::Mutex;
+use common_macros::MallocSizeOf;
 use common_mem_allocator::malloc_size;
 use common_mem_allocator::MallocSizeOf;
 use common_mem_allocator::MallocSizeOfExt;
-use common_mem_derive::*;
 
 #[test]
 fn derive_vec() {

metasrv/Cargo.toml

@@ -20,17 +20,18 @@ simd = ["common-arrow/simd"]
 common-arrow = {path = "../common/arrow"}
 common-datablocks = {path = "../common/datablocks"}
 common-datavalues = {path = "../common/datavalues"}
-common-exception = {path = "../common/exception"}
-common-functions = {path = "../common/functions"}
-common-flight-rpc = {path = "../common/flight-rpc"}
-common-infallible = {path = "../common/infallible"}
-common-planners = {path = "../common/planners"}
-common-base = {path = "../common/base" }
-common-meta-raft-store = {path = "../common/meta/raft-store"}
-common-meta-sled-store = {path = "../common/meta/sled-store"}
-common-meta-flight = {path = "../common/meta/flight" }
-common-meta-types = {path = "../common/meta/types"}
-common-tracing = {path = "../common/tracing"}
+common-exception = { path = "../common/exception" }
+common-functions = { path = "../common/functions" }
+common-flight-rpc = { path = "../common/flight-rpc" }
+common-infallible = { path = "../common/infallible" }
+common-planners = { path = "../common/planners" }
+common-base = { path = "../common/base" }
+common-meta-raft-store = { path = "../common/meta/raft-store" }
+common-meta-sled-store = { path = "../common/meta/sled-store" }
+common-meta-flight = { path = "../common/meta/flight" }
+common-meta-types = { path = "../common/meta/types" }
+common-tracing = { path = "../common/tracing" }
+common-macros = { path = "../common/macros" }
 
 # Github dependencies