Add ability to choose the amount of worker threads for the main session

This commit is contained in:
Dan 2018-01-26 11:41:09 +01:00
parent e542c73966
commit 7234edad5d
2 changed files with 31 additions and 7 deletions

View File

@ -105,7 +105,8 @@ class Client:
phone_code: str or callable = None,
password: str = None,
first_name: str = None,
last_name: str = None):
last_name: str = None,
workers: int = 4):
self.session_name = session_name
self.test_mode = test_mode
@ -115,6 +116,8 @@ class Client:
self.first_name = first_name
self.last_name = last_name
self.workers = workers
self.dc_id = None
self.auth_key = None
self.user_id = None
@ -144,7 +147,14 @@ class Client:
self.load_config()
self.load_session(self.session_name)
self.session = Session(self.dc_id, self.test_mode, self.proxy, self.auth_key, self.config.api_id)
self.session = Session(
self.dc_id,
self.test_mode,
self.proxy,
self.auth_key,
self.config.api_id,
workers=self.workers
)
terms = self.session.start()
@ -243,7 +253,14 @@ class Client:
self.dc_id = e.x
self.auth_key = Auth(self.dc_id, self.test_mode, self.proxy).create()
self.session = Session(self.dc_id, self.test_mode, self.proxy, self.auth_key, self.config.api_id)
self.session = Session(
self.dc_id,
self.test_mode,
self.proxy,
self.auth_key,
self.config.api_id,
workers=self.workers
)
self.session.start()
r = self.send(

View File

@ -60,7 +60,6 @@ class Session:
INITIAL_SALT = 0x616e67656c696361
WORKERS = 4
WAIT_TIMEOUT = 10
MAX_RETRIES = 5
ACKS_THRESHOLD = 8
@ -68,13 +67,21 @@ class Session:
notice_displayed = False
def __init__(self, dc_id: int, test_mode: bool, proxy: type, auth_key: bytes, api_id: str, is_cdn: bool = False):
def __init__(self,
dc_id: int,
test_mode: bool,
proxy: type,
auth_key: bytes,
api_id: str,
is_cdn: bool = False,
workers: int = 2):
if not Session.notice_displayed:
print("Pyrogram v{}, {}".format(__version__, __copyright__))
print("Licensed under the terms of the " + __license__, end="\n\n")
Session.notice_displayed = True
self.is_cdn = is_cdn
self.workers = workers
self.connection = Connection(DataCenter(dc_id, test_mode), proxy)
@ -115,7 +122,7 @@ class Session:
try:
self.connection.connect()
for i in range(self.WORKERS):
for i in range(self.workers):
Thread(target=self.worker, name="Worker#{}".format(i + 1)).start()
Thread(target=self.recv, name="RecvThread").start()
@ -175,7 +182,7 @@ class Session:
self.connection.close()
for i in range(self.WORKERS):
for i in range(self.workers):
self.recv_queue.put(None)
log.debug("Session stopped")