This commit is contained in:
madokax 2024-01-02 11:53:23 +08:00
parent a30acbd87f
commit a9ff5f169f
4 changed files with 52 additions and 37 deletions

View File

@ -1,37 +1,35 @@
import threading import threading
import traceback import traceback
import tracemalloc
from concurrent.futures import ThreadPoolExecutor from concurrent.futures import ThreadPoolExecutor
from queue import Queue from queue import Queue
from config import config from config import config
from globals import GlobalState from globals import GlobalState
from signup import Signup, Interrupted from signup import Signup, Interrupted
from src.pool_manager import ThreadPoolManager
from verify_email import verify_email from verify_email import verify_email
def main(): def main():
email_worker = threading.Thread(target=verify_email) email_worker = threading.Thread(target=verify_email)
email_worker.start() email_worker.start()
max_threads = config['signupWorkerNum'] max_threads = config['signupWorkerNum']
task_queue = Queue(max_threads)
executor = ThreadPoolExecutor(max_threads)
def worker(q, executor):
while True:
task = q.get()
executor.submit(task)
worker_thread = threading.Thread(target=worker, args=(task_queue, executor)) pm = ThreadPoolManager(max_threads)
worker_thread.start()
def signup(): def signup():
s = Signup() s = Signup()
s.sign_up() s.sign_up()
while True: while True:
if GlobalState.exception : if GlobalState.exception :
raise GlobalState.exception raise GlobalState.exception
task_queue.put(signup) pm.add_task(signup)
if __name__ == '__main__': if __name__ == '__main__':
main() main()

35
src/pool_manager.py Normal file
View File

@ -0,0 +1,35 @@
from threading import Thread, Lock
from queue import Queue
from loguru import logger
class ThreadPoolManager:
def __init__(self, max_threads):
self.max_threads = max_threads
self.tasks_queue = Queue(max_threads)
self.threads = []
self.lock = Lock()
self._initialize_threads()
def _initialize_threads(self):
for _ in range(self.max_threads):
thread = Thread(target=self._worker)
thread.daemon = True
thread.start()
self.threads.append(thread)
def add_task(self, func, *args, **kwargs):
self.tasks_queue.put((func, args, kwargs))
def _worker(self):
while True:
func, args, kwargs = self.tasks_queue.get()
try:
func(*args, **kwargs)
except Exception as e:
logger.error(f"An error occurred: {e}")
self.tasks_queue.task_done()
def wait_completion(self):
self.tasks_queue.join()

View File

@ -164,9 +164,6 @@ def get_webdriver() -> WebDriver:
if get_config_headless(): if get_config_headless():
if os.name == 'nt': if os.name == 'nt':
windows_headless = True windows_headless = True
elif os.name == 'darwin':
# debug
windows_headless = False
else: else:
start_xvfb_display() start_xvfb_display()
# For normal headless mode: # For normal headless mode:

View File

@ -1,14 +1,9 @@
import email import email
import imaplib import imaplib
import re import re
import threading
import time import time
from concurrent.futures import ThreadPoolExecutor
from queue import Queue
from func_timeout import FunctionTimedOut,func_timeout
from func_timeout import FunctionTimedOut, func_timeout
import cloudflare_solver import cloudflare_solver
@ -17,22 +12,11 @@ from loguru import logger
from config import config from config import config
from globals import GlobalState from globals import GlobalState
from signup import Interrupted from signup import Interrupted
from src.pool_manager import ThreadPoolManager
from utils import get_webdriver from utils import get_webdriver
max_threads = config['emailWorkerNum'] max_threads = config['emailWorkerNum']
task_queue = Queue(max_threads) pm = ThreadPoolManager(max_threads)
executor = ThreadPoolExecutor(max_threads)
def worker(q, executor):
while True:
task = q.get()
executor.submit(task)
worker_thread = threading.Thread(target=worker, args=(task_queue, executor))
worker_thread.start()
def click_verify_link(link): def click_verify_link(link):
@ -47,8 +31,8 @@ def click_verify_link(link):
finally: finally:
driver.quit() driver.quit()
def verify_email():
def verify_email():
username = config['emailAddr'] username = config['emailAddr']
password = config['emailPassword'] password = config['emailPassword']
imap_server = config['emailImapServer'] imap_server = config['emailImapServer']
@ -57,13 +41,13 @@ def verify_email():
GlobalState.exception = Interrupted("email config error") GlobalState.exception = Interrupted("email config error")
raise GlobalState.exception raise GlobalState.exception
if emailImapPort: if emailImapPort:
mail = imaplib.IMAP4_SSL(imap_server,port=emailImapPort) mail = imaplib.IMAP4_SSL(imap_server, port=emailImapPort)
else: else:
mail = imaplib.IMAP4_SSL(imap_server) mail = imaplib.IMAP4_SSL(imap_server)
try: try:
mail.login(username, password) mail.login(username, password)
except Exception as e: except Exception as e:
GlobalState.exception = Interrupted("email config error") GlobalState.exception = Interrupted("email config error")
raise GlobalState.exception raise GlobalState.exception
logger.info("start to monitor openai verify email") logger.info("start to monitor openai verify email")
@ -104,14 +88,15 @@ def verify_email():
link = re.search(r'href="(https://mandrillapp.com[^"]+)"', html_content) link = re.search(r'href="(https://mandrillapp.com[^"]+)"', html_content)
if link: if link:
link = link.group(1) link = link.group(1)
def task(): pm.add_task(lambda: click_verify_link(link))
click_verify_link(link)
task_queue.put(task)
try: try:
while True: while True:
check_mail() check_mail()
time.sleep(10) time.sleep(10)
finally: finally:
mail.logout() mail.logout()
if __name__ == '__main__': if __name__ == '__main__':
verify_email() verify_email()