ASJPush/Service/CustomizedPushService.py

# @Author    : Rocky
# @File      : CustomizedPushService.py
# @Time      : 2023/10/19 15:49
import logging
import threading
import time
from concurrent.futures import ThreadPoolExecutor

from django.db.models.functions import Substr, Length

from AnsjerPush.config import CONFIG_INFO, CONFIG_TEST, CONFIG_CN, XM_PUSH_CHANNEL_ID
from Model.models import DeviceTypeModel, Device_Info, GatewayPush, CountryModel, SysMsgModel, UserEmailSubscriptions, \
    AppDeviceType, UidSetModel
from Service.CommonService import CommonService
from Service.HuaweiPushService.HuaweiPushService import HuaweiPushObject
from Service.PushService import PushObject

CUSTOMIZED_PUSH_LOGGER = logging.getLogger('customized_push')


class CustomizedPushObject:

    @staticmethod
    def query_push_user(device_name, country, register_period):
        """
        查询需要推送的用户id列表
        @param device_name: 设备型号
        @param country: 国家
        @param register_period: 用户注册年限
        @return: uid_id_list
        """
        # 查询云存设备
        if device_name == 'cloud_storage':
            ipc_types = list(AppDeviceType.objects.filter(model=2).values_list('type', flat=True).distinct())
            uid_list = (UidSetModel.objects.annotate(ucode_char=Substr('ucode', Length('ucode') - 3, 1)).
                        filter(device_type__in=ipc_types, ucode_char__in=['4', '5']).values_list('uid', flat=True))
            device_info_qs = Device_Info.objects.filter(UID__in=list(uid_list))
            # 国外服推给指定国家用户
            if CONFIG_INFO not in [CONFIG_TEST, CONFIG_CN]:
                country_name_list = country.split(',')
                country_id_list = CountryModel.objects.filter(country_name__in=country_name_list). \
                    values_list('id', flat=True)
                device_info_qs = device_info_qs.filter(userID__region_country__in=country_id_list)
        else:
            # 设备型号和国家
            device_name_list = device_name.split(',')
            device_type_list = DeviceTypeModel.objects.filter(name__in=device_name_list).values_list('type', flat=True)
            # 测试和国内服推给所有用户
            if CONFIG_INFO in [CONFIG_TEST, CONFIG_CN]:
                device_info_qs = Device_Info.objects.filter(Type__in=device_type_list)
            else:
                country_name_list = country.split(',')
                country_id_list = CountryModel.objects.filter(country_name__in=country_name_list).\
                    values_list('id', flat=True)
                device_info_qs = Device_Info.objects.filter(Type__in=device_type_list,
                                                            userID__region_country__in=country_id_list)

        # 获取时间范围
        now_time = int(time.time())

        index = register_period.find('-')
        n, m = register_period[:index], register_period[index + 1:]

        if m == '':
            # 0-,所有时间
            if n == '0':
                device_info_qs = device_info_qs.values_list('userID_id', flat=True)
            # n-,n年以上
            else:
                # n年前时间戳转时间字符串
                n_years_seconds = int(n) * 365 * 24 * 60 * 60
                n_year_ago_timestamp = now_time - n_years_seconds
                n_year_ago = CommonService.timestamp_to_str(n_year_ago_timestamp)
                # 注册时间越小越早
                device_info_qs = device_info_qs.filter(userID__data_joined__lte=n_year_ago). \
                    values_list('userID_id', flat=True)
        else:
            # n-m年,(如2-3年)
            n_years_seconds, m_years_seconds = int(n) * 365 * 24 * 60 * 60, int(m) * 365 * 24 * 60 * 60
            n_year_ago_timestamp = now_time - n_years_seconds
            m_year_ago_timestamp = now_time - m_years_seconds
            # 时间戳转时间字符串
            n_year_ago = CommonService.timestamp_to_str(n_year_ago_timestamp)  # 2021
            m_year_ago = CommonService.timestamp_to_str(m_year_ago_timestamp)  # 2020
            # 2020 <= 注册时间 <= 2021
            device_info_qs = device_info_qs. \
                filter(userID__data_joined__gte=m_year_ago, userID__data_joined__lte=n_year_ago). \
                values_list('userID_id', flat=True)

        user_id_list = list(set(device_info_qs))

        return user_id_list

    @classmethod
    def push_and_save_sys_msg(cls, **kwargs):
        """
        推送和保存系统消息（优化版：支持百万级用户）
        @param kwargs: 参数
        @return:
        """
        customized_push_id = kwargs['id']
        user_id_list = kwargs['user_id_list']
        title = kwargs['title']
        msg = kwargs['msg']
        link = kwargs['link']
        icon_link = kwargs['icon_link'] if kwargs['icon_link'] != '' else None

        n_time = int(time.time())
        push_kwargs = {
            'n_time': n_time,
            'title': title,
            'msg': msg,
            'icon_link': icon_link
        }

        # 推送
        if kwargs['push_app'] == 'ZosiSmart':
            app_bundle_id_list = ['com.ansjer.zccloud_a', 'com.ansjer.zccloud']
        else:
            app_bundle_id_list = ['com.ansjer.zccloud_ab', 'com.ansjer.customizede']

        try:
            # 使用流式查询，避免一次性加载所有数据到内存
            gateway_push_qs = GatewayPush.objects.filter(
                user_id__in=user_id_list, app_bundle_id__in=app_bundle_id_list). \
                values('user_id', 'app_bundle_id', 'push_type', 'token_val')
            
            if not gateway_push_qs.exists():
                CUSTOMIZED_PUSH_LOGGER.info('customized_push_id:{}没有需要推送的用户'.format(customized_push_id))
                return

            # 批量查询所有相关用户的邮件订阅状态，减少数据库查询
            unsubscribed_user_ids = set(
                UserEmailSubscriptions.objects.filter(
                    user_id__in=user_id_list, push_sub_status=0
                ).values_list('user_id', flat=True)
            )

            sys_msg_list = []
            saved_user_id_set = set()  # 使用 set 提高查找效率
            gateway_push_list = []
            
            batch_size = 1000  # 每批处理1000条数据
            sys_msg_batch_size = 5000  # 系统消息批量插入大小
            
            # 使用 iterator() 流式查询，分批处理
            for gateway_push in gateway_push_qs.iterator(chunk_size=batch_size):
                user_id = gateway_push['user_id']
                
                # 使用 set 的 O(1) 查找，而不是 list 的 O(n) 查找
                if user_id not in saved_user_id_set:
                    saved_user_id_set.add(user_id)
                    sys_msg_list.append(SysMsgModel(
                        userID_id=user_id, title=title, msg=msg, jumpLink=link, 
                        addTime=n_time, updTime=n_time))
                    
                    # 分批插入系统消息，避免单次插入过多
                    if len(sys_msg_list) >= sys_msg_batch_size:
                        SysMsgModel.objects.bulk_create(sys_msg_list, batch_size=sys_msg_batch_size)
                        sys_msg_list = []
                        CUSTOMIZED_PUSH_LOGGER.info(
                            'customized_push_id:{}已保存{}条系统消息'.format(
                                customized_push_id, len(saved_user_id_set)))
                
                # 检查用户是否取消订阅
                if user_id not in unsubscribed_user_ids:
                    gateway_push_list.append(gateway_push)
            
            # 保存剩余的系统消息
            if sys_msg_list:
                SysMsgModel.objects.bulk_create(sys_msg_list, batch_size=sys_msg_batch_size)
            
            CUSTOMIZED_PUSH_LOGGER.info(
                'customized_push_id:{}共保存{}条系统消息，准备推送{}条'.format(
                    customized_push_id, len(saved_user_id_set), len(gateway_push_list)))
            
            # 异步推送消息
            if gateway_push_list:
                pre_push_kwargs = {
                    'push_kwargs': push_kwargs,
                    'gateway_push_list': gateway_push_list,
                    'customized_push_id': customized_push_id
                }
                pre_push_thread = threading.Thread(
                    target=cls.thr_pool_push,
                    kwargs=pre_push_kwargs)
                pre_push_thread.start()
            
            CUSTOMIZED_PUSH_LOGGER.info('customized_push_id:{}推送任务已启动'.format(customized_push_id))
        except Exception as e:
            CUSTOMIZED_PUSH_LOGGER.info('定制化推送或保存数据异常,'
                                        'error_line:{},error_msg:{}'.format(e.__traceback__.tb_lineno, repr(e)))

    @classmethod
    def thr_pool_push(cls, **kwargs):
        """
        线程池推送（优化版：限制并发数）
        """
        push_kwargs = kwargs['push_kwargs']
        gateway_push_list = kwargs['gateway_push_list']
        customized_push_id = kwargs.get('customized_push_id', 'unknown')
        
        CUSTOMIZED_PUSH_LOGGER.info(
            'customized_push_id:{}线程池推送开始，共{}条'.format(customized_push_id, len(gateway_push_list)))
        
        # 限制最大线程数，避免创建过多线程导致系统资源耗尽
        max_workers = 50  # 根据服务器性能调整
        batch_count = 0
        
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            executor.map(
                lambda gateway_push_kwargs: cls.start_push(push_kwargs, gateway_push_kwargs), 
                gateway_push_list)
            batch_count += 1
        
        CUSTOMIZED_PUSH_LOGGER.info(
            'customized_push_id:{}线程池推送完成'.format(customized_push_id))

    @classmethod
    def start_push(cls, push_kwargs, gateway_push_kwargs):
        title = push_kwargs['title']
        n_time = push_kwargs['n_time']
        msg = push_kwargs['msg']
        icon_link = push_kwargs['icon_link']

        push_type = gateway_push_kwargs['push_type']
        user_id = gateway_push_kwargs['user_id']
        app_bundle_id = gateway_push_kwargs['app_bundle_id']
        token_val = gateway_push_kwargs['token_val']

        push_succeed = cls.push_msg(push_type, app_bundle_id, token_val, n_time, title, msg, icon_link)

        push_status = '成功' if push_succeed else '失败'
        CUSTOMIZED_PUSH_LOGGER.info('{}推送{},push_type:{}'.format(user_id, push_status, push_type))

    @staticmethod
    def push_msg(push_type, app_bundle_id, token_val, n_time, title, msg, icon_link):
        push_kwargs = {
            'nickname': '',
            'event_type': 0,
            'app_bundle_id': app_bundle_id,
            'token_val': token_val,
            'msg_title': title,
            'msg_text': msg,
            'n_time': n_time,
        }
        try:
            # ios
            if push_type == 0:
                push_kwargs['launch_image'] = icon_link
                return PushObject.ios_p8_push(**push_kwargs)
            # gcm
            elif push_type == 1:
                if icon_link is None:
                    icon_link = ''
                push_kwargs['image'] = icon_link
                return PushObject.android_fcm_push_v1(**push_kwargs)
            # 极光
            elif push_type == 2:
                return PushObject.android_jpush(**push_kwargs)
            # 华为
            elif push_type == 3:
                push_kwargs['image_url'] = icon_link
                huawei_push_object = HuaweiPushObject()
                return huawei_push_object.send_push_notify_message(**push_kwargs)
            # 小米
            elif push_type == 4:
                push_kwargs['channel_id'] = XM_PUSH_CHANNEL_ID['service_reminder']
                return PushObject.android_xmpush(**push_kwargs)
            # vivo
            elif push_type == 5:
                return PushObject.android_vivopush(**push_kwargs)
            # oppo
            elif push_type == 6:
                push_kwargs['channel_id'] = 'VALUE_ADDED'
                return PushObject.android_oppopush(**push_kwargs)
            # 魅族
            elif push_type == 7:
                return PushObject.android_meizupush(**push_kwargs)
            else:
                return False
        except Exception as e:
            CUSTOMIZED_PUSH_LOGGER.info('定制化推送异常,'
                                        'error_line:{},error_msg:{}'.format(e.__traceback__.tb_lineno, repr(e)))
            return False