From 96948504652ad00a67bcd1745c649309dd67e019 Mon Sep 17 00:00:00 2001
From: Julien Duponchelle <julien@gns3.net>
Date: Tue, 16 Aug 2016 17:35:43 +0200
Subject: [PATCH] Option for auto close project when no client listen

And fix websocket not correctly detected has closed.

Ref https://github.com/GNS3/gns3-gui/issues/1331
---
 gns3server/controller/notification.py                |  7 +++++++
 .../handlers/api/compute/notification_handler.py     |  5 +++++
 .../handlers/api/controller/project_handler.py       | 12 ++++++++++++
 tests/handlers/api/controller/test_project.py        |  2 ++
 4 files changed, 26 insertions(+)

diff --git a/gns3server/controller/notification.py b/gns3server/controller/notification.py
index f499073b..f89797eb 100644
--- a/gns3server/controller/notification.py
+++ b/gns3server/controller/notification.py
@@ -44,6 +44,13 @@ class Notification:
         yield queue
         self._listeners[project.id].remove(queue)
 
+    def project_has_listeners(self, project):
+        """
+        :param project_id: Project object
+        :returns: True if client listen this project
+        """
+        return project.id in self._listeners and len(self._listeners[project.id]) > 0
+
     def dispatch(self, action, event, compute_id):
         """
         Notification received from compute node. Send it directly
diff --git a/gns3server/handlers/api/compute/notification_handler.py b/gns3server/handlers/api/compute/notification_handler.py
index 19fa64d9..539d0d86 100644
--- a/gns3server/handlers/api/compute/notification_handler.py
+++ b/gns3server/handlers/api/compute/notification_handler.py
@@ -32,11 +32,16 @@ class NotificationHandler:
         ws = WebSocketResponse()
         yield from ws.prepare(request)
 
+        # Process ping / pong and close message
+        asyncio.async(ws.receive())
+
         with notifications.queue() as queue:
             while True:
                 try:
                     notification = yield from queue.get_json(5)
                 except asyncio.futures.CancelledError:
                     break
+                if ws.closed:
+                    break
                 ws.send_str(notification)
         return ws
diff --git a/gns3server/handlers/api/controller/project_handler.py b/gns3server/handlers/api/controller/project_handler.py
index e3ea4b91..91ff62ae 100644
--- a/gns3server/handlers/api/controller/project_handler.py
+++ b/gns3server/handlers/api/controller/project_handler.py
@@ -213,6 +213,9 @@ class ProjectHandler:
                 except asyncio.futures.CancelledError as e:
                     break
 
+        if project.auto_close and not controller.notification.project_has_listeners(project):
+            yield from project.close()
+
     @Route.get(
         r"/projects/{project_id}/notifications/ws",
         description="Receive notifications about projects from a Websocket",
@@ -231,13 +234,22 @@ class ProjectHandler:
         ws = aiohttp.web.WebSocketResponse()
         yield from ws.prepare(request)
 
+        # Process ping / pong and close message
+        asyncio.async(ws.receive())
+
         with controller.notification.queue(project) as queue:
             while True:
                 try:
                     notification = yield from queue.get_json(5)
                 except asyncio.futures.CancelledError as e:
                     break
+                if ws.closed:
+                    break
                 ws.send_str(notification)
+
+        if project.auto_close and not controller.notification.project_has_listeners(project):
+            yield from project.close()
+
         return ws
 
     @Route.get(
diff --git a/tests/handlers/api/controller/test_project.py b/tests/handlers/api/controller/test_project.py
index 23275c5a..d8aecd25 100644
--- a/tests/handlers/api/controller/test_project.py
+++ b/tests/handlers/api/controller/test_project.py
@@ -146,6 +146,7 @@ def test_notification(http_controller, project, controller, loop):
     assert b'"action": "ping"' in response.body
     assert b'"cpu_usage_percent"' in response.body
     assert b'{"action": "node.created", "event": {"a": "b"}}\n' in response.body
+    assert project.status == "opened"
 
 
 def test_notification_invalid_id(http_controller):
@@ -167,6 +168,7 @@ def test_notification_ws(http_controller, controller, project, async_run):
 
     async_run(http_controller.close())
     ws.close()
+    assert project.status == "opened"
 
 
 def test_export(http_controller, tmpdir, loop, project):