jumpserver/apps/common/utils/connection.py

126 lines
3.5 KiB
Python

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)