Make Process#wait asynchronous on Windows (#13908)

Makefile.win

@@ -227,7 +227,7 @@ $(O)\crystal.exe: $(DEPS) $(SOURCES) $(O)\crystal.res
 	@$(call MKDIR,"$(O)")
 	$(call export_vars)
 	$(call export_build_vars)
-	.\bin\crystal build $(FLAGS) -o "$(O)\crystal-next.exe" src\compiler\crystal.cr -D without_openssl -D without_zlib -D without_playground $(if $(USE_PCRE1),-D use_pcre,-D use_pcre2) --link-flags=/PDBALTPATH:crystal.pdb "--link-flags=$(realpath $(O)\crystal.res)"
+	.\bin\crystal build $(FLAGS) -o "$(O)\crystal-next.exe" src\compiler\crystal.cr -D without_openssl -D without_zlib $(if $(USE_PCRE1),-D use_pcre,-D use_pcre2) --link-flags=/PDBALTPATH:crystal.pdb "--link-flags=$(realpath $(O)\crystal.res)"
 	$(call MV,"$(O)\crystal-next.exe","$@")
 	$(call MV,"$(O)\crystal-next.pdb","$(O)\crystal.pdb")
 

src/crystal/system/win32/process.cr

@@ -1,5 +1,6 @@
 require "c/processthreadsapi"
 require "c/handleapi"
+require "c/jobapi2"
 require "c/synchapi"
 require "c/tlhelp32"
 require "process/shell"
@@ -9,6 +10,8 @@ struct Crystal::System::Process
   getter pid : LibC::DWORD
   @thread_id : LibC::DWORD
   @process_handle : LibC::HANDLE
+  @job_object : LibC::HANDLE
+  @completion_key = IO::Overlapped::CompletionKey.new
 
   @@interrupt_handler : Proc(Nil)?
   @@interrupt_count = Crystal::AtomicSemaphore.new
@@ -19,15 +22,62 @@ struct Crystal::System::Process
     @pid = process_info.dwProcessId
     @thread_id = process_info.dwThreadId
     @process_handle = process_info.hProcess
+
+    @job_object = LibC.CreateJobObjectW(nil, nil)
+
+    # enable IOCP notifications
+    config_job_object(
+      LibC::JOBOBJECTINFOCLASS::AssociateCompletionPortInformation,
+      LibC::JOBOBJECT_ASSOCIATE_COMPLETION_PORT.new(
+        completionKey: @completion_key.as(Void*),
+        completionPort: Crystal::Scheduler.event_loop.iocp,
+      ),
+    )
+
+    # but not for any child processes
+    config_job_object(
+      LibC::JOBOBJECTINFOCLASS::ExtendedLimitInformation,
+      LibC::JOBOBJECT_EXTENDED_LIMIT_INFORMATION.new(
+        basicLimitInformation: LibC::JOBOBJECT_BASIC_LIMIT_INFORMATION.new(
+          limitFlags: LibC::JOB_OBJECT_LIMIT_SILENT_BREAKAWAY_OK,
+        ),
+      ),
+    )
+
+    if LibC.AssignProcessToJobObject(@job_object, @process_handle) == 0
+      raise RuntimeError.from_winerror("AssignProcessToJobObject")
+    end
+  end
+
+  private def config_job_object(kind, info)
+    if LibC.SetInformationJobObject(@job_object, kind, pointerof(info), sizeof(typeof(info))) == 0
+      raise RuntimeError.from_winerror("SetInformationJobObject")
+    end
   end
 
   def release
     return if @process_handle == LibC::HANDLE.null
     close_handle(@process_handle)
     @process_handle = LibC::HANDLE.null
+    close_handle(@job_object)
+    @job_object = LibC::HANDLE.null
   end
 
   def wait
+    if LibC.GetExitCodeProcess(@process_handle, out exit_code) == 0
+      raise RuntimeError.from_winerror("GetExitCodeProcess")
+    end
+    return exit_code unless exit_code == LibC::STILL_ACTIVE
+
+    # let `@job_object` do its job
+    # TODO: message delivery is "not guaranteed"; does it ever happen? Are we
+    # stuck forever in that case?
+    # (https://learn.microsoft.com/en-us/windows/win32/api/winnt/ns-winnt-jobobject_associate_completion_port)
+    @completion_key.fiber = ::Fiber.current
+    Crystal::Scheduler.reschedule
+
+    # If the IOCP notification is delivered before the process fully exits,
+    # wait for it
     if LibC.WaitForSingleObject(@process_handle, LibC::INFINITE) != LibC::WAIT_OBJECT_0
       raise RuntimeError.from_winerror("WaitForSingleObject")
     end
@@ -38,7 +88,7 @@ struct Crystal::System::Process
     # waitpid returns, we wait 5 milliseconds to attempt to replicate this behaviour.
     sleep 5.milliseconds
 
-    if LibC.GetExitCodeProcess(@process_handle, out exit_code) == 0
+    if LibC.GetExitCodeProcess(@process_handle, pointerof(exit_code)) == 0
       raise RuntimeError.from_winerror("GetExitCodeProcess")
     end
     if exit_code == LibC::STILL_ACTIVE

src/io/overlapped.cr

@@ -3,6 +3,11 @@ require "c/handleapi"
 require "crystal/system/thread_linked_list"
 
 module IO::Overlapped
+  # :nodoc:
+  class CompletionKey
+    property fiber : Fiber?
+  end
+
   @read_timeout : Time::Span?
   @write_timeout : Time::Span?
 
@@ -61,7 +66,23 @@ module IO::Overlapped
     end
 
     removed.times do |i|
-      OverlappedOperation.schedule(overlapped_entries[i].lpOverlapped) { |fiber| yield fiber }
+      entry = overlapped_entries[i]
+
+      # at the moment only `::Process#wait` uses a non-nil completion key; all
+      # I/O operations, including socket ones, do not set this field
+      case completion_key = Pointer(Void).new(entry.lpCompletionKey).as(CompletionKey?)
+      when Nil
+        OverlappedOperation.schedule(entry.lpOverlapped) { |fiber| yield fiber }
+      else
+        case entry.dwNumberOfBytesTransferred
+        when LibC::JOB_OBJECT_MSG_EXIT_PROCESS, LibC::JOB_OBJECT_MSG_ABNORMAL_EXIT_PROCESS
+          if fiber = completion_key.fiber
+            yield fiber
+          else
+            # the `Process` exits before a call to `#wait`; do nothing
+          end
+        end
+      end
     end
 
     false

src/lib_c/x86_64-windows-msvc/c/jobapi2.cr

@@ -0,0 +1,7 @@
+require "./winnt"
+
+lib LibC
+  fun CreateJobObjectW(lpJobAttributes : SECURITY_ATTRIBUTES*, lpName : LPWSTR) : HANDLE
+  fun SetInformationJobObject(hJob : HANDLE, jobObjectInformationClass : JOBOBJECTINFOCLASS, lpJobObjectInformation : Void*, cbJobObjectInformationLength : DWORD) : BOOL
+  fun AssignProcessToJobObject(hJob : HANDLE, hProcess : HANDLE) : BOOL
+end

src/lib_c/x86_64-windows-msvc/c/winnt.cr

@@ -90,6 +90,51 @@ lib LibC
     WRITE = 0x20006
   end
 
+  enum JOBOBJECTINFOCLASS
+    AssociateCompletionPortInformation = 7
+    ExtendedLimitInformation           = 9
+  end
+
+  struct JOBOBJECT_BASIC_LIMIT_INFORMATION
+    perProcessUserTimeLimit : LARGE_INTEGER
+    perJobUserTimeLimit : LARGE_INTEGER
+    limitFlags : DWORD
+    minimumWorkingSetSize : SizeT
+    maximumWorkingSetSize : SizeT
+    activeProcessLimit : DWORD
+    affinity : ULONG_PTR
+    priorityClass : DWORD
+    schedulingClass : DWORD
+  end
+
+  struct IO_COUNTERS
+    readOperationCount : ULongLong
+    writeOperationCount : ULongLong
+    otherOperationCount : ULongLong
+    readTransferCount : ULongLong
+    writeTransferCount : ULongLong
+    otherTransferCount : ULongLong
+  end
+
+  struct JOBOBJECT_EXTENDED_LIMIT_INFORMATION
+    basicLimitInformation : JOBOBJECT_BASIC_LIMIT_INFORMATION
+    ioInfo : IO_COUNTERS
+    processMemoryLimit : SizeT
+    jobMemoryLimit : SizeT
+    peakProcessMemoryUsed : SizeT
+    peakJobMemoryUsed : SizeT
+  end
+
+  struct JOBOBJECT_ASSOCIATE_COMPLETION_PORT
+    completionKey : Void*
+    completionPort : HANDLE
+  end
+
+  JOB_OBJECT_LIMIT_SILENT_BREAKAWAY_OK = 0x00001000
+
+  JOB_OBJECT_MSG_EXIT_PROCESS          = 7
+  JOB_OBJECT_MSG_ABNORMAL_EXIT_PROCESS = 8
+
   struct CONTEXT
     p1Home : DWORD64
     p2Home : DWORD64