From 56b8306ae7be9eb47241d3335c3465d2ddae1b28 Mon Sep 17 00:00:00 2001
From: Dan Holtby <djholtby@gmail.com>
Date: Tue, 13 Jul 2021 14:18:09 -0400
Subject: [PATCH] fix ws-conn will to work as intended

* changes ws-conn-read-thread to hold a reference to ws-conn weakly, so that
  the ws-conn can be garbage collected even if it is still active
* the will closes the connection if it is still closed (so the client will
  not see a broken-pipe error, but proper websocket disconnect)
* the thread stops if the ws-conn is collected (break is not needed in the will)
---
 net/rfc6455/conn-api.rkt | 94 ++++++++++++++++++++++++----------------
 1 file changed, 57 insertions(+), 37 deletions(-)

diff --git a/net/rfc6455/conn-api.rkt b/net/rfc6455/conn-api.rkt
index cdf40d9..3f897ac 100644
--- a/net/rfc6455/conn-api.rkt
+++ b/net/rfc6455/conn-api.rkt
@@ -133,46 +133,66 @@
 (define ws-executor (make-will-executor))
 (void (thread (lambda () (let loop () (will-execute ws-executor) (loop)))))
 
+;; if ws-conn is gc'd while still active, closes it gracefully
+(define (ws-conn-will ws-conn)
+  (unless (ws-conn-closed? ws-conn)
+    (ws-close! ws-conn)))
+
+;; sets maybe-conn's read-thread-status to e, if maybe-conn is not #f
+;; : (U #f ws-conn-base) Any -> Void
+(define (try-set-ws-conn-base-read-thread-status! maybe-conn e)
+  (when maybe-conn
+    (set-ws-conn-base-read-thread-status! maybe-conn e)))
+
+
 (define (ws-read-thread)
   (thread
    (lambda ()
-     (define thd (current-thread))
-     (match-define (list 'start! ws-conn) (thread-receive))
-     (will-register ws-executor ws-conn (lambda (_) (break-thread thd)))
+     (define ws-conn/weak
+       (make-weak-box
+        (match (thread-receive)
+          [(list 'start! ws-conn) ws-conn])))
+     (will-register ws-executor (weak-box-value ws-conn/weak) ws-conn-will)
+                    
 
      (with-handlers [((lambda (e) #t)
-                      (lambda (e) (set-ws-conn-base-read-thread-status! ws-conn e)))]
+                      (lambda (e) (try-set-ws-conn-base-read-thread-status!
+                                   (weak-box-value ws-conn/weak) e)))]
        (let loop ((backlog #f))
-         (match (thread-receive)
-           [(list 'single payload-type nack-evt ch)
-            (define (deliver raw-item auto-conv)
-              (define item ((or (match payload-type
-                                  ['text bytes->string/utf-8]
-                                  ['binary values]
-                                  ['auto #f])
-                                auto-conv)
-                            raw-item))
-              (sync (handle-evt nack-evt
-                                (lambda (_) (loop (list raw-item auto-conv))))
-                    (handle-evt (channel-put-evt ch item)
-                                (lambda (_) (loop #f)))))
-            (match backlog
-              [#f
-               (sync (handle-evt nack-evt
-                                 (lambda (_) (loop #f)))
-                     (handle-evt (ws-conn-base-ip ws-conn)
-                                 (lambda (_)
-                                   (match (ws-recv** ws-conn)
-                                     [(? eof-object?) (void)] ;; terminate!
-                                     [(list raw-item auto-conv) (deliver raw-item auto-conv)]))))]
-              [(list raw-item auto-conv)
-               (deliver raw-item auto-conv)])]
-           [(list 'stream o)
-            (match backlog
-              [#f
-               (ws-stream** ws-conn o)
-               (loop #f)]
-              [(list raw-item _auto-conv)
-               (display raw-item o)
-               (close-output-port o)
-               (loop #f)])]))))))
+         (define v (thread-receive))
+         (define ws-conn (weak-box-value ws-conn/weak))
+         (when ws-conn
+           (match v
+             [(list 'single payload-type nack-evt ch)
+              (define (deliver raw-item auto-conv)
+                (define item ((or (match payload-type
+                                    ['text bytes->string/utf-8]
+                                    ['binary values]
+                                    ['auto #f])
+                                  auto-conv)
+                              raw-item))
+                (sync (handle-evt nack-evt
+                                  (lambda (_) (loop (list raw-item auto-conv))))
+                      (handle-evt (channel-put-evt ch item)
+                                  (lambda (_) (loop #f)))))
+              (match backlog
+                [#f
+                 (sync (handle-evt nack-evt
+                                   (lambda (_) (loop #f)))
+                       (handle-evt (ws-conn-base-ip ws-conn)
+                                   (lambda (_)
+                                     (match (ws-recv** ws-conn)
+                                       [(? eof-object?) (void)] ;; terminate!
+                                       [(list raw-item auto-conv) (deliver raw-item auto-conv)]))))]
+                [(list raw-item auto-conv)
+                 (deliver raw-item auto-conv)])]
+             [(list 'stream o)
+              (match backlog
+                [#f
+                 (ws-stream** ws-conn o)
+                 (loop #f)]
+                [(list raw-item _auto-conv)
+                 (display raw-item o)
+                 (close-output-port o)
+                 (loop #f)])])))))))
+