Upgrade fetch notify by websocket

pull/289/head
vapao 2021-02-21 01:07:18 +08:00
parent 354d1ae154
commit 702799ce84
6 changed files with 63 additions and 6 deletions

View File

@ -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

View File

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

View File

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

View File

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

View File

@ -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
})

View File

@ -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();
timer = setInterval(() => {
if (ws.readyState === 1) {
ws.send('ping')
} else if (ws.readyState === 3) {
listen()
}
}, 10000)
return () => {
if (interval) clearInterval(interval)
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) {