feat: optimize db connections in thread (#1601)

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

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,6 +242,7 @@ class CompletionService:
def close_pubsub():
with flask_app.app_context():
try:
user = db.session.merge(detached_user)
sleep_iterations = 0
@ -258,6 +259,8 @@ class CompletionService:
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)