Skip to content

Commit 41aae83

Browse files
khvzakkhvzak
committed
Optimize callback creation
Attach only one upvalue to callbacks rather than two. This leads to less lookup to Lua registry.
1 parent fc84e86 commit 41aae83

File tree

4 files changed

+106
-83
lines changed

4 files changed

+106
-83
lines changed

‎benches/benchmark.rs

+17
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,22 @@ fn create_string_table(c: &mut Criterion) {
5959
});
6060
}
6161

62+
fn create_function(c: &mut Criterion) {
63+
let lua = Lua::new();
64+
65+
c.bench_function("create [function] 10", |b| {
66+
b.iter_batched(
67+
|| collect_gc_twice(&lua),
68+
|_| {
69+
for i in 0..10 {
70+
lua.create_function(move |_, ()| Ok(i)).unwrap();
71+
}
72+
},
73+
BatchSize::SmallInput,
74+
);
75+
});
76+
}
77+
6278
fn call_lua_function(c: &mut Criterion) {
6379
let lua = Lua::new();
6480

@@ -258,6 +274,7 @@ criterion_group! {
258274
create_table,
259275
create_array,
260276
create_string_table,
277+
create_function,
261278
call_lua_function,
262279
call_sum_callback,
263280
call_async_sum_callback,

‎src/lua.rs

+62-57
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,8 @@ use crate::string::String;
1919
use crate::table::Table;
2020
use crate::thread::Thread;
2121
use crate::types::{
22-
Callback, HookCallback, Integer, LightUserData, LuaRef, MaybeSend, Number, RegistryKey,
22+
Callback, CallbackUpvalue, HookCallback, Integer, LightUserData, LuaRef, MaybeSend, Number,
23+
RegistryKey,
2324
};
2425
use crate::userdata::{
2526
AnyUserData, MetaMethod, UserData, UserDataCell, UserDataFields, UserDataMethods,
@@ -38,7 +39,7 @@ use std::rc::Rc;
3839

3940
#[cfg(feature = "async")]
4041
use {
41-
crate::types::AsyncCallback,
42+
crate::types::{AsyncCallback, AsyncCallbackUpvalue, AsyncPollUpvalue},
4243
futures_core::{
4344
future::{Future, LocalBoxFuture},
4445
task::{Context, Poll, Waker},
@@ -396,12 +397,13 @@ impl Lua {
396397
// to prevent them from being garbage collected.
397398

398399
init_gc_metatable_for::<Callback>(state, None)?;
399-
init_gc_metatable_for::<Lua>(state, None)?;
400+
init_gc_metatable_for::<CallbackUpvalue>(state, None)?;
400401
init_gc_metatable_for::<Weak<Mutex<ExtraData>>>(state, None)?;
401402
#[cfg(feature = "async")]
402403
{
403404
init_gc_metatable_for::<AsyncCallback>(state, None)?;
404-
init_gc_metatable_for::<LocalBoxFuture<Result<MultiValue>>>(state, None)?;
405+
init_gc_metatable_for::<AsyncCallbackUpvalue>(state, None)?;
406+
init_gc_metatable_for::<AsyncPollUpvalue>(state, None)?;
405407
init_gc_metatable_for::<Option<Waker>>(state, None)?;
406408

407409
// Create empty Waker slot
@@ -1777,22 +1779,22 @@ impl Lua {
17771779
'lua: 'callback,
17781780
{
17791781
unsafe extern "C" fn call_callback(state: *mut ffi::lua_State) -> c_int {
1780-
callback_error2(state, |nargs| {
1781-
let upvalue_idx1 = ffi::lua_upvalueindex(1);
1782-
let upvalue_idx2 = ffi::lua_upvalueindex(2);
1783-
if ffi::lua_type(state, upvalue_idx1) == ffi::LUA_TNIL
1784-
|| ffi::lua_type(state, upvalue_idx2) == ffi::LUA_TNIL
1785-
{
1782+
let get_extra = |state| {
1783+
let upvalue = get_userdata::<CallbackUpvalue>(state, ffi::lua_upvalueindex(1));
1784+
(*upvalue).lua.extra.clone()
1785+
};
1786+
callback_error_ext(state, get_extra, |nargs| {
1787+
let upvalue_idx = ffi::lua_upvalueindex(1);
1788+
if ffi::lua_type(state, upvalue_idx) == ffi::LUA_TNIL {
17861789
return Err(Error::CallbackDestructed);
17871790
}
1788-
let func = get_userdata::<Callback>(state, upvalue_idx1);
1789-
let lua = get_userdata::<Lua>(state, upvalue_idx2);
1791+
let upvalue = get_userdata::<CallbackUpvalue>(state, upvalue_idx);
17901792

17911793
if nargs < ffi::LUA_MINSTACK {
17921794
check_stack(state, ffi::LUA_MINSTACK - nargs)?;
17931795
}
17941796

1795-
let lua = &mut *lua;
1797+
let lua = &mut (*upvalue).lua;
17961798
lua.state = state;
17971799

17981800
let mut args = MultiValue::new();
@@ -1801,7 +1803,7 @@ impl Lua {
18011803
args.push_front(lua.pop_value());
18021804
}
18031805

1804-
let results = (*func)(lua, args)?;
1806+
let results = ((*upvalue).func)(lua, args)?;
18051807
let nresults = results.len() as c_int;
18061808

18071809
check_stack(state, nresults)?;
@@ -1815,12 +1817,13 @@ impl Lua {
18151817

18161818
unsafe {
18171819
let _sg = StackGuard::new(self.state);
1818-
check_stack(self.state, 5)?;
1820+
check_stack(self.state, 4)?;
18191821

1820-
push_gc_userdata::<Callback>(self.state, mem::transmute(func))?;
1821-
push_gc_userdata(self.state, self.clone())?;
1822-
protect_lua(self.state, 2, 1, |state| {
1823-
ffi::lua_pushcclosure(state, call_callback, 2);
1822+
let lua = self.clone();
1823+
let func = mem::transmute(func);
1824+
push_gc_userdata(self.state, CallbackUpvalue { lua, func })?;
1825+
protect_lua(self.state, 1, 1, |state| {
1826+
ffi::lua_pushcclosure(state, call_callback, 1);
18241827
})?;
18251828

18261829
Ok(Function(self.pop_ref()))
@@ -1844,22 +1847,22 @@ impl Lua {
18441847
}
18451848

18461849
unsafe extern "C" fn call_callback(state: *mut ffi::lua_State) -> c_int {
1847-
callback_error2(state, |nargs| {
1848-
let upvalue_idx1 = ffi::lua_upvalueindex(1);
1849-
let upvalue_idx2 = ffi::lua_upvalueindex(2);
1850-
if ffi::lua_type(state, upvalue_idx1) == ffi::LUA_TNIL
1851-
|| ffi::lua_type(state, upvalue_idx2) == ffi::LUA_TNIL
1852-
{
1850+
let get_extra = |state| {
1851+
let upvalue = get_userdata::<AsyncCallbackUpvalue>(state, ffi::lua_upvalueindex(1));
1852+
(*upvalue).lua.extra.clone()
1853+
};
1854+
callback_error_ext(state, get_extra, |nargs| {
1855+
let upvalue_idx = ffi::lua_upvalueindex(1);
1856+
if ffi::lua_type(state, upvalue_idx) == ffi::LUA_TNIL {
18531857
return Err(Error::CallbackDestructed);
18541858
}
1855-
let func = get_userdata::<AsyncCallback>(state, upvalue_idx1);
1856-
let lua = get_userdata::<Lua>(state, upvalue_idx2);
1859+
let upvalue = get_userdata::<AsyncCallbackUpvalue>(state, upvalue_idx);
18571860

18581861
if nargs < ffi::LUA_MINSTACK {
18591862
check_stack(state, ffi::LUA_MINSTACK - nargs)?;
18601863
}
18611864

1862-
let lua = &mut *lua;
1865+
let lua = &mut (*upvalue).lua;
18631866
lua.state = state;
18641867

18651868
let mut args = MultiValue::new();
@@ -1868,35 +1871,34 @@ impl Lua {
18681871
args.push_front(lua.pop_value());
18691872
}
18701873

1871-
let fut = (*func)(lua, args);
1872-
push_gc_userdata(state, fut)?;
1873-
push_gc_userdata(state, lua.clone())?;
1874-
1875-
protect_lua(state, 2, 1, |state| {
1876-
ffi::lua_pushcclosure(state, poll_future, 2);
1874+
let fut = ((*upvalue).func)(lua, args);
1875+
let lua = lua.clone();
1876+
push_gc_userdata(state, AsyncPollUpvalue { lua, fut })?;
1877+
protect_lua(state, 1, 1, |state| {
1878+
ffi::lua_pushcclosure(state, poll_future, 1);
18771879
})?;
18781880

18791881
Ok(1)
18801882
})
18811883
}
18821884

18831885
unsafe extern "C" fn poll_future(state: *mut ffi::lua_State) -> c_int {
1884-
callback_error2(state, |nargs| {
1885-
let upvalue_idx1 = ffi::lua_upvalueindex(1);
1886-
let upvalue_idx2 = ffi::lua_upvalueindex(2);
1887-
if ffi::lua_type(state, upvalue_idx1) == ffi::LUA_TNIL
1888-
|| ffi::lua_type(state, upvalue_idx2) == ffi::LUA_TNIL
1889-
{
1886+
let get_extra = |state| {
1887+
let upvalue = get_userdata::<AsyncPollUpvalue>(state, ffi::lua_upvalueindex(1));
1888+
(*upvalue).lua.extra.clone()
1889+
};
1890+
callback_error_ext(state, get_extra, |nargs| {
1891+
let upvalue_idx = ffi::lua_upvalueindex(1);
1892+
if ffi::lua_type(state, upvalue_idx) == ffi::LUA_TNIL {
18901893
return Err(Error::CallbackDestructed);
18911894
}
1892-
let fut = get_userdata::<LocalBoxFuture<Result<MultiValue>>>(state, upvalue_idx1);
1893-
let lua = get_userdata::<Lua>(state, upvalue_idx2);
1895+
let upvalue = get_userdata::<AsyncPollUpvalue>(state, upvalue_idx);
18941896

18951897
if nargs < ffi::LUA_MINSTACK {
18961898
check_stack(state, ffi::LUA_MINSTACK - nargs)?;
18971899
}
18981900

1899-
let lua = &mut *lua;
1901+
let lua = &mut (*upvalue).lua;
19001902
lua.state = state;
19011903

19021904
// Try to get an outer poll waker
@@ -1910,7 +1912,8 @@ impl Lua {
19101912

19111913
let mut ctx = Context::from_waker(&waker);
19121914

1913-
match (*fut).as_mut().poll(&mut ctx) {
1915+
let fut = &mut (*upvalue).fut;
1916+
match fut.as_mut().poll(&mut ctx) {
19141917
Poll::Pending => {
19151918
check_stack(state, 1)?;
19161919
ffi::lua_pushboolean(state, 0);
@@ -1932,12 +1935,13 @@ impl Lua {
19321935

19331936
let get_poll = unsafe {
19341937
let _sg = StackGuard::new(self.state);
1935-
check_stack(self.state, 5)?;
1938+
check_stack(self.state, 4)?;
19361939

1937-
push_gc_userdata::<AsyncCallback>(self.state, mem::transmute(func))?;
1938-
push_gc_userdata(self.state, self.clone())?;
1939-
protect_lua(self.state, 2, 1, |state| {
1940-
ffi::lua_pushcclosure(state, call_callback, 2);
1940+
let lua = self.clone();
1941+
let func = mem::transmute(func);
1942+
push_gc_userdata(self.state, AsyncCallbackUpvalue { lua, func })?;
1943+
protect_lua(self.state, 1, 1, |state| {
1944+
ffi::lua_pushcclosure(state, call_callback, 1);
19411945
})?;
19421946

19431947
Function(self.pop_ref())
@@ -2287,13 +2291,14 @@ impl<'lua, T: AsRef<[u8]> + ?Sized> AsChunk<'lua> for T {
22872291

22882292
// An optimized version of `callback_error` that does not allocate `WrappedError+Panic` userdata
22892293
// and instead reuses unsed and cached values from previous calls (or allocates new).
2290-
// It assumes that ephemeral `Lua` struct is passed as a 2nd upvalue.
2291-
pub unsafe fn callback_error2<F, R>(state: *mut ffi::lua_State, f: F) -> R
2294+
// It requires `get_extra` function to return `ExtraData` value.
2295+
unsafe fn callback_error_ext<E, F, R>(state: *mut ffi::lua_State, get_extra: E, f: F) -> R
22922296
where
2297+
E: Fn(*mut ffi::lua_State) -> Arc<Mutex<ExtraData>>,
22932298
F: FnOnce(c_int) -> Result<R>,
22942299
{
2295-
let upvalue_idx2 = ffi::lua_upvalueindex(2);
2296-
if ffi::lua_type(state, upvalue_idx2) == ffi::LUA_TNIL {
2300+
let upvalue_idx = ffi::lua_upvalueindex(1);
2301+
if ffi::lua_type(state, upvalue_idx) == ffi::LUA_TNIL {
22972302
return callback_error(state, f);
22982303
}
22992304

@@ -2314,9 +2319,9 @@ where
23142319

23152320
// We cannot shadow Rust errors with Lua ones, so we need to obtain pre-allocated memory
23162321
// to store a wrapped error or panic *before* we proceed.
2317-
let lua = get_userdata::<Lua>(state, upvalue_idx2);
2322+
let extra = get_extra(state);
23182323
let prealloc_err = {
2319-
let mut extra = mlua_expect!((*lua).extra.lock(), "extra is poisoned");
2324+
let mut extra = mlua_expect!(extra.lock(), "extra is poisoned");
23202325
match extra.prealloc_wrapped_errors.pop() {
23212326
Some(index) => PreallocatedError::Cached(index),
23222327
None => {
@@ -2329,7 +2334,7 @@ where
23292334
};
23302335

23312336
let get_prealloc_err = || {
2332-
let mut extra = mlua_expect!((*lua).extra.lock(), "extra is poisoned");
2337+
let mut extra = mlua_expect!(extra.lock(), "extra is poisoned");
23332338
match prealloc_err {
23342339
PreallocatedError::New(ud) => {
23352340
ffi::lua_settop(state, 1);
@@ -2350,7 +2355,7 @@ where
23502355
match catch_unwind(AssertUnwindSafe(|| f(nargs))) {
23512356
Ok(Ok(r)) => {
23522357
// Return unused WrappedError+Panic to the cache
2353-
let mut extra = mlua_expect!((*lua).extra.lock(), "extra is poisoned");
2358+
let mut extra = mlua_expect!(extra.lock(), "extra is poisoned");
23542359
match prealloc_err {
23552360
PreallocatedError::New(_) if extra.prealloc_wrapped_errors.len() < 16 => {
23562361
ffi::lua_rotate(state, 1, -1);

‎src/scope.rs

+10-26
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ use crate::error::{Error, Result};
1313
use crate::ffi;
1414
use crate::function::Function;
1515
use crate::lua::Lua;
16-
use crate::types::{Callback, LuaRef, MaybeSend};
16+
use crate::types::{Callback, CallbackUpvalue, LuaRef, MaybeSend};
1717
use crate::userdata::{
1818
AnyUserData, MetaMethod, UserData, UserDataCell, UserDataFields, UserDataMethods,
1919
};
@@ -25,8 +25,8 @@ use crate::value::{FromLua, FromLuaMulti, MultiValue, ToLua, ToLuaMulti, Value};
2525

2626
#[cfg(feature = "async")]
2727
use {
28-
crate::types::AsyncCallback,
29-
futures_core::future::{Future, LocalBoxFuture},
28+
crate::types::{AsyncCallback, AsyncCallbackUpvalue, AsyncPollUpvalue},
29+
futures_core::future::Future,
3030
futures_util::future::{self, TryFutureExt},
3131
};
3232

@@ -224,7 +224,7 @@ impl<'lua, 'scope> Scope<'lua, 'scope> {
224224
/// use [`Scope::create_userdata`] instead.
225225
///
226226
/// The main limitation that comes from using non-'static userdata is that the produced userdata
227-
/// will no longer have a `TypeId` associated with it, becuase `TypeId` can only work for
227+
/// will no longer have a `TypeId` associated with it, because `TypeId` can only work for
228228
/// 'static types. This means that it is impossible, once the userdata is created, to get a
229229
/// reference to it back *out* of an `AnyUserData` handle. This also implies that the
230230
/// "function" type methods that can be added via [`UserDataMethods`] (the ones that accept
@@ -460,16 +460,11 @@ impl<'lua, 'scope> Scope<'lua, 'scope> {
460460
// We know the destructor has not run yet because we hold a reference to the callback.
461461

462462
ffi::lua_getupvalue(state, -1, 1);
463-
let ud1 = take_userdata::<Callback>(state);
463+
let ud = take_userdata::<CallbackUpvalue>(state);
464464
ffi::lua_pushnil(state);
465465
ffi::lua_setupvalue(state, -2, 1);
466466

467-
ffi::lua_getupvalue(state, -1, 2);
468-
let ud2 = take_userdata::<Lua>(state);
469-
ffi::lua_pushnil(state);
470-
ffi::lua_setupvalue(state, -2, 2);
471-
472-
vec![Box::new(ud1), Box::new(ud2)]
467+
vec![Box::new(ud)]
473468
});
474469
self.destructors
475470
.borrow_mut()
@@ -510,32 +505,21 @@ impl<'lua, 'scope> Scope<'lua, 'scope> {
510505

511506
// Destroy all upvalues
512507
ffi::lua_getupvalue(state, -1, 1);
513-
let ud1 = take_userdata::<AsyncCallback>(state);
508+
let upvalue1 = take_userdata::<AsyncCallbackUpvalue>(state);
514509
ffi::lua_pushnil(state);
515510
ffi::lua_setupvalue(state, -2, 1);
516511

517-
ffi::lua_getupvalue(state, -1, 2);
518-
let ud2 = take_userdata::<Lua>(state);
519-
ffi::lua_pushnil(state);
520-
ffi::lua_setupvalue(state, -2, 2);
521-
522512
ffi::lua_pop(state, 1);
523-
let mut data: Vec<Box<dyn Any>> = vec![Box::new(ud1), Box::new(ud2)];
513+
let mut data: Vec<Box<dyn Any>> = vec![Box::new(upvalue1)];
524514

525515
// Finally, get polled future and destroy it
526516
f.lua.push_ref(&poll_str.0);
527517
if ffi::lua_rawget(state, -2) == ffi::LUA_TFUNCTION {
528518
ffi::lua_getupvalue(state, -1, 1);
529-
let ud3 = take_userdata::<LocalBoxFuture<Result<MultiValue>>>(state);
519+
let upvalue2 = take_userdata::<AsyncPollUpvalue>(state);
530520
ffi::lua_pushnil(state);
531521
ffi::lua_setupvalue(state, -2, 1);
532-
data.push(Box::new(ud3));
533-
534-
ffi::lua_getupvalue(state, -1, 2);
535-
let ud4 = take_userdata::<Lua>(state);
536-
ffi::lua_pushnil(state);
537-
ffi::lua_setupvalue(state, -2, 2);
538-
data.push(Box::new(ud4));
522+
data.push(Box::new(upvalue2));
539523
}
540524

541525
data

‎src/types.rs

+17
Original file line numberDiff line numberDiff line change
@@ -26,10 +26,27 @@ pub struct LightUserData(pub *mut c_void);
2626
pub(crate) type Callback<'lua, 'a> =
2727
Box<dyn Fn(&'lua Lua, MultiValue<'lua>) -> Result<MultiValue<'lua>> + 'a>;
2828

29+
pub(crate) struct CallbackUpvalue<'lua> {
30+
pub(crate) lua: Lua,
31+
pub(crate) func: Callback<'lua, 'static>,
32+
}
33+
2934
#[cfg(feature = "async")]
3035
pub(crate) type AsyncCallback<'lua, 'a> =
3136
Box<dyn Fn(&'lua Lua, MultiValue<'lua>) -> LocalBoxFuture<'lua, Result<MultiValue<'lua>>> + 'a>;
3237

38+
#[cfg(feature = "async")]
39+
pub(crate) struct AsyncCallbackUpvalue<'lua> {
40+
pub(crate) lua: Lua,
41+
pub(crate) func: AsyncCallback<'lua, 'static>,
42+
}
43+
44+
#[cfg(feature = "async")]
45+
pub(crate) struct AsyncPollUpvalue<'lua> {
46+
pub(crate) lua: Lua,
47+
pub(crate) fut: LocalBoxFuture<'lua, Result<MultiValue<'lua>>>,
48+
}
49+
3350
pub(crate) type HookCallback = Arc<RefCell<dyn FnMut(&Lua, Debug) -> Result<()>>>;
3451

3552
#[cfg(feature = "send")]

0 commit comments

Comments
 (0)