Add some locks to fix #109

This commit is contained in:
Tulir Asokan
2018-04-29 00:50:50 +03:00
parent ae792a7b33
commit 4c26d7e59a
+20
View File
@@ -17,6 +17,7 @@
from io import BytesIO from io import BytesIO
import time import time
import logging import logging
import asyncio
import magic import magic
from sqlalchemy.exc import IntegrityError, InvalidRequestError from sqlalchemy.exc import IntegrityError, InvalidRequestError
@@ -125,6 +126,10 @@ async def transfer_thumbnail_to_matrix(client, intent, thumbnail_loc, video, mim
width=width, height=height) width=width, height=height)
transfer_locks = {}
transfer_locks_lock = asyncio.Lock()
async def transfer_file_to_matrix(db, client, intent, location, thumbnail=None, is_sticker=False): async def transfer_file_to_matrix(db, client, intent, location, thumbnail=None, is_sticker=False):
id = _location_to_id(location) id = _location_to_id(location)
if not id: if not id:
@@ -134,6 +139,21 @@ async def transfer_file_to_matrix(db, client, intent, location, thumbnail=None,
if db_file: if db_file:
return db_file return db_file
async with transfer_locks_lock:
try:
lock = transfer_locks[id]
except KeyError:
lock = asyncio.Lock()
transfer_locks[id] = lock
async with lock:
return await _unlocked_transfer_file_to_matrix(db, client, intent, id, location, thumbnail, is_sticker)
async def _unlocked_transfer_file_to_matrix(db, client, intent, id, location, thumbnail, is_sticker):
db_file = DBTelegramFile.query.get(id)
if db_file:
return db_file
try: try:
file = await client.download_file_bytes(location) file = await client.download_file_bytes(location)
except LocationInvalidError: except LocationInvalidError: