Austin/dify
1
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2024-11-16 19:59:50 +08:00
Code Issues Actions 3 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")
PubHandler.pub_error(user, generate_task_id, e)
finally:
db.session.commit()
db.session.remove()
@classmethod
def countdown_and_close(cls, flask_app: Flask, worker_thread, pubsub, detached_user,
@ -242,22 +242,25 @@ class CompletionService:
def close_pubsub():
with flask_app.app_context():
user = db.session.merge(detached_user)
try:
user = db.session.merge(detached_user)
sleep_iterations = 0
while sleep_iterations < timeout and worker_thread.is_alive():
if sleep_iterations > 0 and sleep_iterations % 10 == 0:
PubHandler.ping(user, generate_task_id)
sleep_iterations = 0
while sleep_iterations < timeout and worker_thread.is_alive():
if sleep_iterations > 0 and sleep_iterations % 10 == 0:
PubHandler.ping(user, generate_task_id)
time.sleep(1)
sleep_iterations += 1
time.sleep(1)
sleep_iterations += 1
if worker_thread.is_alive():
PubHandler.stop(user, generate_task_id)
try:
pubsub.close()
except Exception:
pass
if worker_thread.is_alive():
PubHandler.stop(user, generate_task_id)
try:
pubsub.close()
except Exception:
pass
finally:
db.session.remove()
countdown_thread = threading.Thread(target=close_pubsub)
countdown_thread.start()
@ -394,7 +397,7 @@ class CompletionService:
logging.exception(e)
raise
finally:
db.session.commit()
db.session.remove()
try:
pubsub.unsubscribe(generate_channel)
@ -436,7 +439,7 @@ class CompletionService:
logging.exception(e)
raise
finally:
db.session.commit()
db.session.remove()
try:
pubsub.unsubscribe(generate_channel)