From cc169d1c0c1317bd384c8575f74049a43ce0097d Mon Sep 17 00:00:00 2001
From: Milly <milly.ca@gmail.com>
Date: Sun, 1 Jun 2025 21:32:46 +0900
Subject: [PATCH] :bug: Fix race condition in lambda.register with `once`
 option

Ensure the lambda is unregistered before invoking the function when `once`
option is specified, preventing multiple concurrent calls from executing
the function more than once.

Update test to cover the case where the lambda is called multiple times before
the first call resolves.
---
 deno.jsonc         |  1 +
 lambda/mod.ts      |  7 ++-----
 lambda/mod_test.ts | 18 +++++++++++++-----
 3 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/deno.jsonc b/deno.jsonc
index e6cf5845..e0be4d8d 100644
--- a/deno.jsonc
+++ b/deno.jsonc
@@ -47,6 +47,7 @@
     ]
   },
   "imports": {
+    "@core/asyncutil": "jsr:@core/asyncutil@^1.2.0",
     "@core/unknownutil": "jsr:@core/unknownutil@^4.3.0",
     "@denops/core": "jsr:@denops/core@^7.0.0",
     "@denops/test": "jsr:@denops/test@^3.0.1",
diff --git a/lambda/mod.ts b/lambda/mod.ts
index b4bc85d6..dac6e478 100644
--- a/lambda/mod.ts
+++ b/lambda/mod.ts
@@ -120,11 +120,8 @@ export function register(
   const id = `lambda:${denops.name}:${uuid}`;
   if (options.once) {
     denops.dispatcher[id] = async (...args: unknown[]) => {
-      try {
-        return await fn(...args);
-      } finally {
-        unregister(denops, id);
-      }
+      unregister(denops, id);
+      return await fn(...args);
     };
   } else {
     denops.dispatcher[id] = async (...args: unknown[]) => {
diff --git a/lambda/mod_test.ts b/lambda/mod_test.ts
index acdf33a5..9cf0a62f 100644
--- a/lambda/mod_test.ts
+++ b/lambda/mod_test.ts
@@ -6,6 +6,7 @@ import {
   returnsNext,
   spy,
 } from "@std/testing/mock";
+import { flushPromises } from "@core/asyncutil";
 import { assert, is } from "@core/unknownutil";
 import { test } from "@denops/test";
 import { expr, rawString } from "../eval/mod.ts";
@@ -72,16 +73,23 @@ test({
       });
       await t.step("if 'once' option is specified", async (t) => {
         await t.step("registers an oneshot lambda function", async (t) => {
-          const fn = spy(returnsNext(["foo"]));
+          const waiter = Promise.withResolvers<void>();
+          const fn = spy(resolvesNext([waiter.promise.then(() => "foo")]));
           const id = lambda.register(denops, fn, { once: true });
           assertSpyCalls(fn, 0);
-          assertEquals(await denops.dispatch(denops.name, id), "foo");
+
+          // Call the lambda function twice before the first call resolves
+          const firstCall = denops.dispatch(denops.name, id);
+          const secondCall = denops.dispatch(denops.name, id);
+          secondCall.catch(NOOP);
+          await flushPromises();
+          waiter.resolve();
+
+          assertEquals(await firstCall, "foo");
           assertSpyCalls(fn, 1);
 
           await t.step("which will be removed if called once", async () => {
-            const error = await assertRejects(
-              () => denops.dispatch(denops.name, id),
-            );
+            const error = await assertRejects(() => secondCall);
             assertStringIncludes(
               error as string,
               "denops.dispatcher[name] is not a function",