From f78492f4ea9ca9708351aa7d28f7c5eeb6ae0854 Mon Sep 17 00:00:00 2001
From: xicilion <lion@9465.net>
Date: Thu, 17 May 2018 01:08:12 +0800
Subject: [PATCH] ws, bugfix: WebSocket object in the closure was incorrectly
 released.

---
 fibjs/include/WebSocket.h                |  2 +
 fibjs/src/websocket/WebSocket.cpp        |  6 +-
 fibjs/src/websocket/WebSocketHandler.cpp |  2 -
 test/ws_test.js                          | 76 ++++++++++++++++++++++++
 4 files changed, 83 insertions(+), 3 deletions(-)

diff --git a/fibjs/include/WebSocket.h b/fibjs/include/WebSocket.h
index d9d262ec58..de5218a962 100644
--- a/fibjs/include/WebSocket.h
+++ b/fibjs/include/WebSocket.h
@@ -110,6 +110,8 @@ class WebSocket : public WebSocket_base {
 
     int32_t m_code;
     exlib::string m_reason;
+
+    obj_ptr<ValueHolder> m_holder;
 };
 
 } /* namespace fibjs */
diff --git a/fibjs/src/websocket/WebSocket.cpp b/fibjs/src/websocket/WebSocket.cpp
index 77a75719ff..21d56c8adf 100644
--- a/fibjs/src/websocket/WebSocket.cpp
+++ b/fibjs/src/websocket/WebSocket.cpp
@@ -311,6 +311,7 @@ result_t WebSocket_base::_new(exlib::string url, exlib::string protocol, exlib::
 
         virtual int32_t error(int32_t v)
         {
+            m_this->m_holder.Release();
             m_this->isolate_unref();
             m_this->endConnect(1002, "");
             return v;
@@ -325,7 +326,7 @@ result_t WebSocket_base::_new(exlib::string url, exlib::string protocol, exlib::
     };
 
     obj_ptr<WebSocket> sock = new WebSocket(url, protocol, origin);
-    sock->wrap(This);
+    sock->m_holder = new ValueHolder(sock->wrap(This));
 
     (new asyncConnect(sock, sock->holder()))->apost(0);
 
@@ -342,12 +343,15 @@ void WebSocket::startRecv(Isolate* isolate)
             : AsyncState(NULL)
             , m_this(pThis)
         {
+            if (m_this->m_holder == NULL)
+                m_this->m_holder = new ValueHolder(m_this->wrap());
             m_isolate = isolate;
             set(recv);
         }
 
         ~asyncRead()
         {
+            m_this->m_holder.Release();
             m_this->isolate_unref();
         }
 
diff --git a/fibjs/src/websocket/WebSocketHandler.cpp b/fibjs/src/websocket/WebSocketHandler.cpp
index 043dcb6db7..4a55e6b235 100644
--- a/fibjs/src/websocket/WebSocketHandler.cpp
+++ b/fibjs/src/websocket/WebSocketHandler.cpp
@@ -134,8 +134,6 @@ result_t WebSocketHandler::invoke(object_base* v, obj_ptr<Handler_base>& retVal,
             vs[1] = pThis->m_httpreq;
             pHandler->_emit("accept", vs, 2);
 
-            // sock->startRecv();
-
             return CALL_E_PENDDING;
         }
 
diff --git a/test/ws_test.js b/test/ws_test.js
index 6b313b99ea..eb30a9f753 100644
--- a/test/ws_test.js
+++ b/test/ws_test.js
@@ -6,6 +6,7 @@ var io = require('io');
 var http = require('http');
 var mq = require('mq');
 var coroutine = require('coroutine');
+var test_util = require('./test_util');
 
 var base_port = coroutine.vmid * 10000;
 
@@ -763,6 +764,81 @@ describe('ws', () => {
             });
         });
 
+        describe('gc', () => {
+            it("gc on close", () => {
+                var t = false;
+                GC();
+
+                var no1 = test_util.countObject('WebSocket');
+                var httpd = new http.Server(8816 + base_port, {
+                    "/ws": ws.upgrade((s, req) => {
+                        s.onmessage = e => {}
+                    })
+                });
+
+                ss.push(httpd.socket);
+                httpd.run(() => {});
+
+                assert.equal(test_util.countObject('WebSocket'), no1);
+
+                var s = new ws.Socket("ws://127.0.0.1:" + (8816 + base_port) + "/ws", "test");
+
+                s.onopen = () => {
+                    t = true;
+                };
+
+                s.onmessage = e => {}
+
+                for (var i = 0; i < 1000 && !t; i++)
+                    coroutine.sleep(1);
+
+                assert.equal(test_util.countObject('WebSocket'), no1 + 2);
+
+                s.close();
+                s = undefined;
+                GC();
+                coroutine.sleep(10);
+                GC();
+                assert.equal(test_util.countObject('WebSocket'), no1);
+            });
+
+            it("not gc in closure", () => {
+                var t = false;
+                GC();
+
+                var no1 = test_util.countObject('WebSocket');
+                var httpd = new http.Server(8817 + base_port, {
+                    "/ws": ws.upgrade((s, req) => {
+                        ss.push(s);
+                        s.send(new Date());
+                    })
+                });
+
+                ss.push(httpd.socket);
+                httpd.run(() => {});
+
+                GC();
+                assert.equal(test_util.countObject('WebSocket'), no1);
+
+                function test() {
+                    var s = new ws.Socket("ws://127.0.0.1:" + (8817 + base_port) + "/ws", "test");
+                    s.onmessage = e => {
+                        t = true;
+                        s.close();
+                    }
+                }
+
+                test();
+                GC();
+
+                coroutine.sleep(10);
+                assert.equal(test_util.countObject('WebSocket'), no1 + 2);
+                assert.isTrue(t);
+
+                GC();
+                assert.equal(test_util.countObject('WebSocket'), no1 + 1);
+            });
+        });
     });
 });