Update app.py

app.py

@@ -29,6 +29,14 @@ os.makedirs(DATA_DIR, exist_ok=True)
 class DownloadRequest(BaseModel):
     url: str
 
+def clean_filename(url: str) -> str:
+    """URL se sahi extension aur filename nikalne ka jugaad"""
+    path = url.split("?")[0]
+    filename = path.split("/")[-1]
+    if not filename or "." not in filename:
+        return "downloaded_file.bin"
+    return filename
+
 async def cleanup_old_files():
     while True:
         try:
@@ -52,13 +60,15 @@ async def startup_event():
 async def download_file(task_id: str, url: str):
     file_path = tasks[task_id]["file_path"]
     try:
-        connector = aiohttp.TCPConnector(limit=10)
-        async with aiohttp.ClientSession(connector=connector, read_bufsize=256 * 1024) as session:
+        # Buffer limit ko strictly small rakha hai taaki RAM leak na ho
+        connector = aiohttp.TCPConnector(limit=5)
+        async with aiohttp.ClientSession(connector=connector, read_bufsize=64 * 1024) as session:
             async with session.get(url, timeout=None) as response:
                 response.raise_for_status()
                 total_size = int(response.headers.get('Content-Length', 0))
                 tasks[task_id]["total_size"] = total_size
                 
+                # Header check for proper filename extraction (.apk, .mp4, etc.)
                 cd = response.headers.get('Content-Disposition')
                 if cd and 'filename=' in cd:
                     fname = re.findall('filename="([^"]+)"', cd)
@@ -66,10 +76,9 @@ async def download_file(task_id: str, url: str):
                         fname = re.findall('filename=([^;]+)', cd)
                     if fname:
                         new_filename = fname[0]
-                        new_file_path = os.path.join(DATA_DIR, f"{task_id}_{new_filename}")
-                        tasks[task_id]["file_path"] = new_file_path
+                        file_path = os.path.join(DATA_DIR, f"{task_id}_{new_filename}")
+                        tasks[task_id]["file_path"] = file_path
                         tasks[task_id]["original_filename"] = new_filename
-                        file_path = new_file_path
 
                 downloaded = 0
                 start_time = time.time()
@@ -78,7 +87,7 @@ async def download_file(task_id: str, url: str):
                 chunk_counter = 0
                 
                 with open(file_path, 'wb') as f:
-                    async for chunk in response.content.iter_chunked(256 * 1024):
+                    async for chunk in response.content.iter_chunked(64 * 1024):
                         if not chunk:
                             break
                         f.write(chunk)
@@ -94,10 +103,10 @@ async def download_file(task_id: str, url: str):
                         
                         f.flush()
                         chunk_counter += 1
-                        if chunk_counter % 20 == 0:
+                        if chunk_counter % 10 == 0:
                             os.fsync(f.fileno())
                             del chunk
-                            gc.collect()
+                            gc.collect() # Strict garbage collection
                             
                 tasks[task_id]["status"] = "completed"
                 tasks[task_id]["speed"] = 0.0
@@ -110,11 +119,7 @@ async def download_file(task_id: str, url: str):
 @app.post("/start_download")
 async def start_download(req: DownloadRequest, background_tasks: BackgroundTasks):
     task_id = str(uuid.uuid4())
-    filename = req.url.split("/")[-1]
-    if "?" in filename:
-        filename = filename.split("?")[0]
-    if not filename or "." not in filename:
-        filename = "downloaded_file.bin"
+    filename = clean_filename(req.url)
         
     tasks[task_id] = {
         "task_id": task_id,
@@ -166,7 +171,7 @@ async def stream(task_id: str, request: Request):
         def file_iterator(start, length):
             with open(file_path, "rb") as f:
                 f.seek(start)
-                chunk_size = 512 * 1024  # Reduced chunk size for memory safety
+                chunk_size = 128 * 1024  # Hugging Face RAM Optimization
                 while length > 0:
                     read_size = min(chunk_size, length)
                     data = f.read(read_size)
@@ -175,6 +180,7 @@ async def stream(task_id: str, request: Request):
                     yield data
                     length -= len(data)
                     del data
+                    gc.collect()
                     
         headers = {
             "Content-Range": f"bytes {byte1}-{byte1+length-1}/{file_size}",
@@ -215,18 +221,18 @@ async def upload_file(file: UploadFile = File(...)):
     original_filename = file.filename if file.filename else "uploaded_file.bin"
     file_path = os.path.join(DATA_DIR, f"{task_id}_{original_filename}")
     
-    # Using chunked reading directly from the file stream to save memory
     with open(file_path, "wb") as buffer:
         while True:
-            chunk = await file.read(256 * 1024)
+            chunk = await file.read(64 * 1024)
             if not chunk:
                 break
             buffer.write(chunk)
             buffer.flush()
             os.fsync(buffer.fileno())
             del chunk
+            gc.collect()
     
-    await file.close() # Strictly close file to free memory
+    await file.close()
     file_size = os.path.getsize(file_path)
     
     tasks[task_id] = {
@@ -246,11 +252,7 @@ async def upload_file(file: UploadFile = File(...)):
 @app.post("/start_gofile_transfer")
 async def start_gofile_transfer(req: DownloadRequest, background_tasks: BackgroundTasks):
     task_id = str(uuid.uuid4())
-    filename = req.url.split("/")[-1]
-    if "?" in filename:
-        filename = filename.split("?")[0]
-    if not filename or "." not in filename:
-        filename = "transfer_file.bin"
+    filename = clean_filename(req.url)
         
     tasks[task_id] = {
         "task_id": task_id,
@@ -267,8 +269,7 @@ async def start_gofile_transfer(req: DownloadRequest, background_tasks: Backgrou
     background_tasks.add_task(process_gofile_transfer, task_id, req.url)
     return {"task_id": task_id}
 
-# Chunked file reader generator to avoid streaming upload memory spikes
-async def file_sender(file_path, chunk_size=256 * 1024):
+async def file_sender(file_path, chunk_size=64 * 1024):
     with open(file_path, 'rb') as f:
         while True:
             chunk = f.read(chunk_size)
@@ -276,13 +277,13 @@ async def file_sender(file_path, chunk_size=256 * 1024):
                 break
             yield chunk
             del chunk
+            gc.collect()
 
 async def process_gofile_transfer(task_id: str, url: str):
     file_path = tasks[task_id]["file_path"]
     try:
-        # 1. Download Setup
-        connector = aiohttp.TCPConnector(limit=10)
-        async with aiohttp.ClientSession(connector=connector, read_bufsize=256 * 1024) as session:
+        connector = aiohttp.TCPConnector(limit=5)
+        async with aiohttp.ClientSession(connector=connector, read_bufsize=64 * 1024) as session:
             async with session.get(url, timeout=None) as response:
                 response.raise_for_status()
                 total_size = int(response.headers.get('Content-Length', 0))
@@ -295,10 +296,9 @@ async def process_gofile_transfer(task_id: str, url: str):
                         fname = re.findall('filename=([^;]+)', cd)
                     if fname:
                         new_filename = fname[0]
-                        new_file_path = os.path.join(DATA_DIR, f"{task_id}_{new_filename}")
-                        tasks[task_id]["file_path"] = new_file_path
+                        file_path = os.path.join(DATA_DIR, f"{task_id}_{new_filename}")
+                        tasks[task_id]["file_path"] = file_path
                         tasks[task_id]["original_filename"] = new_filename
-                        file_path = new_file_path
 
                 downloaded = 0
                 start_time = time.time()
@@ -307,7 +307,7 @@ async def process_gofile_transfer(task_id: str, url: str):
                 chunk_counter = 0
                 
                 with open(file_path, 'wb') as f:
-                    async for chunk in response.content.iter_chunked(256 * 1024):
+                    async for chunk in response.content.iter_chunked(64 * 1024):
                         if not chunk:
                             break
                         f.write(chunk)
@@ -322,12 +322,11 @@ async def process_gofile_transfer(task_id: str, url: str):
                         
                         f.flush()
                         chunk_counter += 1
-                        if chunk_counter % 20 == 0:
+                        if chunk_counter % 10 == 0:
                             os.fsync(f.fileno())
                             del chunk
                             gc.collect()
         
-        # 2. Upload to GoFile (STREAMING JUGAD WITH ZERO RAM IMPRINT)
         tasks[task_id]["status"] = "uploading_to_gofile"
         tasks[task_id]["speed"] = 0.0
         
@@ -338,7 +337,6 @@ async def process_gofile_transfer(task_id: str, url: str):
             
             upload_url = f"https://{server}.gofile.io/uploadFile"
             
-            # Form data wrapping without buffering entire file array into memory
             data = aiohttp.FormData()
             data.add_field('file', file_sender(file_path), filename=tasks[task_id]["original_filename"])
                 
@@ -361,7 +359,7 @@ async def process_gofile_transfer(task_id: str, url: str):
     finally:
         gc.collect()
 
-# ==================== NEW ZIP MANAGEMENT FEATURES ====================
+# ==================== ZIP MANAGEMENT FEATURES ====================
 
 @app.get("/list_zip/{task_id}")
 async def list_zip_contents(task_id: str):
@@ -410,15 +408,13 @@ async def extract_zip_file(task_id: str, background_tasks: BackgroundTasks):
     return {"status": "extraction_started", "task_id": task_id, "target_folder": extract_folder}
 
 def process_zip_extraction(task_id: str, zip_path: str, target_dir: str):
-    """Memory safe extraction processing member by member loop"""
     try:
         with zipfile.ZipFile(zip_path, 'r') as z:
             file_counter = 0
             for member in z.namelist():
                 z.extract(member, path=target_dir)
                 file_counter += 1
-                # Release lock periodically inside loops
-                if file_counter % 50 == 0:
+                if file_counter % 10 == 0:
                     gc.collect()
         tasks[task_id]["status"] = "extracted"
     except Exception as e:
@@ -429,4 +425,4 @@ def process_zip_extraction(task_id: str, zip_path: str, target_dir: str):
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)