feat: torbox speed improvement + torbox proxy stream fix

g0ldyy · Nov 26, 2024 · 30ecbce · 30ecbce
1 parent 90ce1ad
commit 30ecbce
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 16 deletions.
diff --git a/comet/api/stream.py b/comet/api/stream.py
@@ -213,9 +213,7 @@ async def stream(
                 if "searched" in hash:
                     continue
 
-                all_sorted_ranked_files[hash] = orjson.loads(
-                    result["data"]
-                )
+                all_sorted_ranked_files[hash] = orjson.loads(result["data"])
 
         if len(all_sorted_ranked_files) != 0 and set(indexers).issubset(trackers_found):
             debrid_extension = get_debrid_extension(
@@ -662,22 +660,30 @@ async def close(self):
             range_header = request.headers.get("range", "bytes=0-")
 
             try:
-                response = await session.head(
-                    download_link, headers={"Range": range_header}
-                )
+                if config["debridService"] != "torbox":
+                    response = await session.head(
+                        download_link, headers={"Range": range_header}
+                    )
+                else:
+                    response = await session.get(
+                        download_link, headers={"Range": range_header}
+                    )
             except aiohttp.ClientResponseError as e:
                 if e.status == 503 and config["debridService"] == "alldebrid":
-                        proxy = (
-                            settings.DEBRID_PROXY_URL
-                        ) # proxy is not needed to proxy realdebrid stream
+                    proxy = (
+                        settings.DEBRID_PROXY_URL
+                    )  # proxy is needed only to proxy alldebrid streams
 
-                        response = await session.head(
-                            download_link, headers={"Range": range_header}, proxy=proxy
-                        )
+                    response = await session.head(
+                        download_link, headers={"Range": range_header}, proxy=proxy
+                    )
                 else:
-                    raise
+                    logger.warning(f"Exception while proxying {download_link}: {e}")
+                    return
 
-            if response.status == 206:
+            if response.status == 206 or (
+                response.status == 200 and config["debridService"] == "torbox"
+            ):
                 id = str(uuid.uuid4())
                 await database.execute(
                     f"INSERT  {'OR IGNORE ' if settings.DATABASE_TYPE == 'sqlite' else ''}INTO active_connections (id, ip, content, timestamp) VALUES (:id, :ip, :content, :timestamp){' ON CONFLICT DO NOTHING' if settings.DATABASE_TYPE == 'postgresql' else ''}",

diff --git a/comet/debrid/torbox.py b/comet/debrid/torbox.py
@@ -43,7 +43,7 @@ async def get_instant(self, chunk: list):
     async def get_files(
         self, torrent_hashes: list, type: str, season: str, episode: str, kitsu: bool
     ):
-        chunk_size = 100
+        chunk_size = 50
         chunks = [
             torrent_hashes[i : i + chunk_size]
             for i in range(0, len(torrent_hashes), chunk_size)

diff --git a/comet/utils/general.py b/comet/utils/general.py
@@ -716,7 +716,9 @@ async def add_torrent_to_cache(
     for indexer in indexers:
         hash = f"searched-{indexer}-{name}-{season}-{episode}"
 
-        searched = copy.deepcopy(sorted_ranked_files[list(sorted_ranked_files.keys())[0]])
+        searched = copy.deepcopy(
+            sorted_ranked_files[list(sorted_ranked_files.keys())[0]]
+        )
         searched["infohash"] = hash
         searched["data"]["tracker"] = indexer