From 726fde7e1f5274efd1c246be9bb739e46e3f7955 Mon Sep 17 00:00:00 2001
From: Alex Orlenko <zxteam@protonmail.com>
Date: Wed, 3 Mar 2021 23:21:56 +0000
Subject: [PATCH] Optimise async callbacks (polling)

call async Rust callback [sum] 3 10
                        time:   [59.338 us 59.729 us 60.097 us]
                        change: [-10.336% -8.6212% -6.8003%] (p = 0.00 < 0.05)
                        Performance has improved.
---
 benches/benchmark.rs |  6 +++++-
 src/lua.rs           | 15 ++++++++++-----
 2 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/benches/benchmark.rs b/benches/benchmark.rs
index b74b01b..7f59ce2 100644
--- a/benches/benchmark.rs
+++ b/benches/benchmark.rs
@@ -13,6 +13,7 @@ extern "system" {}
 use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
 use std::time::Duration;
 use tokio::runtime::Runtime;
+use tokio::task;
 
 use mlua::prelude::*;
 
@@ -117,7 +118,10 @@ fn call_sum_callback(c: &mut Criterion) {
 fn call_async_sum_callback(c: &mut Criterion) {
     let lua = Lua::new();
     let callback = lua
-        .create_async_function(|_, (a, b, c): (i64, i64, i64)| async move { Ok(a + b + c) })
+        .create_async_function(|_, (a, b, c): (i64, i64, i64)| async move {
+            task::yield_now().await;
+            Ok(a + b + c)
+        })
         .unwrap();
     lua.globals().set("callback", callback).unwrap();
 
diff --git a/src/lua.rs b/src/lua.rs
index a157c06..9e94e19 100644
--- a/src/lua.rs
+++ b/src/lua.rs
@@ -1735,10 +1735,9 @@ impl Lua {
 
                 match (*fut).as_mut().poll(&mut ctx) {
                     Poll::Pending => {
-                        check_stack(state, 6)?;
+                        check_stack(state, 1)?;
                         ffi::lua_pushboolean(state, 0);
-                        push_gc_userdata(state, AsyncPollPending)?;
-                        Ok(2)
+                        Ok(1)
                     }
                     Poll::Ready(results) => {
                         let results = results?;
@@ -1782,18 +1781,24 @@ impl Lua {
                 ))
             })?,
         )?;
+        env.set("pending", unsafe {
+            let _sg = StackGuard::new(self.state);
+            check_stack(self.state, 5)?;
+            push_gc_userdata(self.state, AsyncPollPending)?;
+            self.pop_value()
+        })?;
 
         // We set `poll` variable in the env table to be able to destroy upvalues
         self.load(
             r#"
             poll = get_poll(...)
-            local poll, yield, unpack = poll, yield, unpack
+            local poll, pending, yield, unpack = poll, pending, yield, unpack
             while true do
                 local ready, res, nres = poll()
                 if ready then
                     return unpack(res, nres)
                 end
-                yield(res)
+                yield(pending)
             end
             "#,
         )