message_push_daemon.py

"""
在`MessageTask`生成`ISOMessage`后，取出`IOSMessage`推送到apns，然后改变`IOSMessage`的状态
"""
import concurrent
import logging
from concurrent.futures import ThreadPoolExecutor

import multiprocessing

from horo import Session
from horo.models import commit
from horo.models.message import IOSMessage, IOSPushError
import time

logger = logging.getLogger('horo.daemon.push_ios_message')


def push_message(message):
    # try:
    message.push()
    # message.status = IOSMessage.STATUS_DELIVERED
    # commit()
    # except IOSPushError:
    #     # push 失败了，状态改成FAIL, FAIL 的是可能要重试的
    #     message.status = IOSMessage.STATUS_FAIL
    #     commit()
    #     Session.rollback()
    #     logger.exception('Fatal! some err occurs when pushing ios message!')


def main():
    """
    尽量跟踪消息推送的状态，SENDING 状态仅作参考， FAIL则是推送失败，考虑重试
    """
    while True:
        messages = IOSMessage.filter_by(
            send_in_background=True, status=IOSMessage.STATUS_UNPROCESSED).offset(0).limit(
            50).all()
        if not messages:
            logger.info('No ios message to push, waiting for 200s till the next round...')
            # close session
            Session.rollback()
            time.sleep(200)
            continue
        logger.info('begin to push messages.')
        # 改为 sending 状态, 如果 push 成功（失败）之后 commit 失败，就一直是这个状态
        for message in messages:
            message.status = IOSMessage.STATUS_SENDING
        commit()
        with ThreadPoolExecutor(max_workers=multiprocessing.cpu_count() * 5) as executor:
            result_future = {executor.submit(push_message, message): message.id for message in messages}
            # 需要在同一个thread commit
            success_ids = set()
            failed_ids = set()
            for future in concurrent.futures.as_completed(result_future):
                message_id = result_future[future]
                try:
                    future.result()
                    success_ids.add(message_id)
                    logger.info('finish push message for {}'.format(message_id))
                except IOSPushError:
                    failed_ids.add(message_id)
                    logger.exception('Fatal! some err occurs when pushing ios message!')
                except Exception:
                    logger.exception('wtf')
            messages = IOSMessage.mget(success_ids | failed_ids)
            for message in messages:
                if message.id in success_ids:
                    message.status = IOSMessage.STATUS_DELIVERED
                else:
                    message.status = IOSMessage.STATUS_FAIL
            commit()

        logger.info('finish pushing messages.')


if __name__ == '__main__':
    main()