Code Updates & Optimisations :Server Upgrade Test

FileStream/server/Functions/downloader.py

@@ -19,80 +19,97 @@ from FileStream.Exceptions import FIleNotFound, InvalidHash
 from FileStream.utils.FileProcessors.custom_ul import  TeleUploader
 
 
-async def media_streamer(request: web.Request, db_id: str, speed:str):
-
-  range_header = request.headers.get("Range", 0)
-  #index = minWORK_LOADS, keyWORK_LOADS.get)
-  #faster_client = MULTI_CLIENTS[index]
-
-  client = await req_client()
-
-  if Telegram.MULTI_CLIENT:
-    logging.info(f"Client {client['index']} is now serving {request.headers.get('X-FORWARDED-FOR',request.remote)}")
-
-  if client['client'] in ACTIVE_CLIENTS:
-    tg_connect = ACTIVE_CLIENTS[client['client']]
-    logging.debug(f"Using cached ByteStreamer object for client {client['index']}")
-
-  else:
-    logging.debug(f"Creating new ByteStreamer object for client {client['index']}")
-    tg_connect = utils.ByteStreamer(client['client'])
-    ACTIVE_CLIENTS[client['client']] = tg_connect
-
-  logging.debug("before calling get_file_properties")
-  file_id = await tg_connect.get_file_properties(db_id, MULTI_CLIENTS)
-  logging.debug("after calling get_file_properties")
-
-  file_size = file_id.file_size
-
-  if range_header:
-    from_bytes, until_bytes = range_header.replace("bytes=", "").split("-")
-    from_bytes = int(from_bytes)
-    until_bytes = int(until_bytes) if until_bytes else file_size - 1
-  else:
-    from_bytes = request.http_range.start or 0
-    until_bytes = (request.http_range.stop or file_size) - 1
-
-  if (until_bytes > file_size) or (from_bytes < 0) or (until_bytes < from_bytes):
-
-    return web.Response(
-      status=416,
-      body="416: Range not satisfiable",
-      headers={"Content-Range": f"bytes */{file_size}"},
-    )
-
-  chunk_size = 4 * 1024 * 1024 if speed == "FAST" else 512 * 1024
-
-  until_bytes = min(until_bytes, file_size - 1)
-
-  offset = from_bytes - (from_bytes % chunk_size)
-  first_part_cut = from_bytes - offset
-  last_part_cut = until_bytes % chunk_size + 1
-
-  req_length = until_bytes - from_bytes + 1
-  part_count = math.ceil(until_bytes / chunk_size) - math.floor(offset / chunk_size)
-
-  body = tg_connect.yield_file(file_id, client['index'], offset, first_part_cut,last_part_cut, part_count, chunk_size)
-
-  mime_type = file_id.mime_type
-  file_name = utils.get_name(file_id)
-  disposition = "attachment"
-
-  if not mime_type:
-    mime_type = mimetypes.guess_type(file_name)[0] or "application/octet-stream"
-
-  # if "video/" in mime_type or "audio/" in mime_type:
-  #     disposition = "inline"
-
-  return web.Response(
-    status=206 if range_header else 200,
-    body=body,
-    headers={
-      "Content-Type": f"{mime_type}",
-      "Content-Range": f"bytes {from_bytes}-{until_bytes}/{file_size}",
-      "Content-Length": str(req_length),
-      "Content-Disposition": f'{disposition}; filename="{file_name}"',
-      "Accept-Ranges": "bytes",
-    },
-  )
-
+async def media_streamer(request: web.Request, db_id: str, speed: str):
+    # Get the Range header from the request, default to 0 if not present
+    range_header = request.headers.get("Range", "0")
+    client = await req_client()
+    # Log client info if multi-client mode
+    if Telegram.MULTI_CLIENT:
+        logging.info(f"Client {client['index']} is now serving {request.headers.get('X-FORWARDED-FOR', request.remote)}")
+
+    # Use an existing ByteStreamer or create a new one
+    tg_connect = ACTIVE_CLIENTS.get(client['client'], None)
+    
+    if tg_connect is None:
+        logging.debug(f"Creating new ByteStreamer object for client {client['index']}")
+        tg_connect = utils.ByteStreamer(client['client'])
+        ACTIVE_CLIENTS[client['client']] = tg_connect
+
+    else:
+        tg_connect.update_last_activity()
+        logging.debug(f"Using cached ByteStreamer object for client {client['index']}")
+
+    try:
+        # Fetch file properties once and use it throughout
+        logging.debug("Fetching file properties")
+        file_id = await tg_connect.get_file_properties(db_id, MULTI_CLIENTS)
+        file_size = file_id.file_size
+
+        # Parse range header efficiently
+        from_bytes, until_bytes = parse_range(range_header, file_size)
+
+        # If range is invalid, return a 416 error
+        if from_bytes is None or until_bytes is None:
+            return web.Response(
+                status=416,
+                body="416: Range not satisfiable",
+                headers={"Content-Range": f"bytes */{file_size}"},
+            )
+
+        # Set chunk size based on speed
+        chunk_size =  1024 * 1024 if speed == "FAST" else 512 * 1024
+        
+
+        # Ensure we don't go past the file size
+        until_bytes = min(until_bytes, file_size - 1)
+
+        # Compute offset and range parts
+        offset, first_part_cut, last_part_cut, part_count = compute_offsets(from_bytes, until_bytes, chunk_size)
+
+        # Request the file chunks
+        body = tg_connect.yield_file(
+            file_id, client['index'], offset, first_part_cut, last_part_cut, part_count, chunk_size
+        )
+
+        # Determine MIME type and filename
+        mime_type = file_id.mime_type or mimetypes.guess_type(file_id.file_name)[0] or "application/octet-stream"
+        file_name = utils.get_name(file_id)
+        disposition = "attachment"
+
+        # Return the response with proper headers and status
+        req_length = until_bytes - from_bytes + 1
+        return web.Response(
+            status=206 if range_header else 200,
+            body=body,
+            headers={
+                "Content-Type": mime_type,
+                "Content-Range": f"bytes {from_bytes}-{until_bytes}/{file_size}",
+                "Content-Length": str(req_length),
+                "Content-Disposition": f'{disposition}; filename="{file_name}"',
+                "Accept-Ranges": "bytes",
+            },
+        )
+    except Exception as e:
+        logging.error(f"Error in media_streamer: {traceback.format_exc()}")
+        raise web.HTTPInternalServerError()  # Re-raise the exception as a server error
+
+
+def parse_range(range_header: str, file_size: int):
+    """Helper function to parse the range header."""
+    try:
+        range_str = range_header.replace("bytes=", "")
+        from_bytes, until_bytes = range_str.split("-")
+        from_bytes = int(from_bytes)
+        until_bytes = int(until_bytes) if until_bytes else file_size - 1
+        return from_bytes, until_bytes
+    except ValueError:
+        return None, None
+
+
+def compute_offsets(from_bytes: int, until_bytes: int, chunk_size: int):
+    """Compute the offsets, cuts, and part counts for file chunking."""
+    offset = from_bytes - (from_bytes % chunk_size)
+    first_part_cut = from_bytes - offset
+    last_part_cut = until_bytes % chunk_size + 1
+    part_count = math.ceil(until_bytes / chunk_size) - math.floor(offset / chunk_size)
+    return offset, first_part_cut, last_part_cut, part_count

FileStream/utils/FileProcessors/custom_dl.py

@@ -23,200 +23,210 @@ from FileStream.Exceptions import FIleNotFound, InvalidHash
 from FileStream.bot import MULTI_CLIENTS, WORK_LOADS, ACTIVE_CLIENTS
 
 
+
 class ByteStreamer:
+    def __init__(self, client: Client):
+        self.clean_timer = 30 * 60  # Cache cleanup timer set to 30 minutes
+        self.client: Client = client
+        self.cached_file_ids: Dict[str, FileId] = {}  # Cache to store file properties by db_id
+        self.last_activity: float = asyncio.get_event_loop().time()  # Track last activity time for the client
+        asyncio.create_task(self.clean_cache())  # Start the cache cleanup task
+
+    def update_last_activity(self):
+        """Update the last activity time to the current time."""
+        self.last_activity = asyncio.get_event_loop().time()
+
+    def get_last_activity(self) -> float:
+        """Get the last activity time of this client."""
+        return self.last_activity
+
+    async def get_file_properties(self, db_id: str, MULTI_CLIENTS) -> FileId:
+        """
+        Returns the properties of a media of a specific message in a FileId class.
+        If the properties are cached, it'll return the cached results.
+        Otherwise, it'll generate the properties from the Message ID and cache them.
+        """
+        if db_id not in self.cached_file_ids:
+            logging.debug("File properties not cached. Generating properties.")
+            await self.generate_file_properties(db_id, MULTI_CLIENTS)  # Generate and cache the file properties
+            logging.debug(f"Cached file properties for file with ID {db_id}")
+        return self.cached_file_ids[db_id]
 
-  def __init__(self, client: Client):
-    self.clean_timer = 30 * 60
-    self.client: Client = client
-    self.cached_file_ids: Dict[str, FileId] = {}
-    asyncio.create_task(self.clean_cache())
-
-  async def get_file_properties(self, db_id: str, MULTI_CLIENTS) -> FileId:
-    """
-        Returns the properties of a media of a specific message in a FIleId class.
-        if the properties are cached, then it'll return the cached results.
-        or it'll generate the properties from the Message ID and cache them.
+    async def generate_file_properties(self, db_id: str, MULTI_CLIENTS) -> FileId:
         """
-    if not db_id in self.cached_file_ids:
-      logging.debug("Before Calling generate_file_properties")
-      await self.generate_file_properties(db_id, MULTI_CLIENTS)
-      logging.debug(f"Cached file properties for file with ID {db_id}")
-    return self.cached_file_ids[db_id]
-
-  async def generate_file_properties(self, db_id: str,
-                                     MULTI_CLIENTS) -> FileId:
-    """
         Generates the properties of a media file on a specific message.
-        returns ths properties in a FIleId class.
+        Returns the properties in a FileId class.
+        """
+        logging.debug("Generating file properties.")
+        file_id = await get_file_ids(self.client, db_id, Message)  # Call the method to get the file properties
+        logging.debug(f"Generated file ID and Unique ID for file with ID {db_id}")
+        self.cached_file_ids[db_id] = file_id  # Cache the file properties
+        logging.debug(f"Cached media file with ID {db_id}")
+        return file_id
+
+    async def generate_media_session(self, client: Client, file_id: FileId) -> Session:
         """
-    logging.debug("Before calling get_file_ids")
-    file_id = await get_file_ids(self.client, db_id, Message)
-    logging.debug(f"Generated file ID and Unique ID for file with ID {db_id}")
-    self.cached_file_ids[db_id] = file_id
-    logging.debug(f"Cached media file with ID {db_id}")
-    return self.cached_file_ids[db_id]
-
-  async def generate_media_session(self, client: Client,
-                                   file_id: FileId) -> Session:
-    """
         Generates the media session for the DC that contains the media file.
         This is required for getting the bytes from Telegram servers.
         """
-
-    media_session = client.media_sessions.get(file_id.dc_id, None)
-
-    if media_session is None:
-      if file_id.dc_id != await client.storage.dc_id():
-        media_session = Session(
-            client,
-            file_id.dc_id,
-            await Auth(client, file_id.dc_id, await
-                       client.storage.test_mode()).create(),
-            await client.storage.test_mode(),
-            is_media=True,
-        )
-        await media_session.start()
-
-        for _ in range(6):
-          exported_auth = await client.invoke(
-              raw.functions.auth.ExportAuthorization(dc_id=file_id.dc_id))
-
-          try:
-            await media_session.invoke(
-                raw.functions.auth.ImportAuthorization(
-                    id=exported_auth.id, bytes=exported_auth.bytes))
-            break
-          except AuthBytesInvalid:
-            logging.debug(
-                f"Invalid authorization bytes for DC {file_id.dc_id}")
-            continue
+        media_session = client.media_sessions.get(file_id.dc_id, None)
+
+        if media_session is None:
+            if file_id.dc_id != await client.storage.dc_id():
+                # Create a new media session if one doesn't exist for this DC ID
+                media_session = Session(
+                    client,
+                    file_id.dc_id,
+                    await Auth(client, file_id.dc_id, await client.storage.test_mode()).create(),
+                    await client.storage.test_mode(),
+                    is_media=True,
+                )
+                await media_session.start()
+
+                # Attempt to import authorization from Telegram's servers
+                for _ in range(6):
+                    exported_auth = await client.invoke(
+                        raw.functions.auth.ExportAuthorization(dc_id=file_id.dc_id))
+
+                    try:
+                        # Import the authorization bytes for the DC
+                        await media_session.invoke(
+                            raw.functions.auth.ImportAuthorization(
+                                id=exported_auth.id, bytes=exported_auth.bytes))
+                        break
+                    except AuthBytesInvalid:
+                        logging.debug(f"Invalid authorization bytes for DC {file_id.dc_id}")
+                        continue
+                else:
+                    await media_session.stop()
+                    raise AuthBytesInvalid
+            else:
+                # Reuse the stored auth key if we're already connected to the correct DC
+                media_session = Session(
+                    client,
+                    file_id.dc_id,
+                    await client.storage.auth_key(),
+                    await client.storage.test_mode(),
+                    is_media=True,
+                )
+                await media_session.start()
+
+            logging.debug(f"Created media session for DC {file_id.dc_id}")
+            client.media_sessions[file_id.dc_id] = media_session  # Cache the media session
         else:
-          await media_session.stop()
-          raise AuthBytesInvalid
-      else:
-        media_session = Session(
-            client,
-            file_id.dc_id,
-            await client.storage.auth_key(),
-            await client.storage.test_mode(),
-            is_media=True,
-        )
-        await media_session.start()
-      logging.debug(f"Created media session for DC {file_id.dc_id}")
-      client.media_sessions[file_id.dc_id] = media_session
-    else:
-      logging.debug(f"Using cached media session for DC {file_id.dc_id}")
-    return media_session
-
-  @staticmethod
-  async def get_location(
-      file_id: FileId
-  ) -> Union[
-      raw.types.InputPhotoFileLocation,
-      raw.types.InputDocumentFileLocation,
-      raw.types.InputPeerPhotoFileLocation,
-  ]:
-    """
-        Returns the file location for the media file.
+            logging.debug(f"Using cached media session for DC {file_id.dc_id}")
+        return media_session
+
+    @staticmethod
+    async def get_location(file_id: FileId) -> Union[
+            raw.types.InputPhotoFileLocation,
+            raw.types.InputDocumentFileLocation,
+            raw.types.InputPeerPhotoFileLocation,
+    ]:
         """
-    file_type = file_id.file_type
-
-    if file_type == FileType.CHAT_PHOTO:
-      if file_id.chat_id > 0:
-        peer = raw.types.InputPeerUser(user_id=file_id.chat_id,
-                                       access_hash=file_id.chat_access_hash)
-      else:
-        if file_id.chat_access_hash == 0:
-          peer = raw.types.InputPeerChat(chat_id=-file_id.chat_id)
+        Returns the file location for the media file based on its type (Photo or Document).
+        """
+        file_type = file_id.file_type
+
+        if file_type == FileType.CHAT_PHOTO:
+            # Handle the case for chat photos
+            if file_id.chat_id > 0:
+                peer = raw.types.InputPeerUser(user_id=file_id.chat_id, access_hash=file_id.chat_access_hash)
+            else:
+                peer = raw.types.InputPeerChannel(
+                    channel_id=utils.get_channel_id(file_id.chat_id),
+                    access_hash=file_id.chat_access_hash,
+                )
+
+            location = raw.types.InputPeerPhotoFileLocation(
+                peer=peer,
+                volume_id=file_id.volume_id,
+                local_id=file_id.local_id,
+                big=file_id.thumbnail_source == ThumbnailSource.CHAT_PHOTO_BIG,
+            )
+        elif file_type == FileType.PHOTO:
+            # Handle regular photos
+            location = raw.types.InputPhotoFileLocation(
+                id=file_id.media_id,
+                access_hash=file_id.access_hash,
+                file_reference=file_id.file_reference,
+                thumb_size=file_id.thumbnail_size,
+            )
         else:
-          peer = raw.types.InputPeerChannel(
-              channel_id=utils.get_channel_id(file_id.chat_id),
-              access_hash=file_id.chat_access_hash,
-          )
-
-      location = raw.types.InputPeerPhotoFileLocation(
-          peer=peer,
-          volume_id=file_id.volume_id,
-          local_id=file_id.local_id,
-          big=file_id.thumbnail_source == ThumbnailSource.CHAT_PHOTO_BIG,
-      )
-    elif file_type == FileType.PHOTO:
-      location = raw.types.InputPhotoFileLocation(
-          id=file_id.media_id,
-          access_hash=file_id.access_hash,
-          file_reference=file_id.file_reference,
-          thumb_size=file_id.thumbnail_size,
-      )
-    else:
-      location = raw.types.InputDocumentFileLocation(
-          id=file_id.media_id,
-          access_hash=file_id.access_hash,
-          file_reference=file_id.file_reference,
-          thumb_size=file_id.thumbnail_size,
-      )
-    return location
-
-  async def yield_file(
-      self,
-      file_id: FileId,
-      index: int,
-      offset: int,
-      first_part_cut: int,
-      last_part_cut: int,
-      part_count: int,
-      chunk_size: int,
-  ) -> Union[str, None]:
-
-    client = self.client
-    WORK_LOADS[index] += 1
-    logging.debug(f"Starting to yielding file with client {index}.")
-    media_session = await self.generate_media_session(client, file_id)
-
-    current_part = 1
-
-    location = await self.get_location(file_id)
-
-    try:
-      r = await media_session.invoke(
-          raw.functions.upload.GetFile(location=location,
-                                       offset=offset,
-                                       limit=chunk_size), )
-      if isinstance(r, raw.types.upload.File):
-        while True:
-          chunk = r.bytes
-          if not chunk:
-            break
-          elif part_count == 1:
-            yield chunk[first_part_cut:last_part_cut]
-          elif current_part == 1:
-            yield chunk[first_part_cut:]
-          elif current_part == part_count:
-            yield chunk[:last_part_cut]
-          else:
-            yield chunk
-
-          current_part += 1
-          offset += chunk_size
-
-          if current_part > part_count:
-            break
-
-          r = await media_session.invoke(
-              raw.functions.upload.GetFile(location=location,
-                                           offset=offset,
-                                           limit=chunk_size), )
-    except (TimeoutError, AttributeError):
-      pass
-    finally:
-      logging.debug(f"Finished yielding file with {current_part} parts.")
-      WORK_LOADS[index] -= 1
-
-  async def clean_cache(self) -> None:
-    """
-        function to clean the cache to reduce memory usage
+            # Handle document files
+            location = raw.types.InputDocumentFileLocation(
+                id=file_id.media_id,
+                access_hash=file_id.access_hash,
+                file_reference=file_id.file_reference,
+                thumb_size=file_id.thumbnail_size,
+            )
+        return location
+
+    async def yield_file(
+            self,
+            file_id: FileId,
+            index: int,
+            offset: int,
+            first_part_cut: int,
+            last_part_cut: int,
+            part_count: int,
+            chunk_size: int,
+    ) -> Union[str, None]:
+        """
+        Yields the file in chunks based on the specified range and chunk size.
+        This method streams the file from Telegram's server, breaking it into smaller parts.
+        """
+        client = self.client
+        WORK_LOADS[index] += 1  # Increase the workload for this client
+        logging.debug(f"Starting to yield file with client {index}.")
+        media_session = await self.generate_media_session(client, file_id)
+
+        current_part = 1
+        location = await self.get_location(file_id)
+
+        try:
+            # Fetch the file chunks
+            r = await media_session.invoke(
+                raw.functions.upload.GetFile(location=location, offset=offset, limit=chunk_size), )
+
+            if isinstance(r, raw.types.upload.File):
+                # Stream the file in chunks
+                while True:
+                    chunk = r.bytes
+                    if not chunk:
+                        break
+                    elif part_count == 1:
+                        yield chunk[first_part_cut:last_part_cut]
+                    elif current_part == 1:
+                        yield chunk[first_part_cut:]
+                    elif current_part == part_count:
+                        yield chunk[:last_part_cut]
+                    else:
+                        yield chunk
+
+                    current_part += 1
+                    offset += chunk_size
+
+                    if current_part > part_count:
+                        break
+
+                    r = await media_session.invoke(
+                        raw.functions.upload.GetFile(location=location, offset=offset, limit=chunk_size), )
+        except (TimeoutError, AttributeError):
+            pass
+        except Exception as e:
+            logging.info(f"Error at Bytestreamer Generating Chunk : {e}")
+        finally:
+            logging.debug(f"Finished yielding file with {current_part} parts.")
+            WORK_LOADS[index] -= 1  # Decrease the workload for this client
+
+    async def clean_cache(self) -> None:
         """
-    while True:
-      await asyncio.sleep(self.clean_timer)
-      print("** Caches Cleared :", self.cached_file_ids)
-      self.cached_file_ids.clear()
-      logging.debug("Cleaned the cache")
+        Function to clean the cache to reduce memory usage.
+        This method will be called periodically to clear the cached file properties.
+        """
+        await asyncio.sleep(self.clean_timer)  # Wait for the cleanup interval
+        logging.info("*** Cleaning cached file IDs...")
+        self.cached_file_ids.clear()  # Clear the cache
+        logging.debug("Cache cleaned.")
+