mirror of https://github.com/jumpserver/jumpserver
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
126 lines
3.5 KiB
126 lines
3.5 KiB
import json
|
|
import threading
|
|
import time
|
|
|
|
import redis
|
|
from django.core.cache import cache
|
|
from redis.client import PubSub
|
|
|
|
from common.db.utils import safe_db_connection
|
|
from common.utils import get_logger
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
def get_redis_client(db=0):
|
|
client = cache.client.get_client()
|
|
assert isinstance(client, redis.Redis)
|
|
return client
|
|
|
|
|
|
class RedisPubSub:
|
|
def __init__(self, ch, db=10):
|
|
self.ch = ch
|
|
self.db = db
|
|
self.redis = get_redis_client(db)
|
|
|
|
def subscribe(self, _next, error=None, complete=None):
|
|
ps = self.redis.pubsub()
|
|
ps.subscribe(self.ch)
|
|
sub = Subscription(self, ps)
|
|
sub.keep_handle_msg(_next, error, complete)
|
|
return sub
|
|
|
|
def resubscribe(self, _next, error=None, complete=None):
|
|
self.redis = get_redis_client(self.db)
|
|
self.subscribe(_next, error, complete)
|
|
|
|
def publish(self, data):
|
|
data_json = json.dumps(data)
|
|
self.redis.publish(self.ch, data_json)
|
|
return True
|
|
|
|
|
|
class Subscription:
|
|
def __init__(self, pb: RedisPubSub, sub: PubSub):
|
|
self.pb = pb
|
|
self.ch = pb.ch
|
|
self.sub = sub
|
|
self.unsubscribed = False
|
|
|
|
def _handle_msg(self, _next, error, complete):
|
|
"""
|
|
handle arg is the pub published
|
|
:param _next: next msg handler
|
|
:param error: error msg handler
|
|
:param complete: complete msg handler
|
|
:return:
|
|
"""
|
|
msgs = self.sub.listen()
|
|
|
|
if error is None:
|
|
error = lambda m, i: None
|
|
|
|
if complete is None:
|
|
complete = lambda: None
|
|
|
|
try:
|
|
for msg in msgs:
|
|
if msg["type"] != "message":
|
|
continue
|
|
item = None
|
|
try:
|
|
item_json = msg['data'].decode()
|
|
item = json.loads(item_json)
|
|
|
|
with safe_db_connection():
|
|
_next(item)
|
|
except Exception as e:
|
|
error(msg, item)
|
|
logger.error('Subscribe handler handle msg error: {}'.format(e))
|
|
except Exception as e:
|
|
if self.unsubscribed:
|
|
logger.debug('Subscription unsubscribed')
|
|
else:
|
|
logger.error('Consume msg error: {}'.format(e))
|
|
self.retry(_next, error, complete)
|
|
return
|
|
|
|
try:
|
|
complete()
|
|
except Exception as e:
|
|
logger.error('Complete subscribe error: {}'.format(e))
|
|
pass
|
|
|
|
try:
|
|
self.unsubscribe()
|
|
except Exception as e:
|
|
logger.error("Redis observer close error: {}".format(e))
|
|
|
|
def keep_handle_msg(self, _next, error, complete):
|
|
t = threading.Thread(target=self._handle_msg, args=(_next, error, complete))
|
|
t.daemon = True
|
|
t.start()
|
|
return t
|
|
|
|
def unsubscribe(self):
|
|
self.unsubscribed = True
|
|
try:
|
|
self.sub.close()
|
|
except Exception as e:
|
|
logger.error('Unsubscribe msg error: {}'.format(e))
|
|
|
|
def retry(self, _next, error, complete):
|
|
logger.info('Retry subscribe channel: {}'.format(self.ch))
|
|
times = 0
|
|
|
|
while True:
|
|
try:
|
|
self.unsubscribe()
|
|
self.pb.resubscribe(_next, error, complete)
|
|
break
|
|
except Exception as e:
|
|
logger.error('Retry #{} {} subscribe channel error: {}'.format(times, self.ch, e))
|
|
times += 1
|
|
time.sleep(times * 2)
|