diff --git a/src/bitworld/native_websocket.nim b/src/bitworld/native_websocket.nim index 00be255..2744b4b 100644 --- a/src/bitworld/native_websocket.nim +++ b/src/bitworld/native_websocket.nim @@ -2,6 +2,7 @@ ## owns receive. Sends and libcurl calls are serialized without a hidden reader. import std/[atomics, base64, locks, monotimes, os, posix, sequtils, sha1, strutils, sysrand, times, unicode, uri] +from system/ansi_c import c_malloc, c_free import libcurl except Option import native_stop @@ -19,7 +20,10 @@ type handle: PCurl socket: cint ioLock, sendLock: Lock - incoming, fragments, pending: string + incoming, fragments: string + # Writer ownership may move between joined threads; Nim thread heaps cannot. + pending: ptr UncheckedArray[byte] + pendingLength, pendingOffset: int fragmentOpcode: int pongState: Atomic[PongState] pongData: string @@ -117,25 +121,36 @@ proc receiveBytes(ws: NativeWebSocket, deadline: MonoTime, let ready = ws.waitSocket(POLLIN, deadline, purpose) if ready.kind != wsReady: return ready +proc installPending(ws: NativeWebSocket, data: string) = + doAssert ws.pendingLength == 0 + ws.pending = cast[ptr UncheckedArray[byte]](c_malloc(csize_t(data.len))) + doAssert ws.pending != nil, "WebSocket frame allocation failed" + ws.pendingLength = data.len + copyMem(ws.pending, data[0].unsafeAddr, data.len) + proc flushPending(ws: NativeWebSocket, deadline: MonoTime, purpose: IoPurpose): WebSocketResult = - while ws.pending.len > 0: + while ws.pendingOffset < ws.pendingLength: if interrupted(purpose): return WebSocketResult(kind: wsInterrupted) if getMonoTime() >= deadline: return WebSocketResult(kind: wsDeadline) var sent: csize_t acquire(ws.ioLock) var code: cint withBlockedPipe: - code = curlSend(ws.handle, ws.pending[0].addr, - csize_t(ws.pending.len), sent.addr) + code = curlSend(ws.handle, ws.pending[ws.pendingOffset].addr, + csize_t(ws.pendingLength - ws.pendingOffset), sent.addr) release(ws.ioLock) if code == 0: - if sent > 0: ws.pending.delete(0 .. int(sent) - 1) + ws.pendingOffset += int(sent) elif code != CurlAgain: return WebSocketResult(kind: wsFailure, error: "WebSocket send failed") - if ws.pending.len > 0: + if ws.pendingOffset < ws.pendingLength: let ready = ws.waitSocket(POLLOUT, deadline, purpose) if ready.kind != wsReady: return ready + c_free(ws.pending) + ws.pending = nil + ws.pendingLength = 0 + ws.pendingOffset = 0 if ws.pongState.load() == psQueued: ws.pongState.store(psSent) WebSocketResult(kind: wsReady) @@ -173,7 +188,7 @@ proc sendFrame(ws: NativeWebSocket, opcode: int, data: string, for value in mask: frame.add char(value) for index, value in data: frame.add char(ord(value) xor int(mask[index mod 4])) - ws.pending = move(frame) + ws.installPending(frame) if opcode == 10: ws.pongState.store(psQueued) return ws.flushPending(deadline, purpose) finally: @@ -280,6 +295,10 @@ proc closeNativeWebSocket*(ws: NativeWebSocket) = doAssert not ws.closed, "Native WebSocket already closed" ws.closed = true easy_cleanup(ws.handle) + c_free(ws.pending) + ws.pending = nil + ws.pendingLength = 0 + ws.pendingOffset = 0 deinitLock(ws.ioLock) deinitLock(ws.sendLock) @@ -334,9 +353,10 @@ proc connectNativeWebSocket*(url: string, deadline: MonoTime, if parsed.port.len > 0: host.add ":" & parsed.port var path = if parsed.path.len == 0: "/" else: parsed.path if parsed.query.len > 0: path.add "?" & parsed.query - ws.pending = "GET " & path & " HTTP/1.1\r\nHost: " & host & + let upgradeRequest = "GET " & path & " HTTP/1.1\r\nHost: " & host & "\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Key: " & key & "\r\nSec-WebSocket-Version: 13\r\n\r\n" + ws.installPending(upgradeRequest) let sent = ws.flushPending(deadline, ipPlayer) if sent.kind != wsReady: return WebSocketConnection(kind: sent.kind, error: sent.error) var endHeaders = ws.incoming.find("\r\n\r\n") diff --git a/tests/support/native_websocket_probe.nim b/tests/support/native_websocket_probe.nim index a39be02..a16fed1 100644 --- a/tests/support/native_websocket_probe.nim +++ b/tests/support/native_websocket_probe.nim @@ -1,9 +1,9 @@ import std/[json, monotimes, os, strutils, times] import libcurl import bitworld/[native_stop, native_websocket] -type Sender = tuple[socket: ptr NativeWebSocket, status: ptr WebSocketKind, count, bytes: int] +type Sender = tuple[socket: ptr NativeWebSocket, status: ptr WebSocketKind, count, bytes, deadlineMs: int] proc sendMany(argument: Sender) {.thread.} = - let deadline = getMonoTime() + initDuration(seconds = 2) + let deadline = getMonoTime() + initDuration(milliseconds = argument.deadlineMs) for index in 0 ..< argument.count: let sent = argument.socket[].sendNativeText((if argument.bytes > 0: repeat('x', argument.bytes) else: $index), deadline) argument.status[] = sent.kind @@ -13,7 +13,8 @@ proc main() = let mode = paramStr(2) let started = getMonoTime() let connected = connectNativeWebSocket(paramStr(1), started + initDuration(milliseconds = 600), 16 * 1024 * 1024) - echo $(%*{"connect": $connected.kind, "error": connected.error, "curl_version": $libcurl.version()}) + echo $(%*{"connect": $connected.kind, "error": connected.error, "curl_version": $libcurl.version(), + "elapsed_ms": (getMonoTime() - started).inMilliseconds}) if connected.kind != wsReady: quit(0) let socket = connected.socket defer: closeNativeWebSocket(socket) @@ -40,7 +41,7 @@ proc main() = var owned = socket var status: WebSocketKind var thread: Thread[Sender] - createThread(thread, sendMany, (owned.addr, status.addr, 100, 0)) + createThread(thread, sendMany, (owned.addr, status.addr, 100, 0, 2000)) var count = 0 let deadline = getMonoTime() + initDuration(seconds = 2) for index in 0 ..< 100: @@ -49,11 +50,35 @@ proc main() = inc count joinThread(thread) echo $(%*{"kind": $status, "messages": count}) + elif mode == "alternating": + var owned = socket + var status: WebSocketKind + var thread: Thread[Sender] + let deadline = getMonoTime() + initDuration(seconds = 4) + for index in 0 ..< 20: + createThread(thread, sendMany, (owned.addr, status.addr, 1, 128, 2000)) + joinThread(thread) + doAssert status == wsReady + # The sender's thread heap is gone before the main owner sends/pongs. + doAssert socket.sendCleanupText("stopped", deadline).kind == wsReady + let received = socket.receiveCleanupText(deadline) + doAssert received.kind == wsMessage and received.data == "received" + echo $(%*{"kind": $status, "messages": 20}) + elif mode == "partial-exited-worker": + var owned = socket + var status: WebSocketKind + var thread: Thread[Sender] + createThread(thread, sendMany, (owned.addr, status.addr, 1, 8 * 1024 * 1024, 80)) + joinThread(thread) + doAssert status == wsDeadline + let deadline = getMonoTime() + initDuration(seconds = 4) + report(socket.sendCleanupText("stopped", deadline)) + report(socket.receiveCleanupText(deadline)) elif mode == "pong-timeout": var owned = socket var status: WebSocketKind var thread: Thread[Sender] - createThread(thread, sendMany, (owned.addr, status.addr, 1, 8 * 1024 * 1024)) + createThread(thread, sendMany, (owned.addr, status.addr, 1, 8 * 1024 * 1024, 2000)) report(socket.receiveNativeText(getMonoTime() + initDuration(milliseconds = 250))) joinThread(thread) doAssert status == wsReady diff --git a/tests/test_native_websocket.py b/tests/test_native_websocket.py index 8cd8d09..b1e751e 100644 --- a/tests/test_native_websocket.py +++ b/tests/test_native_websocket.py @@ -162,9 +162,17 @@ def echo_many(connection): connection.sendall(frame(str(index).encode())) rows = run_case('concurrent-reader-writer', 'concurrent', echo_many) assert rows[1]['kind'] == 'wsReady' and rows[1]['messages'] == 100 -started = time.monotonic() +def alternate_owners(connection): + for index in range(20): + assert client_frame(connection) == (1, b'x' * 128) + assert client_frame(connection) == (1, b'stopped') + connection.sendall(frame(b'joined-ping', opcode=9)) + assert client_frame(connection) == (10, b'joined-ping') + connection.sendall(frame(b'received')) +rows = run_case('alternating-exited-worker-main-stop-pong', 'alternating', alternate_owners) +assert rows[1]['kind'] == 'wsReady' and rows[1]['messages'] == 20 rows = run_case('upgrade-absolute-deadline', 'normal', lambda connection: None, stall=True) -assert rows[0]['connect'] == 'wsDeadline' and time.monotonic() - started < 1 +assert rows[0]['connect'] == 'wsDeadline' and rows[0]['elapsed_ms'] < 1000 def pong_after_timeout(connection): prefix = exact(connection, 1) @@ -184,6 +192,8 @@ def partial_send(connection): rows = run_case('partial-send-flush-before-stop', 'partial-send', partial_send) assert rows[1]['kind'] == 'wsDeadline' and rows[2]['kind'] == 'wsReady' assert rows[3]['data'] == 'evidence_received' +rows = run_case('partial-exited-worker-main-cleanup', 'partial-exited-worker', partial_send) +assert rows[1]['kind'] == 'wsReady' and rows[2]['data'] == 'evidence_received' with tempfile.TemporaryDirectory(prefix='native-websocket-tls-', dir=Path(PROBE).parent) as directory: cert = Path(directory) / 'cert.pem' key = Path(directory) / 'key.pem'