Skip to content

Commit 2665f6a

Browse files
committed
feat(iocp): prefer async connect
1 parent 9f785d2 commit 2665f6a

File tree

1 file changed

+72
-7
lines changed

1 file changed

+72
-7
lines changed

src/backend/iocp.zig

Lines changed: 72 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -545,13 +545,58 @@ pub const Loop = struct {
545545
.close => |v| .{ .result = .{ .close = windows.CloseHandle(v.fd) } },
546546

547547
.connect => |*v| action: {
548-
const result = windows.ws2_32.connect(asSocket(v.socket), &v.addr.any, @as(i32, @intCast(v.addr.getOsSockLen())));
549-
if (result != 0) {
550-
const err = windows.ws2_32.WSAGetLastError();
551-
break :action switch (err) {
552-
else => .{ .result = .{ .connect = windows.unexpectedWSAError(err) } },
548+
const as_socket = asSocket(v.socket);
549+
// Associate our socket with loop's completion port.
550+
self.associate_fd(v.socket) catch unreachable;
551+
552+
// ConnectEx requires socket to be initially bound.
553+
// https://github.com/tigerbeetle/tigerbeetle/blob/main/src/io/windows.zig#L467
554+
{
555+
const inaddr_any = std.mem.zeroes([4]u8);
556+
const bind_addr = std.net.Address.initIp4(inaddr_any, 0);
557+
// NOTE: This may return many other errors; we should extend `ConnectError` set.
558+
posix.bind(as_socket, &bind_addr.any, bind_addr.getOsSockLen()) catch unreachable;
559+
}
560+
561+
// NOTE: This can be declared in somewhere else; it all happens in comptime though so no issue putting it here.
562+
const LPFN_CONNECTEX = *const fn (
563+
Socket: windows.ws2_32.SOCKET,
564+
SockAddr: *const windows.ws2_32.sockaddr,
565+
SockLen: posix.socklen_t,
566+
SendBuf: ?*const anyopaque,
567+
SendBufLen: windows.DWORD,
568+
BytesSent: *windows.DWORD,
569+
Overlapped: *windows.OVERLAPPED,
570+
) callconv(.winapi) windows.BOOL;
571+
572+
// Dynamically load the ConnectEx function.
573+
const ConnectEx = windows.loadWinsockExtensionFunction(LPFN_CONNECTEX, as_socket, windows.ws2_32.WSAID_CONNECTEX) catch |err| switch (err) {
574+
error.OperationNotSupported => unreachable, // Something other than sockets has given.
575+
error.FileDescriptorNotASocket => unreachable, // Must be preferred on a socket.
576+
error.ShortRead => unreachable,
577+
error.Unexpected => break :action .{ .result = .{ .connect = error.Unexpected } },
578+
};
579+
580+
// Connect attempt.
581+
var bytes_transferred: windows.DWORD = 0;
582+
const result = ConnectEx(as_socket, &v.addr.any, v.addr.getOsSockLen(), null, 0, &bytes_transferred, &completion.overlapped);
583+
584+
// If ConnectEx returns `windows.TRUE`, it means operation completed immediately.
585+
// Which is most of the time not the case; we should check it anyways though!
586+
if (result == windows.FALSE) {
587+
// NOTE: This may return many other errors; we should extend `ConnectError` set.
588+
break :action switch (windows.ws2_32.WSAGetLastError()) {
589+
.WSA_IO_PENDING, .WSAEWOULDBLOCK, .WSA_IO_INCOMPLETE => .{ .submitted = {} }, // Operation will be completed in the future.
590+
else => |err| .{ .result = .{ .connect = windows.unexpectedWSAError(err) } },
553591
};
554592
}
593+
594+
// Surprisingly, we connected immediately.
595+
// The following is necessary for various functions (shutdown, getsockopt, setsockopt, getsockname, getpeername) to work.
596+
// https://learn.microsoft.com/en-us/windows/win32/api/mswsock/nc-mswsock-lpfn_connectex#remarks
597+
// https://stackoverflow.com/questions/13598530/connectex-requires-the-socket-to-be-initially-bound-but-to-what
598+
_ = windows.ws2_32.setsockopt(asSocket(v.socket), windows.ws2_32.SOL.SOCKET, windows.ws2_32.SO.UPDATE_CONNECT_CONTEXT, null, 0);
599+
555600
break :action .{ .result = .{ .connect = {} } };
556601
},
557602

@@ -960,7 +1005,7 @@ pub const Completion = struct {
9601005
/// operation for the completion.
9611006
pub fn perform(self: *Completion) Result {
9621007
return switch (self.op) {
963-
.noop, .close, .connect, .shutdown, .timer, .cancel => {
1008+
.noop, .close, .shutdown, .timer, .cancel => {
9641009
std.log.warn("perform op={s}", .{@tagName(self.op)});
9651010
unreachable;
9661011
},
@@ -985,7 +1030,27 @@ pub const Completion = struct {
9851030
return .{ .accept = self.op.accept.internal_accept_socket.? };
9861031
},
9871032

988-
.read => |*v| {
1033+
.connect => |*v| r: {
1034+
const as_socket = asSocket(v.socket);
1035+
var transferred: windows.DWORD = 0;
1036+
var flags: windows.DWORD = 0;
1037+
const result = windows.ws2_32.WSAGetOverlappedResult(as_socket, &self.overlapped, &transferred, windows.FALSE, &flags);
1038+
1039+
// Connected successfully.
1040+
if (result == windows.TRUE) {
1041+
// The following is necessary for various functions (shutdown, getsockopt, setsockopt, getsockname, getpeername) to work.
1042+
// https://learn.microsoft.com/en-us/windows/win32/api/mswsock/nc-mswsock-lpfn_connectex#remarks
1043+
// https://stackoverflow.com/questions/13598530/connectex-requires-the-socket-to-be-initially-bound-but-to-what
1044+
_ = windows.ws2_32.setsockopt(as_socket, windows.ws2_32.SOL.SOCKET, windows.ws2_32.SO.UPDATE_CONNECT_CONTEXT, null, 0);
1045+
1046+
break :r .{ .connect = {} };
1047+
}
1048+
1049+
// We got an error.
1050+
break :r .{ .connect = windows.unexpectedWSAError(windows.ws2_32.WSAGetLastError()) };
1051+
},
1052+
1053+
.read => |*v| r: {
9891054
var bytes_transferred: windows.DWORD = 0;
9901055
const result = windows.kernel32.GetOverlappedResult(v.fd, &self.overlapped, &bytes_transferred, windows.FALSE);
9911056
if (result == windows.FALSE) {

0 commit comments

Comments
 (0)