jumpserver/apps/common/utils/connection.py

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)