net: Check for closing status when iterating Listener (#3309)

std/http/server.ts: Use listener.next() instead of listener.accept()

4 files changed, 46 insertions, 6 deletions

diff --git a/cli/js/net.ts b/cli/js/net.ts
index f463da60b..8109934a4 100644
--- a/cli/js/net.ts
+++ b/cli/js/net.ts
@@ -82,7 +82,8 @@ export class ListenerImpl implements Listener {
   constructor(
     readonly rid: number,
     private transport: Transport,
-    private localAddr: string
+    private localAddr: string,
+    private closing: boolean = false
   ) {}
 
   async accept(): Promise<Conn> {
@@ -91,6 +92,7 @@ export class ListenerImpl implements Listener {
   }
 
   close(): void {
+    this.closing = true;
     close(this.rid);
   }
 
@@ -102,10 +104,20 @@ export class ListenerImpl implements Listener {
   }
 
   async next(): Promise<IteratorResult<Conn>> {
-    return {
-      done: false,
-      value: await this.accept()
-    };
+    if (this.closing) {
+      return { value: undefined, done: true };
+    }
+    return await this.accept()
+      .then(value => ({ value, done: false }))
+      .catch(e => {
+        // It wouldn't be correct to simply check this.closing here.
+        // TODO: Get a proper error kind for this case, don't check the message.
+        // The current error kind is Other.
+        if (e.message == "Listener has been closed") {
+          return { value: undefined, done: true };
+        }
+        throw e;
+      });
   }
 
   [Symbol.asyncIterator](): AsyncIterator<Conn> {
diff --git a/cli/js/net_test.ts b/cli/js/net_test.ts
index 33f4f7d07..dc7451434 100644
--- a/cli/js/net_test.ts
+++ b/cli/js/net_test.ts
@@ -77,6 +77,18 @@ testPerm({ net: true }, async function netDialListen(): Promise<void> {
   conn.close();
 });
 
+testPerm({ net: true }, async function netListenCloseWhileIterating(): Promise<
+  void
+> {
+  const listener = Deno.listen({ port: 8000 });
+  const nextWhileClosing = listener[Symbol.asyncIterator]().next();
+  listener.close();
+  assertEquals(await nextWhileClosing, { value: undefined, done: true });
+
+  const nextAfterClosing = listener[Symbol.asyncIterator]().next();
+  assertEquals(await nextAfterClosing, { value: undefined, done: true });
+});
+
 /* TODO(ry) Re-enable this test.
 testPerm({ net: true }, async function netListenAsyncIterator(): Promise<void> {
   const listener = Deno.listen(":4500");
diff --git a/std/http/server.ts b/std/http/server.ts
index deedb0f12..ce32511b1 100644
--- a/std/http/server.ts
+++ b/std/http/server.ts
@@ -371,7 +371,9 @@ export class Server implements AsyncIterable<ServerRequest> {
   ): AsyncIterableIterator<ServerRequest> {
     if (this.closing) return;
     // Wait for a new connection.
-    const conn = await this.listener.accept();
+    const { value, done } = await this.listener.next();
+    if (done) return;
+    const conn = value as Conn;
     // Try to accept another connection and add it to the multiplexer.
     mux.add(this.acceptConnAndIterateHttpRequests(mux));
     // Yield the requests that arrive on the just-accepted connection.
diff --git a/std/http/server_test.ts b/std/http/server_test.ts
index 5baeaa144..c8d4080ca 100644
--- a/std/http/server_test.ts
+++ b/std/http/server_test.ts
@@ -14,6 +14,7 @@ import {
   ServerRequest,
   writeResponse,
   readRequest,
+  serve,
   parseHTTPVersion
 } from "./server.ts";
 import { delay } from "../util/async.ts";
@@ -580,4 +581,17 @@ test({
   }
 });
 
+test({
+  name: "[http] close server while iterating",
+  async fn(): Promise<void> {
+    const server = serve(":8123");
+    const nextWhileClosing = server[Symbol.asyncIterator]().next();
+    server.close();
+    assertEquals(await nextWhileClosing, { value: undefined, done: true });
+
+    const nextAfterClosing = server[Symbol.asyncIterator]().next();
+    assertEquals(await nextAfterClosing, { value: undefined, done: true });
+  }
+});
+
 runIfMain(import.meta);