From 329a7eb13252d289ac19745986d7c5d5ac71c4cc Mon Sep 17 00:00:00 2001
From: Mikhail Kornilovich <kornilovichmihail@gmail.com>
Date: Fri, 10 Apr 2026 23:56:11 +0300
Subject: [PATCH] fix windows memory limiter: instead of malloc return null
 make job object IO completion port

---
 runner/limiter_windows.go | 73 ++++++++++++++++++++++++++++++++-------
 1 file changed, 61 insertions(+), 12 deletions(-)

diff --git a/runner/limiter_windows.go b/runner/limiter_windows.go
index a7ba22c..5e9bd52 100644
--- a/runner/limiter_windows.go
+++ b/runner/limiter_windows.go
@@ -13,6 +13,7 @@ import (
 type windowsLimiter struct {
 	memLimit int64
 	job      windows.Handle
+	iocp     windows.Handle
 	peak     int64
 	exceeded bool
 }
@@ -22,9 +23,14 @@ func newLimiter(memLimit int64) limiter {
 }
 
 const (
-	jobObjectExtendedLimitInformationClass = 9
-	jobObjectLimitProcessMemory            = 0x00000100
-	jobObjectLimitKillOnJobClose           = 0x00002000
+	jobObjectExtendedLimitInformationClass           = 9
+	jobObjectAssociateCompletionPortInformationClass = 7
+
+	jobObjectLimitProcessMemory  = 0x00000100
+	jobObjectLimitKillOnJobClose = 0x00002000
+
+	jobObjectMsgProcessMemoryLimit = 9
+	jobObjectMsgJobMemoryLimit     = 10
 )
 
 type ioCounters struct {
@@ -57,6 +63,11 @@ type jobObjectExtendedLimitInformation struct {
 	PeakJobMemoryUsed     uintptr
 }
 
+type jobObjectAssociateCompletionPort struct {
+	CompletionKey  uintptr
+	CompletionPort windows.Handle
+}
+
 func (l *windowsLimiter) prepare(cmd *exec.Cmd) error {
 	if l.memLimit <= 0 {
 		return nil
@@ -78,9 +89,31 @@ func (l *windowsLimiter) prepare(cmd *exec.Cmd) error {
 		uint32(unsafe.Sizeof(info)),
 	); err != nil {
 		windows.CloseHandle(job)
-		return fmt.Errorf("SetInformationJobObject: %w", err)
+		return fmt.Errorf("SetInformationJobObject(extended): %w", err)
 	}
+
+	iocp, err := windows.CreateIoCompletionPort(windows.InvalidHandle, 0, 0, 1)
+	if err != nil {
+		windows.CloseHandle(job)
+		return fmt.Errorf("CreateIoCompletionPort: %w", err)
+	}
+	assoc := jobObjectAssociateCompletionPort{
+		CompletionKey:  uintptr(job),
+		CompletionPort: iocp,
+	}
+	if _, err := windows.SetInformationJobObject(
+		job,
+		jobObjectAssociateCompletionPortInformationClass,
+		uintptr(unsafe.Pointer(&assoc)),
+		uint32(unsafe.Sizeof(assoc)),
+	); err != nil {
+		windows.CloseHandle(iocp)
+		windows.CloseHandle(job)
+		return fmt.Errorf("SetInformationJobObject(iocp): %w", err)
+	}
+
 	l.job = job
+	l.iocp = iocp
 	return nil
 }
 
@@ -107,26 +140,42 @@ func (l *windowsLimiter) collect() limitStats {
 	if l.job == 0 {
 		return limitStats{}
 	}
+
+	if l.iocp != 0 {
+		for {
+			var bytes uint32
+			var key uintptr
+			var overlapped *windows.Overlapped
+			err := windows.GetQueuedCompletionStatus(l.iocp, &bytes, &key, &overlapped, 0)
+			if err != nil {
+				break
+			}
+			if bytes == jobObjectMsgProcessMemoryLimit || bytes == jobObjectMsgJobMemoryLimit {
+				l.exceeded = true
+			}
+		}
+	}
+
 	var info jobObjectExtendedLimitInformation
 	var ret uint32
-	err := windows.QueryInformationJobObject(
+	if err := windows.QueryInformationJobObject(
 		l.job,
 		jobObjectExtendedLimitInformationClass,
 		uintptr(unsafe.Pointer(&info)),
 		uint32(unsafe.Sizeof(info)),
 		&ret,
-	)
-	if err != nil {
-		return limitStats{}
-	}
-	l.peak = int64(info.PeakProcessMemoryUsed)
-	if l.memLimit > 0 && l.peak >= l.memLimit {
-		l.exceeded = true
+	); err == nil {
+		l.peak = int64(info.PeakProcessMemoryUsed)
 	}
+
 	return limitStats{PeakMemory: l.peak, MemoryExceeded: l.exceeded}
 }
 
 func (l *windowsLimiter) cleanup() {
+	if l.iocp != 0 {
+		windows.CloseHandle(l.iocp)
+		l.iocp = 0
+	}
 	if l.job != 0 {
 		windows.CloseHandle(l.job)
 		l.job = 0