Own pending WebSocket bytes across writer threads

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")

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

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'