Resume transfers + switch to Redis streams + clean

Author: codeSamuraii

lib/metadata.py

@@ -1,12 +1,12 @@
 from starlette.datastructures import Headers
-from pydantic import BaseModel, Field, field_validator, ByteSize, StrictStr, ConfigDict, AliasChoices
+from pydantic import BaseModel, Field, field_validator, ByteSize, ConfigDict, AliasChoices
 from typing import Optional, Self, Annotated
 
 
 class FileMetadata(BaseModel):
-    name: StrictStr = Field(description="File name", min_length=2, max_length=255)
+    name: str = Field(description="File name", min_length=1, max_length=255)
     size: ByteSize = Field(description="Size in bytes", gt=0)
-    type: StrictStr = Field(description="MIME type", default='application/octet-stream')
+    type: str = Field(description="MIME type", default='application/octet-stream')
 
     model_config = ConfigDict(title="File transfer metadata", alias_generator=lambda s: f'file_{s}', populate_by_name=True, validate_by_name=True)
 
@@ -29,7 +29,7 @@ def get_from_http_headers(cls, headers: Headers, filename: str) -> Self:
         return cls(
             name=filename,
             size=headers.get('content-length', '0'),
-            type=headers.get('content-type', '') or None
+            type=headers.get('content-type', '')
         )
 
     @classmethod

lib/store.py

@@ -19,10 +19,10 @@ def __init__(self, transfer_id: str):
         self.transfer_id = transfer_id
         self.redis = self.get_redis()
 
-        self._k_queue = self.key('queue')
-        self._k_meta = self.key('metadata')
-        self._k_cleanup = f'cleanup:{transfer_id}'
-        self._k_receiver_connected = self.key('receiver_connected')
+        self._k_stream = self.key('stream')
+        self._k_metadata = self.key('metadata')
+        self._k_position = self.key('position')
+        self._k_progress = self.key('progress')
 
     @classmethod
     def get_redis(cls) -> redis.Redis:
@@ -36,26 +36,26 @@ def key(self, name: str) -> str:
         """Get the Redis key for this transfer with the provided name."""
         return f'transfer:{self.transfer_id}:{name}'
 
-    ## Queue operations ##
+    async def add_chunk(self, data: bytes) -> None:
+        """Add chunk to stream."""
+        # No maxlen limit - streams auto-expire after 5 minutes
+        await self.redis.xadd(self._k_stream, {'data': data})
 
-    async def _wait_for_queue_space(self, maxsize: int) -> None:
-        while await self.redis.llen(self._k_queue) >= maxsize:
-            await anyio.sleep(0.5)
+    async def stream_chunks(self, timeout_ms: int = 20000):
+        """Stream chunks from last position."""
+        position = await self.redis.get(self._k_position)
+        last_id = position.decode() if position else '0'
 
-    async def put_in_queue(self, data: bytes, maxsize: int = 16, timeout: float = 20.0) -> None:
-        """Add data to the transfer queue with backpressure control."""
-        with anyio.fail_after(timeout):
-            await self._wait_for_queue_space(maxsize)
-        await self.redis.lpush(self._k_queue, data)
-
-    async def get_from_queue(self, timeout: float = 20.0) -> bytes:
-        """Get data from the transfer queue with timeout."""
-        result = await self.redis.brpop([self._k_queue], timeout=timeout)
-        if not result:
-            raise TimeoutError("Timeout waiting for data")
+        while True:
+            result = await self.redis.xread({self._k_stream: last_id}, block=timeout_ms)
+            if not result:
+                raise TimeoutError("Stream read timeout")
 
-        _, data = result
-        return data
+            _, messages = result[0]
+            for message_id, fields in messages:
+                last_id = message_id
+                await self.redis.set(self._k_position, last_id, ex=300)
+                yield fields[b'data']
 
     ## Event operations ##
 
@@ -99,80 +99,35 @@ async def wait_for_event(self, event_name: str, timeout: float = 300.0) -> None:
             await pubsub.unsubscribe(event_key)
             await pubsub.aclose()
 
-    ## Metadata operations ##
-
     async def set_metadata(self, metadata: str) -> None:
         """Store transfer metadata."""
-        challenge = random.randbytes(8)
-        await self.redis.set(self._k_meta, challenge, nx=True)
-        if await self.redis.get(self._k_meta) == challenge:
-            await self.redis.set(self._k_meta, metadata, ex=300)
-        else:
-            raise KeyError("Metadata already set for this transfer.")
+        if not await self.redis.set(self._k_metadata, metadata, nx=True, ex=300):
+            raise KeyError("Transfer already exists")
 
     async def get_metadata(self) -> str | None:
-        """Retrieve transfer metadata."""
-        return await self.redis.get(self._k_meta)
-
-    ## Transfer state operations ##
-
-    async def set_receiver_connected(self) -> bool:
-        """
-        Mark that a receiver has connected for this transfer.
-        Returns True if the flag was set, False if it was already created.
-        """
-        return bool(await self.redis.set(self._k_receiver_connected, '1', ex=300, nx=True))
-
-    async def is_receiver_connected(self) -> bool:
-        """Check if a receiver has already connected."""
-        return await self.redis.exists(self._k_receiver_connected) > 0
-
-    async def set_completed(self) -> None:
-        """Mark the transfer as completed."""
-        await self.redis.set(f'completed:{self.transfer_id}', '1', ex=300, nx=True)
-
-    async def is_completed(self) -> bool:
-        """Check if the transfer is marked as completed."""
-        return await self.redis.exists(f'completed:{self.transfer_id}') > 0
-
-    async def set_interrupted(self) -> None:
-        """Mark the transfer as interrupted."""
-        await self.redis.set(f'interrupt:{self.transfer_id}', '1', ex=300, nx=True)
-        await self.redis.ltrim(self._k_queue, 0, 0)
-
-    async def is_interrupted(self) -> bool:
-        """Check if the transfer was interrupted."""
-        return await self.redis.exists(f'interrupt:{self.transfer_id}') > 0
-
-    ## Cleanup operations ##
-
-    async def cleanup_started(self) -> bool:
-        """
-        Check if cleanup has already been initiated for this transfer.
-        This uses a set/get pattern with challenge to avoid race conditions.
-        """
-        challenge = random.randbytes(8)
-        await self.redis.set(self._k_cleanup, challenge, ex=60, nx=True)
-        if await self.redis.get(self._k_cleanup) == challenge:
-            return False
-        return True
-
-    async def cleanup(self) -> int:
-        """Remove all keys related to this transfer."""
-        if await self.cleanup_started():
-            return 0
+        """Get transfer metadata."""
+        return await self.redis.get(self._k_metadata)
 
-        pattern = self.key('*')
-        keys_to_delete = set()
+    async def save_progress(self, bytes_downloaded: int) -> None:
+        """Save download progress."""
+        await self.redis.set(self._k_progress, str(bytes_downloaded), ex=300)
 
+    async def get_progress(self) -> int:
+        """Get download progress."""
+        progress = await self.redis.get(self._k_progress)
+        return int(progress) if progress else 0
+
+    async def cleanup(self) -> None:
+        """Delete all transfer data."""
+        pattern = self.key('*')
         cursor = 0
+        keys = []
+
         while True:
-            cursor, keys = await self.redis.scan(cursor, match=pattern)
-            keys_to_delete |= set(keys)
+            cursor, batch = await self.redis.scan(cursor, match=pattern)
+            keys.extend(batch)
             if cursor == 0:
                 break
 
-        if keys_to_delete:
-            self.debug(f"- Cleaning up {len(keys_to_delete)} keys")
-            return await self.redis.delete(*keys_to_delete)
-        return 0
+        if keys:
+            await self.redis.delete(*keys)

lib/transfer.py

@@ -37,12 +37,14 @@ def __init__(self, uid: str, file: FileMetadata):
 
     @classmethod
     async def create(cls, uid: str, file: FileMetadata):
+        """Create a new transfer using the provided identifier and file metadata."""
         transfer = cls(uid, file)
         await transfer.store.set_metadata(file.to_json())
         return transfer
 
     @classmethod
     async def get(cls, uid: str):
+        """Fetch a transfer from the store using the provided identifier."""
         store = Store(uid)
         metadata_json = await store.get_metadata()
         if not metadata_json:
@@ -58,122 +60,81 @@ def _format_uid(uid: str):
     def get_file_info(self):
         return self.file.name, self.file.size, self.file.type
 
-    async def wait_for_event(self, event_name: str, timeout: float = 300.0):
-        await self.store.wait_for_event(event_name, timeout)
+    async def notify_receiver_connected(self):
+        """Notify sender that receiver connected."""
+        await self.store.set_event('receiver_connected')
 
-    async def set_client_connected(self):
-        self.debug(f"▼ Notifying sender that receiver is connected...")
-        await self.store.set_event('client_connected')
+    async def wait_for_receiver(self):
+        """Wait for receiver to connect."""
+        self.info(f"△ Waiting for receiver...")
+        await self.store.wait_for_event('receiver_connected')
+        self.debug(f"△ Receiver connected")
 
-    async def wait_for_client_connected(self):
-        self.info(f"△ Waiting for client to connect...")
-        await self.wait_for_event('client_connected')
-        self.debug(f"△ Received client connected notification.")
-
-    async def is_receiver_connected(self) -> bool:
-        return await self.store.is_receiver_connected()
-
-    async def set_receiver_connected(self) -> bool:
-        return await self.store.set_receiver_connected()
-
-    async def is_interrupted(self) -> bool:
-        return await self.store.is_interrupted()
-
-    async def set_interrupted(self):
-        await self.store.set_interrupted()
-
-    async def is_completed(self) -> bool:
-        return await self.store.is_completed()
-
-    async def set_completed(self):
-        await self.store.set_completed()
-
-    async def collect_upload(self, stream: AsyncIterator[bytes], on_error: Callable[[Exception | str], Awaitable[None]]) -> None:
+    async def consume_upload(self, stream: AsyncIterator[bytes], on_error: Callable[[Exception | str], Awaitable[None]]) -> None:
+        """Consume upload stream and add chunks to Redis stream."""
         self.bytes_uploaded = 0
 
         try:
             async for chunk in stream:
                 if not chunk:
-                    self.debug(f"△ Empty chunk received, ending upload.")
                     break
 
-                if await self.is_interrupted():
-                    raise TransferError("Transfer was interrupted by the receiver.", propagate=False)
-
-                await self.store.put_in_queue(chunk)
+                await self.store.add_chunk(chunk)
                 self.bytes_uploaded += len(chunk)
 
             if self.bytes_uploaded < self.file.size:
-                raise TransferError("Received less data than expected.", propagate=True)
+                raise TransferError("Incomplete upload", propagate=True)
 
-            self.debug(f"△ End of upload, sending done marker.")
-            await self.store.put_in_queue(self.DONE_FLAG)
+            await self.store.add_chunk(self.DONE_FLAG)
+            self.debug(f"△ Upload complete: {self.bytes_uploaded} bytes")
 
-        except (ClientDisconnect, WebSocketDisconnect) as e:
-            self.error(f"△ Unexpected upload error: {e}")
-            await self.store.put_in_queue(self.DEAD_FLAG)
+        except (ClientDisconnect, WebSocketDisconnect):
+            self.error(f"△ Sender disconnected")
+            await self.store.add_chunk(self.DEAD_FLAG)
 
-        except TimeoutError as e:
-            self.warning(f"△ Timeout during upload.", exc_info=True)
-            await on_error("Timeout during upload.")
+        except TimeoutError:
+            self.warning(f"△ Upload timeout")
+            await on_error("Upload timeout")
 
         except TransferError as e:
-            self.warning(f"△ Upload error: {e}")
             if e.propagate:
-                await self.store.put_in_queue(self.DEAD_FLAG)
-            else:
-                await on_error(e)
+                await self.store.add_chunk(self.DEAD_FLAG)
+            await on_error(e)
 
-        finally:
-            await anyio.sleep(1.0)
+    async def produce_download(self, on_error: Callable[[Exception | str], Awaitable[None]]) -> AsyncIterator[bytes]:
+        """Produce download stream from Redis stream."""
+        self.bytes_downloaded = await self.store.get_progress()
 
-    async def supply_download(self, on_error: Callable[[Exception | str], Awaitable[None]]) -> AsyncIterator[bytes]:
-        self.bytes_downloaded = 0
+        if self.bytes_downloaded > 0:
+            self.info(f"▼ Resuming from byte {self.bytes_downloaded}")
 
         try:
-            while True:
-                chunk = await self.store.get_from_queue()
-
+            async for chunk in self.store.stream_chunks():
                 if chunk == self.DEAD_FLAG:
-                    raise TransferError("Sender disconnected.")
-
-                if chunk == self.DONE_FLAG and self.bytes_downloaded < self.file.size:
-                    raise TransferError("Received less data than expected.")
+                    raise TransferError("Sender disconnected")
 
-                elif chunk == self.DONE_FLAG:
-                    self.debug(f"▼ Done marker received, ending download.")
+                if chunk == self.DONE_FLAG:
+                    if self.bytes_downloaded >= self.file.size:
+                        self.debug(f"▼ Download complete: {self.bytes_downloaded} bytes")
                     break
 
                 self.bytes_downloaded += len(chunk)
+                await self.store.save_progress(self.bytes_downloaded)
                 yield chunk
 
-        except Exception as e:
-            self.error(f"▼ Unexpected download error!", exc_info=True)
-            self.debug("Debug info:", stack_info=True)
-            await on_error(e)
-
         except TransferError as e:
-            self.warning(f"▼ Download error")
+            await on_error(e)
+        except Exception as e:
+            self.error(f"▼ Download error", exc_info=True)
             await on_error(e)
 
     async def cleanup(self):
-        try:
-            with anyio.fail_after(30.0):
-                await self.store.cleanup()
-        except TimeoutError:
-            self.warning(f"- Cleanup timed out.")
-            pass
+        """Clean up transfer data."""
+        await self.store.cleanup()
 
     async def finalize_download(self):
-        # self.debug("▼ Finalizing download...")
-        if self.bytes_downloaded < self.file.size and not await self.is_interrupted():
-            self.warning("▼ Client disconnected before download was complete.")
-            await self.set_interrupted()
-
-        await self.cleanup()
-        # self.debug("▼ Finalizing download...")
-        if self.bytes_downloaded < self.file.size and not await self.is_interrupted():
-            self.warning("▼ Client disconnected before download was complete.")
-            await self.set_interrupted()
-
-        await self.cleanup()
+        """Finalize download and cleanup if complete."""
+        if self.bytes_downloaded < self.file.size:
+            self.info(f"▼ Download paused at {self.bytes_downloaded}/{self.file.size} bytes")
+        else:
+            await self.cleanup()

static/index.html

@@ -54,7 +54,7 @@ <h2>Send a file</h2>
             <section class="section">
                 <div class="code-section">
                     <h3>Using cURL</h3>
-                    <p>You can use the <code class="inline-highlight">curl</code> command to transfer from your terminal. 100 MiB maximum.</p>
+                    <p>You can use the <code class="inline-highlight">curl</code> command to transfer from your terminal. 1 GiB maximum.</p>
 
                     <div class="code-block">
                         <code><span class="code-comment"># Send</span>