AI/dify
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/langgenius/dify.git synced 2025-08-14 23:45:55 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: optimize db connections in thread (#1601)

Browse Source
This commit is contained in:
takatost 2023-11-22 16:55:59 +08:00 committed by GitHub
parent b930716745
commit 1f58f15bff
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 19 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

35
api/services/completion_service.py
View File

@ -232,7 +232,7 @@ class CompletionService:
logging.exception("Unknown Error in completion") logging.exception("Unknown Error in completion")
PubHandler.pub_error(user, generate_task_id, e) PubHandler.pub_error(user, generate_task_id, e)
finally: finally:
db.session.commit() db.session.remove()
@classmethod @classmethod
def countdown_and_close(cls, flask_app: Flask, worker_thread, pubsub, detached_user, def countdown_and_close(cls, flask_app: Flask, worker_thread, pubsub, detached_user,
@ -242,22 +242,25 @@ class CompletionService:
def close_pubsub(): def close_pubsub():
with flask_app.app_context(): with flask_app.app_context():
user = db.session.merge(detached_user) try:
user = db.session.merge(detached_user)
sleep_iterations = 0 sleep_iterations = 0
while sleep_iterations < timeout and worker_thread.is_alive(): while sleep_iterations < timeout and worker_thread.is_alive():
if sleep_iterations > 0 and sleep_iterations % 10 == 0: if sleep_iterations > 0 and sleep_iterations % 10 == 0:
PubHandler.ping(user, generate_task_id) PubHandler.ping(user, generate_task_id)
time.sleep(1) time.sleep(1)
sleep_iterations += 1 sleep_iterations += 1
if worker_thread.is_alive(): if worker_thread.is_alive():
PubHandler.stop(user, generate_task_id) PubHandler.stop(user, generate_task_id)
try: try:
pubsub.close() pubsub.close()
except Exception: except Exception:
pass pass
finally:
db.session.remove()
countdown_thread = threading.Thread(target=close_pubsub) countdown_thread = threading.Thread(target=close_pubsub)
countdown_thread.start() countdown_thread.start()
@ -394,7 +397,7 @@ class CompletionService:
logging.exception(e) logging.exception(e)
raise raise
finally: finally:
db.session.commit() db.session.remove()
try: try:
pubsub.unsubscribe(generate_channel) pubsub.unsubscribe(generate_channel)
@ -436,7 +439,7 @@ class CompletionService:
logging.exception(e) logging.exception(e)
raise raise
finally: finally:
db.session.commit() db.session.remove()
try: try:
pubsub.unsubscribe(generate_channel) pubsub.unsubscribe(generate_channel)