mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-16 04:29:31 +00:00
84 lines
3.1 KiB
Python
84 lines
3.1 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import ssl
|
|
from collections.abc import Awaitable, Callable
|
|
from contextlib import suppress
|
|
from threading import Thread
|
|
from typing import Final
|
|
|
|
Handler = Callable[[asyncio.StreamReader, asyncio.StreamWriter], Awaitable[None]]
|
|
|
|
|
|
async def close_writer(writer: asyncio.StreamWriter) -> None:
|
|
writer.close()
|
|
with suppress(ConnectionError, OSError):
|
|
await writer.wait_closed()
|
|
|
|
|
|
class LoopbackServer:
|
|
def __init__(self, port: int, handler: Handler, ssl_context: ssl.SSLContext | None = None) -> None:
|
|
self.loop: Final = asyncio.new_event_loop()
|
|
self._port: Final = port
|
|
self._handler: Final = handler
|
|
self._ssl_context: Final = ssl_context
|
|
self._server: asyncio.Server | None = None
|
|
self._thread: Final = Thread(target=self._run, name=f"integration-relay-{port}", daemon=True)
|
|
self._connections: tuple[asyncio.Task[None], ...] = ()
|
|
|
|
def start(self) -> None:
|
|
try:
|
|
self._thread.start()
|
|
asyncio.run_coroutine_threadsafe(self._bind(), self.loop).result(timeout=10)
|
|
except BaseException:
|
|
self.stop()
|
|
raise
|
|
|
|
def stop(self) -> None:
|
|
if not self._thread.is_alive():
|
|
if not self.loop.is_closed():
|
|
self.loop.close()
|
|
return
|
|
try:
|
|
asyncio.run_coroutine_threadsafe(self._close(), self.loop).result(timeout=10)
|
|
finally:
|
|
self.loop.call_soon_threadsafe(self.loop.stop)
|
|
self._thread.join(10)
|
|
assert not self._thread.is_alive(), f"Loopback server thread on port {self._port} did not stop"
|
|
|
|
def _run(self) -> None:
|
|
asyncio.set_event_loop(self.loop)
|
|
try:
|
|
self.loop.run_forever()
|
|
finally:
|
|
self.loop.run_until_complete(self.loop.shutdown_asyncgens())
|
|
self.loop.close()
|
|
|
|
async def _bind(self) -> None:
|
|
self._server = await asyncio.start_server(self._serve, "127.0.0.1", self._port, ssl=self._ssl_context)
|
|
|
|
async def _serve(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
|
|
task: Final = asyncio.current_task()
|
|
assert task is not None
|
|
self._connections = (*self._connections, task)
|
|
try:
|
|
await self._handler(reader, writer)
|
|
finally:
|
|
await close_writer(writer)
|
|
self._connections = tuple(connection for connection in self._connections if connection is not task)
|
|
|
|
async def _close(self) -> None:
|
|
server: Final = self._server
|
|
if server is not None:
|
|
server.close()
|
|
connections: Final = self._connections
|
|
for task in connections:
|
|
task.cancel()
|
|
await asyncio.gather(*connections, return_exceptions=True)
|
|
current: Final = asyncio.current_task()
|
|
remaining: Final = tuple(task for task in asyncio.all_tasks(self.loop) if task is not current)
|
|
for task in remaining:
|
|
task.cancel()
|
|
await asyncio.gather(*remaining, return_exceptions=True)
|
|
if server is not None:
|
|
await server.wait_closed()
|