How to implement a select for Lua?

[sxyazi]

I have three asynchronous functions, a(), b(), and c().

I want to implement a select() function that can select and return the one that finishes the fastest among a, b, and c:

#[tokio::main]
async fn main() -> anyhow::Result<()> {
  let lua = mlua::Lua::new();
  let globals = lua.globals();

  globals.set(
    "a",
    lua.create_async_function(|_, ()| async move {
      tokio::time::sleep(std::time::Duration::from_millis(500)).await;
      Ok(1)
    })?,
  )?;
  globals.set(
    "b",
    lua.create_async_function(|_, ()| async move {
      tokio::time::sleep(std::time::Duration::from_millis(300)).await;
      Ok("b")
    })?,
  )?;
  globals.set(
    "c",
    lua.create_async_function(|_, ()| async move {
      tokio::time::sleep(std::time::Duration::from_millis(1000)).await;
      Ok(())
    })?,
  )?;
  globals.set(
    "select",
    lua.create_async_function(|_, ()| async move {
      // ...
    })?,
  )?;

  lua
    .load(
      r#"
local value = select({ a(), b(), c() })
print(value)  -- "b"
"#,
    )
    .exec_async()
    .await;

  return Ok(());
}

Is this possible? Any help would be greatly appreciated!

[khvzak]

It's possible, eg. you can use futures_util::future::select_ok helper:

    let sleep = LuaFunction::wrap_async(|_, secs: f64| async move {
        tokio::time::sleep(tokio::time::Duration::from_secs_f64(secs)).await;
        Ok(secs)
    });
    let select = LuaFunction::wrap_async(|_, futs: Variadic<LuaFunction>| async move {
        let (res, _) = futures_util::future::select_ok(
            futs.into_iter()
                .map(|f| Box::pin(f.call_async::<_, LuaValue>(()))),
        )
        .await?;
        Ok(res)
    });

    lua.load(mlua::chunk! {
        local res = $select(function() return $sleep(0.1) end, function() return $sleep(0.2) end)
        print(res)
    }).exec_async().await?;

[sxyazi]

This is exactly what I wanted, thank you very much! I never thought it would be so neat.