Telethon/telethon/update_state.py

import logging
import threading
from collections import deque
from datetime import datetime
from threading import RLock, Event, Thread

from .tl import types as tl


class UpdateState:
    """Used to hold the current state of processed updates.
       To retrieve an update, .poll() should be called.
    """
    def __init__(self, polling):
        self._polling = polling
        self._workers = 4
        self._worker_threads = []

        self.handlers = []
        self._updates_lock = RLock()
        self._updates_available = Event()
        self._updates = deque()

        self._logger = logging.getLogger(__name__)

        # https://core.telegram.org/api/updates
        self._state = tl.updates.State(0, 0, datetime.now(), 0, 0)

        # TODO Rename "polling" to some other variable
        # that signifies "running background threads".
        if polling:
            self._setup_workers()

    def can_poll(self):
        """Returns True if a call to .poll() won't lock"""
        return self._updates_available.is_set()

    def poll(self):
        """Polls an update or blocks until an update object is available"""
        if not self._polling:
            raise ValueError('Updates are not being polled hence not saved.')

        self._updates_available.wait()
        with self._updates_lock:
            if not self._updates_available.is_set():
                return

            update = self._updates.popleft()
            if not self._updates:
                self._updates_available.clear()

        if isinstance(update, Exception):
            raise update  # Some error was set through .set_error()

        return update

    # TODO How should this be handled with background worker threads?
    def get_polling(self):
        return self._polling

    def set_polling(self, polling):
        self._polling = polling
        if polling:
            self._setup_workers()
        else:
            with self._updates_lock:
                self._updates.clear()
            self._stop_workers()

    polling = property(fget=get_polling, fset=set_polling)

    def get_workers(self):
        return self._workers

    def set_workers(self, n):
        self._stop_workers()
        self._workers = n
        self._setup_workers()

    workers = property(fget=get_workers, fset=set_workers)

    def _stop_workers(self):
        """Raises "StopIterationException" on the worker threads to stop them,
           and also clears all of them off the list
        """
        if self._worker_threads:
            pass

        self.set_error(StopIteration())
        for t in self._worker_threads:
            t.join()

        self._worker_threads.clear()

    def _setup_workers(self):
        if self._worker_threads:
            # There already are workers
            return

        for i in range(self._workers):
            thread = Thread(
                target=UpdateState._worker_loop,
                name='UpdateWorker{}'.format(i),
                daemon=True,
                args=(self, i)
            )
            self._worker_threads.append(thread)
            thread.start()

    def _worker_loop(self, wid):
        while True:
            try:
                update = self.poll()
                # TODO Maybe people can add different handlers per update type
                for handler in self.handlers:
                    handler(update)
            except StopIteration:
                break
            except Exception as e:
                # We don't want to crash a worker thread due to any reason
                self._logger.debug(
                    '[ERROR] Unhandled exception on worker {}'.format(wid), e
                )

    def set_error(self, error):
        """Sets an error, so that the next call to .poll() will raise it.
           Can be (and is) used to pass exceptions between threads.
        """
        with self._updates_lock:
            # Insert at the beginning so the very next poll causes an error
            # TODO Should this reset the pts and such?
            self._updates.appendleft(error)
            self._updates_available.set()

    def check_error(self):
        with self._updates_lock:
            if self._updates and isinstance(self._updates[0], Exception):
                raise self._updates.popleft()

    def process(self, update):
        """Processes an update object. This method is normally called by
           the library itself.
        """
        if not self._polling and not self.handlers:
            return

        with self._updates_lock:
            if isinstance(update, tl.updates.State):
                self._state = update
                return  # Nothing else to be done

            pts = getattr(update, 'pts', self._state.pts)
            if hasattr(update, 'pts') and pts <= self._state.pts:
                return  # We already handled this update

            self._state.pts = pts
            if self._polling:
                self._updates.append(update)
                self._updates_available.set()