Upgrade fetch notify by websocket

spug_api/apps/notify/models.py

@@ -4,6 +4,7 @@
 from django.db import models
 from django.core.cache import cache
 from libs import ModelMixin, human_datetime
+from libs.channel import Channel
 import time
 
 
@@ -31,6 +32,7 @@ class Notify(models.Model, ModelMixin):
         if not with_quiet or time.time() - cache.get('spug:notify_quiet', 0) > 3600:
             cache.set('spug:notify_quiet', time.time())
             cls.objects.create(source=source, title=title, type=type, content=content)
+            Channel.send_notify(title, content)
 
     def __repr__(self):
         return '<Notify %r>' % self.title

spug_api/consumer/consumers.py

@@ -3,6 +3,7 @@
 # Released under the AGPL-3.0 License.
 from channels.generic.websocket import WebsocketConsumer
 from django_redis import get_redis_connection
+from asgiref.sync import async_to_sync
 from apps.host.models import Host
 from threading import Thread
 import json
@@ -88,3 +89,18 @@ class SSHConsumer(WebsocketConsumer):
         self.chan = self.ssh.invoke_shell(term='xterm')
         self.chan.transport.set_keepalive(30)
         Thread(target=self.loop_read).start()
+
+
+class NotifyConsumer(WebsocketConsumer):
+    def connect(self):
+        async_to_sync(self.channel_layer.group_add)('notify', self.channel_name)
+        self.accept()
+
+    def disconnect(self, code):
+        async_to_sync(self.channel_layer.group_discard)('notify', self.channel_name)
+
+    def receive(self, **kwargs):
+        self.send(text_data='pong')
+
+    def notify_message(self, event):
+        self.send(text_data=json.dumps(event))

spug_api/consumer/routing.py

@@ -10,5 +10,6 @@ ws_router = AuthMiddleware(
     URLRouter([
         path('ws/exec/<str:token>/', ExecConsumer),
         path('ws/ssh/<int:id>/', SSHConsumer),
+        path('ws/notify/', NotifyConsumer),
     ])
 )

spug_api/libs/channel.py

@@ -25,3 +25,12 @@ class Channel:
             'pkey': pkey
         }
         async_to_sync(layer.send)('ssh_exec', message)
+
+    @staticmethod
+    def send_notify(title, content):
+        message = {
+            'type': 'notify.message',
+            'title': title,
+            'content': content
+        }
+        async_to_sync(layer.group_send)('notify', message)

spug_api/spug/routing.py

@@ -2,11 +2,12 @@
 # Copyright: (c) <spug.dev@gmail.com>
 # Released under the AGPL-3.0 License.
 from channels.routing import ProtocolTypeRouter, ChannelNameRouter
-from consumer import routing, executors
+from consumer import routing, executors, consumers
 
 application = ProtocolTypeRouter({
     'channel': ChannelNameRouter({
         'ssh_exec': executors.SSHExecutor,
+        'notify_message': consumers.NotifyConsumer,
     }),
     'websocket': routing.ws_router
 })

spug_web/src/layout/Notification.js

@@ -1,11 +1,12 @@
 import React, { useState, useEffect } from 'react';
-import { Menu, List, Dropdown, Badge } from 'antd';
+import { Menu, List, Dropdown, Badge, Button, notification } from 'antd';
 import { CheckOutlined, NotificationOutlined } from '@ant-design/icons';
-import { http } from 'libs';
+import { http, X_TOKEN } from 'libs';
 import moment from 'moment';
 import styles from './layout.module.less';
 
-let interval;
+let ws = {readyState: 3};
+let timer;
 
 export default function () {
   const [loading, setLoading] = useState(false);
@@ -14,10 +15,19 @@ export default function () {
 
   useEffect(() => {
     fetch();
-    interval = setInterval(fetch, 60000);
+    listen();
-    return () => {
+    timer = setInterval(() => {
-      if (interval) clearInterval(interval)
+      if (ws.readyState === 1) {
+        ws.send('ping')
+      } else if (ws.readyState === 3) {
+        listen()
       }
+    }, 10000)
+    return () => {
+      if (timer) clearInterval(timer);
+      if (ws.close) ws.close()
+    }
+    // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [])
 
   function fetch() {
@@ -30,6 +40,24 @@ export default function () {
       .finally(() => setLoading(false))
   }
 
+  function listen() {
+    const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
+    ws = new WebSocket(`${protocol}//${window.location.host}/api/ws/notify/?x-token=${X_TOKEN}`);
+    ws.onopen = () => ws.send('ok');
+    ws.onmessage = e => {
+      if (e.data === 'pong') {
+      } else {
+        fetch();
+        const {title, content} = JSON.parse(e.data);
+        const key = `open${Date.now()}`;
+        const btn = (
+          <Button type="primary" size="small" onClick={() => notification.close(key)}>知道了</Button>
+        );
+        notification.warning({message: title, description: content, btn, key, top: 64, duration: null})
+      }
+    }
+  }
+
   function handleRead(e, item) {
     e.stopPropagation();
     if (reads.indexOf(item.id) === -1) {