fix: a bunch of selectors

Author: algebraic-dev

src/Std/Internal/Async/Signal.lean

@@ -242,19 +242,21 @@ Create a `Selector` that resolves once `s` has received the signal. Note that ca
 if it hasn't already started.
 -/
 def selector (s : Signal.Waiter) : IO (Selector Unit) := do
-  let signalWaiter ← s.wait
   return {
     tryFn := do
+      let signalWaiter : AsyncTask _ ← async s.wait
       if ← IO.hasFinished signalWaiter then
         return some ()
       else
         return none
 
     registerFn waiter := do
+      let signalWaiter ← s.wait
       discard <| AsyncTask.mapIO (x := signalWaiter) fun _ => do
         let lose := return ()
         let win promise := promise.resolve (.ok ())
         waiter.race lose win
 
-    unregisterFn := s.stop
+    unregisterFn := s.native.cancel
+
   }

src/Std/Internal/Async/TCP.lean

@@ -70,6 +70,40 @@ def accept (s : Server) : Async Client := do
   |> Async.ofPromise
   |>.map Client.ofNative
 
+/--
+Tries to accepts an incoming connection.
+-/
+@[inline]
+def tryAccept (s : Server) : Async (Option Client) := do
+  let res ← s.native.tryAccept
+  let socket ← Async.ofExcept res
+  return Client.ofNative <$> socket
+
+/--
+Creates a `Selector` that resolves once `s` has a connetion available. Calling this function
+does not starts the connection wait, so it must not be called in parallel with `accept`.
+-/
+def acceptSelector (s : TCP.Socket.Server) : Async (Selector Client) := do
+  return {
+    tryFn := s.tryAccept
+    registerFn waiter := do
+      let task ← s.native.accept
+      -- If we get cancelled the promise will be dropped so prepare for that
+      IO.chainTask (t := task.result?) fun res => do
+        match res with
+        | none => return ()
+        | some res =>
+          let lose := return ()
+          let win promise := do
+            try
+              let result ← IO.ofExcept res
+              promise.resolve (.ok (Client.ofNative result))
+            catch e =>
+              promise.resolve (.error e)
+          waiter.race lose win
+    unregisterFn := s.native.cancelAccept
+  }
+
 /--
 Gets the local address of the server socket.
 -/

src/Std/Internal/UV/Signal.lean

@@ -78,6 +78,15 @@ This function has different behavior depending on the state of the `Signal`:
 @[extern "lean_uv_signal_stop"]
 opaque stop (signal : @& Signal) : IO Unit
 
+/--
+This function has different behavior depending on the state of the `Signal`:
+- If it is initial or finished this is a no-op.
+- If it's running then it drops the accept promise and if it's not repeatable it finishes
+  the execution of the timer.
+-/
+@[extern "lean_uv_signal_cancel"]
+opaque cancel (signal : @& Signal) : IO Unit
+
 end Signal
 
 end UV

src/Std/Internal/UV/TCP.lean

@@ -96,6 +96,18 @@ Accepts an incoming connection on a listening TCP socket.
 @[extern "lean_uv_tcp_accept"]
 opaque accept (socket : @& Socket) : IO (IO.Promise (Except IO.Error Socket))
 
+/--
+Tries to accepts an incoming connection on a listening TCP socket.
+-/
+@[extern "lean_uv_tcp_try_accept"]
+opaque tryAccept (socket : @& Socket) : IO ((Except IO.Error (Option Socket)))
+
+/--
+Cancels the accept request of a socket.
+-/
+@[extern "lean_uv_tcp_cancel_accept"]
+opaque cancelAccept (socket : @& Socket) : IO Unit
+
 /--
 Shuts down an incoming connection on a listening TCP socket.
 -/

src/runtime/uv/signal.cpp

@@ -219,8 +219,8 @@ extern "C" LEAN_EXPORT lean_obj_res lean_uv_signal_stop(b_obj_arg obj, obj_arg /
     }
 }
 
-/* Std.Internal.UV.Signal.cancelNext (signal : @& Signal) : IO Unit */
-extern "C" LEAN_EXPORT lean_obj_res lean_uv_signal_cancel_next(b_obj_arg obj, obj_arg /* w */) {
+/* Std.Internal.UV.Signal.cancel (signal : @& Signal) : IO Unit */
+extern "C" LEAN_EXPORT lean_obj_res lean_uv_signal_cancel(b_obj_arg obj, obj_arg /* w */) {
     lean_uv_signal_object * signal = lean_to_uv_signal(obj);
 
     // It's locking here to avoid changing the state during other operations.
@@ -267,6 +267,13 @@ extern "C" LEAN_EXPORT lean_obj_res lean_uv_signal_stop(b_obj_arg signal, obj_ar
     );
 }
 
+/* Std.Internal.UV.Signal.cancel (signal : @& Signal) : IO Unit */
+extern "C" LEAN_EXPORT lean_obj_res lean_uv_signal_cancel(b_obj_arg obj, obj_arg /* w */) {
+    lean_always_assert(
+        false && ("Please build a version of Lean4 with libuv to invoke this.")
+    );
+}
+
 #endif
 
 }

src/runtime/uv/signal.h

@@ -51,5 +51,6 @@ static inline lean_uv_signal_object* lean_to_uv_signal(lean_object * o) { return
 extern "C" LEAN_EXPORT lean_obj_res lean_uv_signal_mk(uint32_t signum_obj, uint8_t repeating, obj_arg /* w */);
 extern "C" LEAN_EXPORT lean_obj_res lean_uv_signal_next(b_obj_arg signal, obj_arg /* w */);
 extern "C" LEAN_EXPORT lean_obj_res lean_uv_signal_stop(b_obj_arg signal, obj_arg /* w */);
+extern "C" LEAN_EXPORT lean_obj_res lean_uv_signal_cancel(b_obj_arg obj, obj_arg /* w */);
 
 }

src/runtime/uv/tcp.cpp

@@ -231,7 +231,7 @@ extern "C" LEAN_EXPORT lean_obj_res lean_uv_tcp_send(b_obj_arg socket, obj_arg d
         lean_dec(socket);
         lean_dec(data_array);
         free(bufs);
-        
+
         free(write_uv->data);
         free(write_uv);
 
@@ -524,6 +524,37 @@ extern "C" LEAN_EXPORT lean_obj_res lean_uv_tcp_accept(b_obj_arg socket, obj_arg
     return lean_io_result_mk_ok(promise);
 }
 
+/* Std.Internal.UV.TCP.Socket.tryAccept (socket : @& Socket) : IO (Except IO.Error (Option Socket)) */
+extern "C" LEAN_EXPORT lean_obj_res lean_uv_tcp_try_accept(b_obj_arg socket, obj_arg /* w */) {
+    lean_uv_tcp_socket_object* tcp_socket = lean_to_uv_tcp_socket(socket);
+
+    // Locking early prevents potential parallelism issues setting m_promise_accept.
+    event_loop_lock(&global_ev);
+
+    if (tcp_socket->m_promise_accept != nullptr) {
+        return lean_io_result_mk_error(lean_decode_uv_error(UV_EALREADY, mk_string("parallel accept is not allowed! consider binding multiple sockets to the same address and accepting on them instead")));
+    }
+
+    lean_object* client = lean_io_result_take_value(lean_uv_tcp_new(lean_box(0)));
+    lean_uv_tcp_socket_object* client_socket = lean_to_uv_tcp_socket(client);
+
+    int result = uv_accept((uv_stream_t*)tcp_socket->m_uv_tcp, (uv_stream_t*)client_socket->m_uv_tcp);
+
+    if (result < 0 && result != UV_EAGAIN) {
+        event_loop_unlock(&global_ev);
+        lean_dec(client);
+        return lean_io_result_mk_error(lean_decode_uv_error(result, NULL));
+    } else if (result >= 0) {
+        event_loop_unlock(&global_ev);
+        return lean_io_result_mk_ok(mk_except_ok(lean::mk_option_some(client)));
+    } else {
+        event_loop_unlock(&global_ev);
+        lean_dec(client);
+        return lean_io_result_mk_ok(mk_except_ok(lean::mk_option_none()));
+    }
+}
+
+
 
 /* Std.Internal.UV.TCP.Socket.cancelAccept (socket : @& Socket) : IO Unit */
 extern "C" LEAN_EXPORT lean_obj_res lean_uv_tcp_cancel_accept(b_obj_arg socket, obj_arg /* w */) {
@@ -718,6 +749,12 @@ extern "C" LEAN_EXPORT lean_obj_res lean_uv_tcp_listen(b_obj_arg socket, int32_t
     );
 }
 
+extern "C" LEAN_EXPORT lean_obj_res lean_uv_tcp_cancel_accept(b_obj_arg socket, obj_arg /* w */) {
+    lean_always_assert(
+        false && ("Please build a version of Lean4 with libuv to invoke this.")
+    );
+}
+
 extern "C" LEAN_EXPORT lean_obj_res lean_uv_tcp_accept(b_obj_arg socket, obj_arg /* w */) {
     lean_always_assert(
         false && ("Please build a version of Lean4 with libuv to invoke this.")

src/runtime/uv/tcp.h

@@ -35,7 +35,6 @@ typedef struct {
 // Tcp socket object manipulation functions.
 static inline lean_object* lean_uv_tcp_socket_new(lean_uv_tcp_socket_object* s) { return lean_alloc_external(g_uv_tcp_socket_external_class, s); }
 static inline lean_uv_tcp_socket_object* lean_to_uv_tcp_socket(lean_object* o) { return (lean_uv_tcp_socket_object*)(lean_get_external_data(o)); }
-
 #endif
 
 // =======================================
@@ -50,7 +49,10 @@ extern "C" LEAN_EXPORT lean_obj_res lean_uv_tcp_cancel_recv(b_obj_arg socket, ob
 extern "C" LEAN_EXPORT lean_obj_res lean_uv_tcp_bind(b_obj_arg socket, b_obj_arg addr, obj_arg /* w */);
 extern "C" LEAN_EXPORT lean_obj_res lean_uv_tcp_listen(b_obj_arg socket, int32_t backlog, obj_arg /* w */);
 extern "C" LEAN_EXPORT lean_obj_res lean_uv_tcp_accept(b_obj_arg socket, obj_arg /* w */);
+extern "C" LEAN_EXPORT lean_obj_res lean_uv_tcp_cancel_accept(b_obj_arg socket, obj_arg /* w */);
+extern "C" LEAN_EXPORT lean_obj_res lean_uv_tcp_try_accept(b_obj_arg socket, obj_arg /* w */);
 extern "C" LEAN_EXPORT lean_obj_res lean_uv_tcp_shutdown(b_obj_arg socket, obj_arg /* w */);
+extern "C" LEAN_EXPORT lean_obj_res lean_uv_tcp_try_accept(b_obj_arg socket, obj_arg /* w */);
 
 // =======================================
 // TCP Socket Utility Functions