cache replyto chat (#69)

* cache replyto chat so no need to replyto every msg * Mod modify cache replied chat logic * Update efb_telegram_master/cache.py Co-Authored-By: Eana Hufwe <ilove@1a23.com> * Update efb_telegram_master/cache.py Co-Authored-By: Eana Hufwe <ilove@1a23.com> * Update efb_telegram_master/cache.py Co-Authored-By: Eana Hufwe <ilove@1a23.com> * Update efb_telegram_master/cache.py Co-Authored-By: Eana Hufwe <ilove@1a23.com> * Update efb_telegram_master/cache.py Co-Authored-By: Eana Hufwe <ilove@1a23.com> * Update efb_telegram_master/cache.py Co-Authored-By: Eana Hufwe <ilove@1a23.com> * Update efb_telegram_master/slave_message.py Co-Authored-By: Eana Hufwe <ilove@1a23.com> * Update efb_telegram_master/cache.py Co-Authored-By: Eana Hufwe <ilove@1a23.com> * fix code review problem
ehForwarderBot · Sep 15, 2019 · e438c7f · e438c7f
1 parent d25c1f5
commit e438c7f
Show file tree

Hide file tree

Showing 4 changed files with 63 additions and 1 deletion.
diff --git a/efb_telegram_master/__init__.py b/efb_telegram_master/__init__.py
@@ -38,6 +38,7 @@
 from .rpc_utils import RPCUtilities
 from .slave_message import SlaveMessageProcessor
 from .utils import ExperimentalFlagsManager, EFBChannelChatIDStr
+from .cache import LocalCache
 
 
 class TelegramChannel(EFBChannel):
@@ -125,6 +126,7 @@ def __init__(self, instance_id: InstanceID = None):
         # Initialize managers
         self.flag: ExperimentalFlagsManager = ExperimentalFlagsManager(self)
         self.db: DatabaseManager = DatabaseManager(self)
+        self.cache: LocalCache = LocalCache()
         self.bot_manager: TelegramBotManager = TelegramBotManager(self)
         self.chat_binding: ChatBindingManager = ChatBindingManager(self)
         self.commands: CommandsManager = CommandsManager(self)

diff --git a/efb_telegram_master/cache.py b/efb_telegram_master/cache.py
@@ -0,0 +1,45 @@
+# coding: utf-8
+# modified from [messud4312]https://my.oschina.net/u/914655/blog/1799159
+
+import collections
+import time
+import weakref
+
+class LocalCache():
+
+    # Wrapping dict as it requires subclassing for weak reference.
+    class Dict(dict):
+        def __del__(self):
+            pass
+
+    def __init__(self, maxlen=20):
+        self.weak = weakref.WeakValueDictionary()
+        self.strong = collections.deque(maxlen=maxlen)
+
+    @staticmethod
+    def now_time():
+        return int(time.time())
+
+    def get(self, key):
+        val = self.weak.get(key, None)
+        if val is not None:
+            expire = val['expire']
+            value = val['value']
+            if self.now_time() > expire:
+                return None
+            else:
+                return value
+        else:
+            return None
+
+    def set(self, key, value, expire=3600):
+        # strong_ref prevent object from being collected by gc.
+        self.weak[key] = strong_ref = LocalCache.Dict({
+            'expire': self.now_time() + expire,
+            'value': value
+        })
+        # Enqueue the element and waiting to be collected by gc once popped.
+        self.strong.append(strong_ref)
+
+    def remove(self, key):
+        return self.weak.pop(key, None)
diff --git a/efb_telegram_master/master_message.py b/efb_telegram_master/master_message.py
@@ -30,6 +30,7 @@
     from . import TelegramChannel
     from .bot_manager import TelegramBotManager
     from .db import DatabaseManager, MsgLog
+    from .cache import LocalCache
 
 
 class MasterMessageProcessor(LocaleMixin):
@@ -60,6 +61,7 @@ def __init__(self, channel: 'TelegramChannel'):
         self.channel: 'TelegramChannel' = channel
         self.bot: 'TelegramBotManager' = channel.bot_manager
         self.db: 'DatabaseManager' = channel.db
+        self.cache: 'LocalCache' = channel.cache
         self.bot.dispatcher.add_handler(MessageHandler(
             (Filters.text | Filters.photo | Filters.sticker | Filters.document |
              Filters.venue | Filters.location | Filters.audio | Filters.voice | Filters.video) &
@@ -112,7 +114,7 @@ def msg(self, update: Update, context: CallbackContext):
         reply_to = bool(getattr(message, "reply_to_message", None))
         private_chat = message.chat.id == message.from_user.id
 
-        if (private_chat or multi_slaves) and not reply_to:
+        if (private_chat or multi_slaves) and not reply_to and not self.cache.get(message.chat.id):
             candidates = self.db.get_recent_slave_chats(message.chat.id) or \
                          self.db.get_chat_assoc(master_uid=utils.chat_id_to_str(self.channel_id, message.chat.id))[:5]
             if candidates:
@@ -170,6 +172,7 @@ def process_telegram_message(self, update: Update, context: CallbackContext,
         reply_to = bool(getattr(message, "reply_to_message", None))
 
         # Process predefined target (slave) chat.
+        cached_dest = self.cache.get(message.chat.id)
         if channel_id and chat_id:
             destination = utils.chat_id_to_str(channel_id, chat_id)
             if target_msg is not None:
@@ -189,10 +192,13 @@ def process_telegram_message(self, update: Update, context: CallbackContext,
                     message.reply_to_message.message_id))
                 if dest_msg:
                     destination = dest_msg.slave_origin_uid
+                    self.cache.set(message.chat.id, destination)
                 else:
                     return self.bot.reply_error(update,
                                                 self._("Message is not found in database. "
                                                        "Please try with another one. (UC03)"))
+            elif cached_dest:
+                destination = cached_dest
             else:
                 return self.bot.reply_error(update,
                                             self._("Please reply to an incoming message. (UC04)"))
@@ -204,10 +210,13 @@ def process_telegram_message(self, update: Update, context: CallbackContext,
                         message.reply_to_message.message_id))
                     if dest_msg:
                         destination = dest_msg.slave_origin_uid
+                        self.cache.set(message.chat.id, destination)
                     else:
                         return self.bot.reply_error(update,
                                                     self._("Message is not found in database. "
                                                            "Please try with another one. (UC05)"))
+                elif cached_dest:
+                    destination = cached_dest
                 else:
                     return self.bot.reply_error(update,
                                                 self._("This group is linked to multiple remote chats. "

diff --git a/efb_telegram_master/slave_message.py b/efb_telegram_master/slave_message.py
@@ -38,6 +38,7 @@
     from . import TelegramChannel
     from .bot_manager import TelegramBotManager
     from .db import DatabaseManager
+    from .cache import LocalCache
 
 OldMsgID = Tuple[TelegramChatID, TelegramMessageID]
 
@@ -51,6 +52,7 @@ def __init__(self, channel: 'TelegramChannel'):
         self.logger: logging.Logger = logging.getLogger(__name__)
         self.flag: utils.ExperimentalFlagsManager = self.channel.flag
         self.db: 'DatabaseManager' = channel.db
+        self.cache: 'LocalCache' = channel.cache
 
     def send_message(self, msg: EFBMsg) -> EFBMsg:
         """
@@ -269,6 +271,10 @@ def get_slave_msg_dest(self, msg: EFBMsg) -> Tuple[str, Optional[TelegramChatID]
         msg_template = self.generate_message_template(msg, tg_chat, multi_slaves)
         self.logger.debug("[%s] Message is sent to Telegram chat %s, with header \"%s\".",
                           xid, tg_dest, msg_template)
+
+        if self.cache.get(tg_dest) != chat_uid:
+            self.cache.remove(tg_dest)
+
         return msg_template, tg_dest
 
     def slave_message_text(self, msg: EFBMsg, tg_dest: TelegramChatID, msg_template: str, reactions: str,