Update main.py

main.py

@@ -1,109 +1,122 @@
 from flask import Flask, request, jsonify, Response
 import random
 import string
-import time
-import threading
 import os
-from collections import defaultdict
+import threading
+import time
+from pathlib import Path
 
 app = Flask(__name__)
+app.config['UPLOAD_FOLDER'] = 'temp_uploads'
+Path(app.config['UPLOAD_FOLDER']).mkdir(exist_ok=True)
 
-# Конфигурация
+# Конфигурация памяти
 MAX_MEMORY = 12 * 1024 * 1024 * 1024  # 12 GB
-CHUNK_SIZE = 5 * 1024 * 1024  # 5 MB
-FILE_STORAGE = "./temp_uploads"
-os.makedirs(FILE_STORAGE, exist_ok=True)
+CHUNK_SIZE = 100 * 1024 * 1024  # 4 MB
+MAX_FILE_SIZE = 20 * 1024 * 1024 * 1024  # 20 GB
+
+transfers = {}
+active_memory = 0
+lock = threading.Lock()
 
-active_transfers = defaultdict(dict)
-transfer_lock = threading.Lock()
+class TransferManager:
+    def __init__(self):
+        self.transfers = {}
+        self.file_locks = {}
+    
+    def create_transfer(self, filename, filesize):
+        if filesize > MAX_FILE_SIZE:
+            raise ValueError("File size exceeds maximum limit")
+            
+        transfer_id = ''.join(random.choices(string.ascii_letters + string.digits, k=16))
+        file_path = os.path.join(app.config['UPLOAD_FOLDER'], transfer_id)
+        
+        with lock:
+            self.transfers[transfer_id] = {
+                'filename': filename,
+                'filesize': filesize,
+                'received': 0,
+                'file_path': file_path,
+                'created_at': time.time(),
+                'status': 'uploading'
+            }
+            self.file_locks[transfer_id] = threading.Lock()
+        
+        return transfer_id
 
-def generate_transfer_id():
-    return ''.join(random.choices(string.ascii_letters + string.digits, k=16))
+manager = TransferManager()
 
 @app.route('/init_upload', methods=['POST'])
 def init_upload():
     filename = request.json.get('filename')
     filesize = request.json.get('filesize')
     
-    transfer_id = generate_transfer_id()
-    filepath = os.path.join(FILE_STORAGE, transfer_id)
-    
-    with transfer_lock:
-        active_transfers[transfer_id] = {
-            'filename': filename,
-            'filesize': filesize,
-            'received': 0,
-            'filepath': filepath,
-            'created_at': time.time()
-        }
-    
-    # Создаем пустой файл
-    with open(filepath, 'wb') as f:
-        f.truncate(filesize)
-    
-    return jsonify({'transfer_id': transfer_id})
+    try:
+        transfer_id = manager.create_transfer(filename, filesize)
+        return jsonify({'transfer_id': transfer_id})
+    except Exception as e:
+        return jsonify({'error': str(e)}), 400
 
 @app.route('/upload/<transfer_id>', methods=['POST'])
 def upload_chunk(transfer_id):
-    if transfer_id not in active_transfers:
+    if transfer_id not in manager.transfers:
         return jsonify({'error': 'Invalid transfer ID'}), 404
     
     chunk = request.data
-    start = int(request.headers.get('Content-Range').split('-')[0])
+    chunk_size = len(chunk)
     
-    filepath = active_transfers[transfer_id]['filepath']
+    with lock:
+        transfer = manager.transfers[transfer_id]
+        if transfer['status'] != 'uploading':
+            return jsonify({'error': 'Transfer completed'}), 400
     
     try:
-        with open(filepath, 'r+b') as f:
-            f.seek(start)
+        with manager.file_locks[transfer_id], open(transfer['file_path'], 'ab') as f:
             f.write(chunk)
+            transfer['received'] += chunk_size
         
-        with transfer_lock:
-            active_transfers[transfer_id]['received'] += len(chunk)
-            
-        return jsonify({'status': 'success', 'received': active_transfers[transfer_id]['received']})
-    
+        if transfer['received'] >= transfer['filesize']:
+            transfer['status'] = 'ready'
+        
+        return jsonify({
+            'received': transfer['received'],
+            'status': transfer['status']
+        })
     except Exception as e:
         return jsonify({'error': str(e)}), 500
 
 @app.route('/download/<transfer_id>', methods=['GET'])
 def download_file(transfer_id):
     def generate():
-        filepath = active_transfers[transfer_id]['filepath']
-        with open(filepath, 'rb') as f:
+        file_path = manager.transfers.get(transfer_id, {}).get('file_path')
+        if not file_path or not os.path.exists(file_path):
+            yield b'File not found'
+            return
+
+        with open(file_path, 'rb') as f:
             while True:
                 chunk = f.read(CHUNK_SIZE)
                 if not chunk:
                     break
                 yield chunk
         
-        # Удаляем файл после отправки
-        os.remove(filepath)
-        del active_transfers[transfer_id]
+        # Автоочистка через 1 час
+        threading.Timer(3600, cleanup_transfer, args=[transfer_id]).start()
     
     return Response(
         generate(),
         mimetype='application/octet-stream',
-        headers={
-            'Content-Disposition': f'attachment; filename="{active_transfers[transfer_id]["filename"]}"',
-            'Content-Length': str(active_transfers[transfer_id]['filesize'])
-        }
+        headers={'Content-Disposition': f'attachment; filename="{manager.transfers[transfer_id]["filename"]}"'}
     )
 
-# Очистка зависших передач
-def cleanup():
-    while True:
-        time.sleep(3600)
-        now = time.time()
-        for tid in list(active_transfers.keys()):
-            if now - active_transfers[tid]['created_at'] > 86400:  # 24 часа
-                try:
-                    os.remove(active_transfers[tid]['filepath'])
-                    del active_transfers[tid]
-                except:
-                    pass
-
-threading.Thread(target=cleanup, daemon=True).start()
+def cleanup_transfer(transfer_id):
+    with lock:
+        if transfer_id in manager.transfers:
+            try:
+                os.remove(manager.transfers[transfer_id]['file_path'])
+            except:
+                pass
+            del manager.transfers[transfer_id]
 
 if __name__ == '__main__':
     app.run(host='0.0.0.0', port=5000, threaded=True)