use proc_macro2::TokenStream;
use quote::{format_ident, quote};
use std::{collections::VecDeque, env, fs::File, io::Write, path::Path};
use syn::{
parenthesized,
parse::{Parse, ParseStream},
punctuated::Punctuated,
token::Paren,
Block, ExprLit, FnArg, Ident, Lit, Pat, ReturnType, Token, Type, Visibility,
};
#[derive(Debug)]
pub struct GwasmFn {
vis: Visibility,
fn_token: Token![fn],
ident: Ident,
paren_token: Paren,
args: Punctuated<FnArg, Token![,]>,
ret: ReturnType,
body: Box<Block>,
}
impl Parse for GwasmFn {
fn parse(input: ParseStream) -> syn::Result<Self> {
let content;
Ok(GwasmFn {
vis: input.parse()?,
fn_token: input.parse()?,
ident: input.parse()?,
paren_token: parenthesized!(content in input),
args: content.parse_terminated(FnArg::parse)?,
ret: input.parse()?,
body: input.parse()?,
})
}
}
fn validate_extract_args(input: impl IntoIterator<Item = FnArg>) -> Vec<(Box<Pat>, Box<Type>)> {
let mut args = vec![];
for arg in input {
let (pat, ty) = match arg {
FnArg::Typed(arg) => {
if arg.attrs.len() > 0 {
panic!("attributes around function arguments are unsupported");
}
let pat = arg.pat;
let ty = arg.ty;
(pat, ty)
}
_ => panic!("functions taking 'self' are unsupported"),
};
args.push((pat, ty));
}
args
}
fn validate_extract_return_type(output: &ReturnType) -> Box<Type> {
match output {
ReturnType::Default => panic!("functions returning unit type () are unsupported"),
ReturnType::Type(_, tt) => tt.clone(),
}
}
#[derive(Debug)]
pub struct GwasmAttr {
ident: Ident,
eq_token: Token![=],
value: ExprLit,
}
impl Parse for GwasmAttr {
fn parse(input: ParseStream) -> syn::Result<Self> {
Ok(GwasmAttr {
ident: input.parse()?,
eq_token: input.parse()?,
value: input.parse()?,
})
}
}
#[derive(Debug)]
pub struct GwasmAttrs(Punctuated<GwasmAttr, Token![,]>);
impl Parse for GwasmAttrs {
fn parse(input: ParseStream) -> syn::Result<Self> {
Ok(GwasmAttrs(input.parse_terminated(GwasmAttr::parse)?))
}
}
#[derive(Debug, Default)]
struct GwasmParams {
run_local: Option<bool>,
budget: Option<u64>,
timeout: Option<u64>, subnet: Option<String>,
}
pub(super) fn remote_fn_impl(attrs: GwasmAttrs, f: GwasmFn, preserved: TokenStream) -> TokenStream {
let mut params = GwasmParams::default();
for attr in attrs.0.into_iter() {
let attr_str = attr.ident.to_string();
match attr_str.as_str() {
"run_local" => {
let lit = attr.value.lit;
match lit {
Lit::Str(s) => params
.run_local
.replace(s.value().parse().expect("correct value")),
Lit::Bool(b) => params.run_local.replace(b.value),
x => panic!(
"invalid attribute value '{:#?}': expected string or bool",
x
),
};
}
"budget" => {
let lit = attr.value.lit;
match lit {
Lit::Str(s) => params
.budget
.replace(s.value().parse().expect("correct value")),
Lit::Int(i) => params
.budget
.replace(i.base10_parse().expect("correct value")),
x => panic!("invalid attribute value '{:#?}': expected string or int", x),
};
}
"timeout" => {
let lit = attr.value.lit;
match lit {
Lit::Str(s) => params
.timeout
.replace(s.value().parse().expect("correct value")),
Lit::Int(i) => params
.timeout
.replace(i.base10_parse().expect("correct value")),
x => panic!("invalid attribute value '{:#?}': expected string or int", x),
};
}
"subnet" => {
let lit = attr.value.lit;
match lit {
Lit::Str(s) => params.subnet.replace(s.value()),
x => panic!("invalid attribute value '{:#?}': expected string or int", x),
};
}
x => panic!(
"unexpected attribute '{}': expected 'budget', 'timeout', or 'subnet'",
x
),
}
}
let args = validate_extract_args(f.args.iter().map(|x| x.clone()));
let fn_vis = f.vis;
let fn_ident = f.ident;
let fn_args = f.args;
let return_type = validate_extract_return_type(&f.ret);
let run_local = params.run_local.unwrap_or(false);
let budget = params.budget.unwrap_or(100);
let timeout = params.timeout.unwrap_or(10 * 60);
let subnet = params.subnet.unwrap_or("devnet-alpha.2".to_string());
let mut local_input_args = vec![];
let mut remote_input_args = vec![];
let input_file_names: Vec<_> = (0..args.len()).map(|i| format!("in{}", i)).collect();
for (name, (arg, _)) in input_file_names.iter().zip(args.iter()) {
let ts = quote! {
let input_path = vol.join(#name);
let serialized = serde_json::to_vec(&#arg).context("serializing input data")?;
fs::write(&input_path, serialized).context("writing serialized data to file")?;
};
local_input_args.push(ts);
let ts = quote! {
let input_path = workspace.path().join(#name);
let serialized = serde_json::to_vec(&#arg).context("serializing input data")?;
fs::write(&input_path, serialized).context("writing serialized data to file")?;
};
remote_input_args.push(ts);
}
let mut local_input_paths = vec![];
let mut remote_input_paths = vec![];
let mut remote_input_names = vec![];
for name in input_file_names {
let ts = quote! {
["/workdir/", #name].join(""),
};
local_input_paths.push(ts);
let ts = quote! {
format!("/workdir/{}", #name)
};
remote_input_names.push(ts);
let ts = quote! {
upload(workspace.path().join(#name), format!("/workdir/{}", #name));
};
remote_input_paths.push(ts);
}
let output = {
if run_local {
quote! {
#fn_vis async fn #fn_ident(#fn_args) -> std::result::Result<#return_type, gfaas::Error> {
use gfaas::__private::anyhow::{anyhow, Context};
use gfaas::__private::tokio::task;
use gfaas::__private::tempfile::tempdir;
use gfaas::__private::ya_runtime_wasi;
use gfaas::__private::package::Package;
use gfaas::__private::serde_json;
use std::{fs, env, path::PathBuf};
task::spawn_blocking(move || {
let workspace = tempdir().context("creating temp dir")?;
let exe_path = env::current_exe().context("extracting path to the current exe")?;
let parent = exe_path
.parent()
.ok_or_else(|| anyhow!("path to the current exe without parent: '{}'", exe_path.display()))?;
let module_name = format!("{}", stringify!(#fn_ident));
let wasm = parent.join(format!("{}.wasm", module_name));
let package_path = workspace.path().join("pkg.zip");
let mut package = Package::new();
package.add_module_from_path(wasm).context("adding Wasm module from path")?;
package.write(&package_path).context("saving Yagna zip package to file")?;
ya_runtime_wasi::deploy(workspace.path(), &package_path).context("deploying Yagna package")?;
ya_runtime_wasi::start(workspace.path()).context("executing Yagna start command")?;
let deployment = ya_runtime_wasi::DeployFile::load(workspace.path()).context("loading deployed Yagna package")?;
let vol = deployment
.vols()
.find(|vol| vol.path.starts_with("/workdir"))
.map(|vol| workspace.path().join(&vol.name))
.context("extracting workdir path from Yagna package")?;
let output_file_name = "out".to_owned();
let output_path = vol.join(&output_file_name);
#(#local_input_args)*
ya_runtime_wasi::run(
workspace.path(),
&module_name,
vec![
#(#local_input_paths)*
["/workdir/", &output_file_name].join(""),
],
).context("executing Yagna run command")?;
let output_data = fs::read(output_path).context("reading output data from file")?;
let res = serde_json::from_slice(&output_data).context("deserializing output data")?;
Ok(res)
}).await?
}
}
} else {
quote! {
#fn_vis async fn #fn_ident(#fn_args) -> std::result::Result<#return_type, gfaas::Error> {
use gfaas::__private::anyhow::{self, Context, anyhow};
use gfaas::__private::futures::future::{select, FutureExt};
use gfaas::__private::tempfile::tempdir;
use gfaas::__private::package::Package;
use gfaas::__private::yarapi::{commands, requestor::{self, CommandList, Image::Wasm, Requestor}};
use gfaas::__private::ya_agreement_utils::{constraints, ConstraintKey, Constraints};
use gfaas::__private::serde_json;
use std::{fs, env, path::{Path, PathBuf}, collections::HashMap};
let workspace = tempdir().context("creating temp dir")?;
let exe_path = env::current_exe().context("extracting path to the current exe")?;
let parent = exe_path
.parent()
.ok_or_else(|| anyhow!("path to the current exe without parent: '{}'", exe_path.display()))?;
let module_name = format!("{}", stringify!(#fn_ident));
let wasm = parent.join(format!("{}.wasm", module_name));
let package_path = workspace.path().join("pkg.zip");
let mut package = Package::new();
package.add_module_from_path(wasm).context("adding Wasm module from path")?;
package.write(&package_path).context("saving Yagna zig package to file")?;
let output_path = workspace.path().join("out");
#(#remote_input_args)*
Requestor::new(
"custom",
Wasm((0, 0, 0).into()),
requestor::Package::Archive(package_path)
)
.with_subnet(#subnet)
.with_max_budget_ngnt(#budget)
.with_timeout(std::time::Duration::from_secs(#timeout))
.with_constraints(constraints![
"golem.inf.mem.gib" > 0.5,
"golem.inf.storage.gib" > 1.0,
])
.with_tasks(vec![commands! {
#(#remote_input_paths)*
run(module_name, #(#remote_input_names),*, "/workdir/out");
download("/workdir/out", &output_path);
}].into_iter())
.on_completed(|activity_id, output| {
println!("{} => {:#?}", activity_id, output);
})
.run()
.await?;
let output_data = fs::read(&output_path).context("reading output data from file")?;
let res = serde_json::from_slice(&output_data).context("deserializing output data")?;
Ok(res)
}
}
}
};
let mut inputs = vec![];
let mut input_args = VecDeque::with_capacity(args.len());
for i in 0..args.len() {
let in_ident = format_ident!("in{}", i);
let ts = quote! {
let next_arg = args.pop().unwrap();
let #in_ident = fs::read(next_arg).unwrap();
let #in_ident = serde_json::from_slice(&#in_ident).unwrap();
};
inputs.push(ts);
input_args.push_front(quote!(#in_ident));
}
let args_in_order = input_args.as_slices().0;
let contents = quote! {
#preserved
fn main() {
use std::fs;
use std::env;
let mut args: Vec<_> = env::args().collect();
let out = args.pop().unwrap();
#(#inputs)*
let res = #fn_ident(#(#args_in_order),*);
let serialized = serde_json::to_vec(&res).unwrap();
fs::write(out, &serialized).unwrap();
}
};
let out_dir = env::var("GFAAS_OUT_DIR")
.expect("GFAAS_OUT_DIR should be defined. Did you build the project with gfaas tool?");
let out_path = Path::new(&out_dir)
.join("gfaas_modules")
.join("src")
.join("bin")
.join(format!("{}.rs", fn_ident.to_string()));
let mut out = File::create(out_path).unwrap_or_else(|_| {
panic!(
"generating Wasm src file {}",
[&out_dir, "gfaas.rs"].join("/")
)
});
writeln!(out, "{}", contents).unwrap();
output
}