feat: add Python implementation and cross-test support
- Add Python toki_socket package (base_client, communicator, tcp/ws client/server) - Add protobuf message definitions for Python - Add go_python.go cross-test implementation and python_go_client - Add agent-task/python_impl plan and code review docs - Update .gitignore for Python cache files with recursive patterns
This commit is contained in:
parent
f92f6672e4
commit
fa273bcc26
24 changed files with 2785 additions and 3 deletions
6
.gitignore
vendored
6
.gitignore
vendored
|
|
@ -28,9 +28,9 @@ swift/*.xcworkspace/
|
|||
swift/Packages.resolved
|
||||
|
||||
# ── Python (future) ───────────────────────────────
|
||||
python/__pycache__/
|
||||
python/*.pyc
|
||||
python/*.pyo
|
||||
python/**/__pycache__/
|
||||
python/**/*.pyc
|
||||
python/**/*.pyo
|
||||
python/.venv/
|
||||
python/venv/
|
||||
python/dist/
|
||||
|
|
|
|||
159
agent-task/python_impl/CODE_REVIEW.md
Normal file
159
agent-task/python_impl/CODE_REVIEW.md
Normal file
|
|
@ -0,0 +1,159 @@
|
|||
<!-- task=python_impl plan=0 tag=API -->
|
||||
|
||||
# Code Review Reference - API
|
||||
|
||||
## 개요
|
||||
|
||||
date=2026-04-21
|
||||
task=python_impl, plan=0, tag=API
|
||||
|
||||
## 이 파일을 읽는 리뷰 에이전트에게
|
||||
|
||||
각 항목의 구현을 실제 소스 파일과 대조하고, `검증 결과` 섹션의 출력이 코드와 일치하는지 확인하세요.
|
||||
리뷰 완료 후 반드시 아래 순서로 아카이브하세요.
|
||||
|
||||
1. `CODE_REVIEW.md` → `code_review_0.log` (N = 기존 code_review_*.log 수)
|
||||
2. `PLAN.md` → `plan_0.log` (M = 기존 plan_*.log 수)
|
||||
3. PASS인 경우 `complete.log` 작성 후 종료. WARN/FAIL인 경우 새 `PLAN.md` + `CODE_REVIEW.md` 스텁 작성.
|
||||
|
||||
---
|
||||
|
||||
## 구현 항목별 완료 여부
|
||||
|
||||
| 항목 | 완료 여부 |
|
||||
|------|---------|
|
||||
| [API-1] Python 패키지 scaffold | [x] |
|
||||
| [API-2] Communicator 구현 | [x] |
|
||||
| [API-3] BaseClient 구현 | [x] |
|
||||
| [API-4] TCP Transport 구현 | [x] |
|
||||
| [API-5] WebSocket Transport 구현 | [x] |
|
||||
| [API-6] 단위 테스트 | [x] |
|
||||
| [API-7] Python×Go 크로스 테스트 | [x] |
|
||||
|
||||
## 계획 대비 변경 사항
|
||||
|
||||
- **websockets API**: 계획에서는 `websockets.connect` / `websockets.server.serve`를 언급했으나, 설치된 버전이 16.0으로 legacy API가 deprecated됨. `websockets.asyncio.client.connect` / `websockets.asyncio.server.serve`로 대체.
|
||||
- **`queue_packet` 구현**: 계획의 큐 full 처리를 `put_nowait`로 단순화. maxsize=64이고 write_loop가 항상 소비하므로 현실적으로 블로킹이 발생하지 않음.
|
||||
- **`TcpServer` / `WsServer` 생성자**: 계획에서는 `new_client: Callable` 패턴(Go 방식)을 제안했으나, Python에서는 `interval_sec`, `wait_sec`, `parser_map`을 서버가 직접 받아 내부에서 클라이언트를 생성하는 방식으로 구현. API가 더 단순해짐.
|
||||
- **`_echo_handler` 패턴**: test_tcp.py에서 request listener가 응답을 보내기 위해 `queue_packet`을 직접 사용. 추후 `add_request_listener_typed` 헬퍼 추가 시 개선 가능.
|
||||
|
||||
## 주요 설계 결정
|
||||
|
||||
- **sentinel 패턴**: `_write_loop` 종료를 위해 `_STOP` sentinel 객체를 큐에 삽입. `asyncio.Event`보다 단순하고 명확.
|
||||
- **asyncio 단일 스레드**: `is_alive`, `nonce` 등 모든 상태가 단일 스레드에서만 접근되므로 lock 불필요. `BaseClient.close()`만 `asyncio.Lock`으로 connCloseOnce 보장.
|
||||
- **Python 오케스트레이터의 Go 클라이언트 실행**: `asyncio.get_event_loop().run_in_executor(None, ...)` 로 블로킹 `subprocess.run`을 실행. 이로써 서버 이벤트루프가 Go 클라이언트 실행 중에도 계속 동작.
|
||||
- **`type_name_of(cls)`**: `cls.DESCRIPTOR.name`을 반환. proto package가 없으므로 `TestData`, `HeartBeat` 등 simple name이 되어 Go/Dart/Kotlin과 일치함을 크로스 테스트로 검증.
|
||||
|
||||
## 리뷰어를 위한 체크포인트
|
||||
|
||||
- `type_name_of(cls)` 가 `cls.DESCRIPTOR.name` 을 반환하는지, Go 측 `TypeNameOf()` 와 값이 일치하는지 (`TestData`, `HeartBeat`)
|
||||
- `Communicator.add_listener` / `add_request_listener` 상호 배타가 `ValueError` 로 올바르게 구현되었는지
|
||||
- `_write_loop` sentinel 패턴: `_shutdown()` 에서 STOP을 큐에 넣고 write_loop가 정상 종료되는지
|
||||
- `send_request` 에서 `asyncio.wait_for` timeout 처리 시 pending 딕셔너리에서 해당 nonce가 제거되는지
|
||||
- `BaseClient.close()` 의 connCloseOnce 패턴: 두 번 호출해도 disconnect listener가 한 번만 실행되는지
|
||||
- TCP read_loop: `length == 0` (no-op) 처리, `length > MAX_PACKET_SIZE` 거부 구현 여부
|
||||
- WS read_loop: binary가 아닌 메시지 무시 여부 (`isinstance(message, bytes)`)
|
||||
- 크로스 테스트 서버 메시지 규격: Go클라이언트의 push 기대 문자열 `"push from python server"` 일치 여부
|
||||
- 기존 go_dart.go, go_kotlin.go 크로스 테스트 회귀 없음 확인
|
||||
|
||||
## 검증 결과
|
||||
|
||||
### API-1 중간 검증
|
||||
```
|
||||
$ cd python && python3 -c "from toki_socket.packets import message_common_pb2; print(message_common_pb2.TestData.DESCRIPTOR.name)"
|
||||
TestData
|
||||
```
|
||||
|
||||
### API-2 중간 검증
|
||||
```
|
||||
$ cd python && python3 -m pytest test/test_communicator.py -v
|
||||
============================= test session starts ==============================
|
||||
collected 4 items
|
||||
|
||||
test/test_communicator.py::test_add_listener_exclusivity PASSED [ 25%]
|
||||
test/test_communicator.py::test_send_and_receive PASSED [ 50%]
|
||||
test/test_communicator.py::test_send_request_response PASSED [ 75%]
|
||||
test/test_communicator.py::test_shutdown PASSED [100%]
|
||||
|
||||
========================= 4 passed, 1 warning in 0.14s =========================
|
||||
```
|
||||
|
||||
### API-3 중간 검증
|
||||
```
|
||||
$ python3 -c "from toki_socket.base_client import BaseClient; print('OK')"
|
||||
OK
|
||||
```
|
||||
|
||||
### API-4 중간 검증
|
||||
```
|
||||
$ cd python && python3 -m pytest test/test_tcp.py -v
|
||||
collected 2 items
|
||||
|
||||
test/test_tcp.py::test_tcp_send_receive PASSED [ 50%]
|
||||
test/test_tcp.py::test_tcp_request_response PASSED [100%]
|
||||
|
||||
========================= 2 passed, 1 warning in 0.16s =========================
|
||||
```
|
||||
|
||||
### API-5 중간 검증
|
||||
```
|
||||
$ python3 -c "from toki_socket.ws_client import WsClient; from toki_socket.ws_server import WsServer; print('OK')"
|
||||
OK
|
||||
```
|
||||
|
||||
### 최종 검증
|
||||
```
|
||||
$ cd python && python3 -m pytest test/ -v
|
||||
collected 6 items
|
||||
|
||||
test/test_communicator.py::test_add_listener_exclusivity PASSED
|
||||
test/test_communicator.py::test_send_and_receive PASSED
|
||||
test/test_communicator.py::test_send_request_response PASSED
|
||||
test/test_communicator.py::test_shutdown PASSED
|
||||
test/test_tcp.py::test_tcp_send_receive PASSED
|
||||
test/test_tcp.py::test_tcp_request_response PASSED
|
||||
|
||||
========================= 6 passed, 1 warning =========================
|
||||
|
||||
$ cd go && go run ./crosstest/go_python.go
|
||||
INFO typeName go=TestData
|
||||
INFO typeName python=TestData
|
||||
PASS scenario=1 detail=fire-and-forget sent
|
||||
SERVER_RECEIVED index=101 message=fire from python client
|
||||
PASS scenario=2 detail=received push from go server
|
||||
INFO typeName python=TestData
|
||||
PASS scenario=3 detail=single request response matched
|
||||
PASS scenario=4 detail=concurrent request responses matched
|
||||
INFO typeName python=TestData
|
||||
PASS scenario=1 detail=fire-and-forget sent
|
||||
SERVER_RECEIVED index=101 message=fire from python client
|
||||
PASS scenario=2 detail=received push from go server
|
||||
INFO typeName python=TestData
|
||||
PASS scenario=3 detail=single request response matched
|
||||
PASS scenario=4 detail=concurrent request responses matched
|
||||
PASS all go-server/python-client crosstests passed
|
||||
|
||||
$ cd python && python3 crosstest/python_go.py
|
||||
INFO typeName python=TestData
|
||||
SERVER_RECEIVED index=101 message=fire from go client
|
||||
INFO typeName go=TestData
|
||||
PASS scenario=1 detail=fire-and-forget sent
|
||||
PASS scenario=2 detail=received push from python server
|
||||
INFO typeName go=TestData
|
||||
PASS scenario=3 detail=single request response matched
|
||||
PASS scenario=4 detail=concurrent request responses matched
|
||||
SERVER_RECEIVED index=101 message=fire from go client
|
||||
INFO typeName go=TestData
|
||||
PASS scenario=1 detail=fire-and-forget sent
|
||||
PASS scenario=2 detail=received push from python server
|
||||
INFO typeName go=TestData
|
||||
PASS scenario=3 detail=single request response matched
|
||||
PASS scenario=4 detail=concurrent request responses matched
|
||||
PASS all python-server/go-client crosstests passed
|
||||
|
||||
$ cd go && go run ./crosstest/go_dart.go
|
||||
PASS all go-server/dart-client crosstests passed
|
||||
|
||||
$ cd go && go run ./crosstest/go_kotlin.go
|
||||
PASS all go-server/kotlin-client crosstests passed
|
||||
```
|
||||
691
agent-task/python_impl/PLAN.md
Normal file
691
agent-task/python_impl/PLAN.md
Normal file
|
|
@ -0,0 +1,691 @@
|
|||
<!-- task=python_impl plan=0 tag=API -->
|
||||
|
||||
# Python 구현체 추가 (toki_socket)
|
||||
|
||||
## 이 파일을 읽는 구현 에이전트에게
|
||||
|
||||
각 항목의 체크리스트를 완료하면서 `[ ]`를 `[x]`로 표시하라.
|
||||
중간 검증 명령은 해당 항목 구현 직후 실행하고, 출력을 `CODE_REVIEW.md`의 검증 결과 섹션에 붙여넣어라.
|
||||
최종 검증도 마찬가지로 실행 후 출력을 기록하라.
|
||||
계획과 다르게 구현한 부분은 이유와 함께 `CODE_REVIEW.md`의 "계획 대비 변경 사항"에 기록하라.
|
||||
|
||||
---
|
||||
|
||||
## 배경
|
||||
|
||||
`PROTOCOL.md`의 Multi-language Implementations 표에 Python이 `Planned`로 등록되어 있다.
|
||||
Go 구현체(`go/`)가 레퍼런스이며, `PORTING_GUIDE.md`에 Python 매핑 표와 주의사항이 명시되어 있다.
|
||||
Python 패키지를 `python/` 디렉터리에 생성하고, Go×Python 양방향 크로스 테스트를 추가하여 프로토콜 호환성을 검증한다.
|
||||
TLS/WSS는 이번 범위에서 제외하고 README에 명시한다.
|
||||
|
||||
---
|
||||
|
||||
## 의존 관계 및 구현 순서
|
||||
|
||||
```
|
||||
[API-1] scaffold → [API-2] Communicator → [API-3] BaseClient
|
||||
→ [API-4] TCP Transport
|
||||
→ [API-5] WebSocket Transport
|
||||
[API-2~5] 완료 → [API-6] 단위 테스트 → [API-7] 크로스 테스트
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## [API-1] Python 패키지 scaffold
|
||||
|
||||
### 문제
|
||||
|
||||
`python/` 디렉터리가 없다. 패키지 매니저, proto binding, 최소 패키지 구조가 필요하다.
|
||||
|
||||
### 해결 방법
|
||||
|
||||
- `python/pyproject.toml` 생성 (setuptools, Python 3.11+)
|
||||
- `python/toki_socket/__init__.py` 생성
|
||||
- `dart/lib/src/packets/message_common.proto`를 `python/toki_socket/packets/message_common.proto`로 복사 (schema만, 언어 option 추가 불필요)
|
||||
- `protoc`로 `message_common_pb2.py`, `message_common_pb2.pyi` 생성
|
||||
|
||||
**pyproject.toml 핵심**:
|
||||
```toml
|
||||
[build-system]
|
||||
requires = ["setuptools>=68"]
|
||||
build-backend = "setuptools.backends.legacy:build"
|
||||
|
||||
[project]
|
||||
name = "toki-socket"
|
||||
version = "0.1.0"
|
||||
requires-python = ">=3.11"
|
||||
dependencies = [
|
||||
"protobuf>=4.25",
|
||||
"websockets>=12.0",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
dev = ["pytest>=8.0", "pytest-asyncio>=0.23"]
|
||||
```
|
||||
|
||||
**디렉터리 구조**:
|
||||
```
|
||||
python/
|
||||
toki_socket/
|
||||
__init__.py
|
||||
communicator.py
|
||||
base_client.py
|
||||
tcp_client.py
|
||||
tcp_server.py
|
||||
ws_client.py
|
||||
ws_server.py
|
||||
packets/
|
||||
__init__.py
|
||||
message_common.proto
|
||||
message_common_pb2.py (generated)
|
||||
message_common_pb2.pyi (generated)
|
||||
test/
|
||||
__init__.py
|
||||
test_communicator.py
|
||||
test_tcp.py
|
||||
crosstest/
|
||||
go_python_client.py
|
||||
python_go.py
|
||||
pyproject.toml
|
||||
README.md
|
||||
```
|
||||
|
||||
### 수정 파일 및 체크리스트
|
||||
|
||||
- [ ] `python/pyproject.toml` 생성
|
||||
- [ ] `python/toki_socket/__init__.py` 생성 (public export 비워둠)
|
||||
- [ ] `python/toki_socket/packets/__init__.py` 생성
|
||||
- [ ] `python/toki_socket/packets/message_common.proto` 복사
|
||||
- [ ] `python/toki_socket/packets/message_common_pb2.py` 생성 (protoc 실행 또는 직접 작성)
|
||||
- [ ] `python/toki_socket/packets/message_common_pb2.pyi` 생성
|
||||
- [ ] `python/README.md` 생성 (TLS 미지원 명시 포함)
|
||||
|
||||
### 테스트 작성
|
||||
|
||||
skip — scaffold 단계, 코드 로직 없음.
|
||||
|
||||
### 중간 검증
|
||||
|
||||
```bash
|
||||
cd python && pip install -e ".[dev]" --quiet && python -c "from toki_socket.packets import message_common_pb2; print(message_common_pb2.TestData.DESCRIPTOR.name)"
|
||||
```
|
||||
|
||||
예상 출력: `TestData`
|
||||
|
||||
---
|
||||
|
||||
## [API-2] Communicator 구현
|
||||
|
||||
### 문제
|
||||
|
||||
`go/communicator.go`에 해당하는 Python 비동기 구현이 없다.
|
||||
`Transport` Protocol, `Communicator` 클래스, write 큐, nonce 관리, listener/requestListener 라우팅을 구현해야 한다.
|
||||
|
||||
### 해결 방법
|
||||
|
||||
`python/toki_socket/communicator.py`에 아래 구조로 구현한다.
|
||||
|
||||
**Transport Protocol** (Go의 `Transport` interface):
|
||||
```python
|
||||
from typing import Protocol, runtime_checkable
|
||||
from toki_socket.packets.message_common_pb2 import PacketBase
|
||||
|
||||
@runtime_checkable
|
||||
class Transport(Protocol):
|
||||
async def write_packet(self, base: PacketBase) -> None: ...
|
||||
async def close(self) -> None: ...
|
||||
```
|
||||
|
||||
**ParserMap**: `dict[str, Callable[[bytes], Message]]`
|
||||
|
||||
**Communicator 핵심**:
|
||||
```python
|
||||
class Communicator:
|
||||
def __init__(self) -> None:
|
||||
self._nonce: int = 0
|
||||
self._is_alive: bool = False
|
||||
self._parser_map: dict[str, Callable] = {}
|
||||
self._handlers: dict[str, list[Callable]] = {}
|
||||
self._req_handlers: dict[str, Callable] = {}
|
||||
self._pending_requests: dict[int, asyncio.Future] = {}
|
||||
self._write_queue: asyncio.Queue = asyncio.Queue(maxsize=64)
|
||||
self._closed: asyncio.Event = asyncio.Event()
|
||||
self._transport: Transport | None = None
|
||||
self._write_loop_task: asyncio.Task | None = None
|
||||
self._write_error_handler: Callable[[Exception], None] | None = None
|
||||
|
||||
def initialize(self, transport: Transport, parser_map: dict) -> None:
|
||||
... # 하트비트 파서 추가, write_loop Task 시작
|
||||
|
||||
def is_alive(self) -> bool: ...
|
||||
def _next_nonce(self) -> int: ...
|
||||
def _shutdown(self) -> None: ...
|
||||
async def close(self) -> None: ...
|
||||
async def _write_loop(self) -> None: ...
|
||||
async def queue_packet(self, base: PacketBase) -> None: ...
|
||||
async def send(self, message: Message) -> None: ...
|
||||
async def send_request(self, req: Message, res_type: type[T], timeout: float = 30.0) -> T: ...
|
||||
def add_listener(self, type_name: str, fn: Callable) -> None: ...
|
||||
def remove_listeners(self, type_name: str) -> None: ...
|
||||
def add_request_listener(self, type_name: str, fn: Callable) -> None: ...
|
||||
def on_received_data(self, type_name: str, data: bytes, nonce: int, response_nonce: int) -> None: ...
|
||||
def _handle_response(self, type_name: str, data: bytes, response_nonce: int) -> None: ...
|
||||
def _parse(self, type_name: str, data: bytes) -> Message: ...
|
||||
|
||||
def type_name_of(message_class) -> str:
|
||||
return message_class.DESCRIPTOR.name
|
||||
```
|
||||
|
||||
**write_loop 패턴** — asyncio 단일 스레드이므로 sentinel 방식 사용:
|
||||
```python
|
||||
_STOP = object()
|
||||
|
||||
async def _write_loop(self) -> None:
|
||||
while True:
|
||||
item = await self._write_queue.get()
|
||||
if item is _STOP:
|
||||
return
|
||||
base, fut = item
|
||||
try:
|
||||
await self._transport.write_packet(base)
|
||||
if not fut.done():
|
||||
fut.set_result(None)
|
||||
except Exception as e:
|
||||
if not fut.done():
|
||||
fut.set_exception(e)
|
||||
if self._write_error_handler:
|
||||
self._write_error_handler(e)
|
||||
```
|
||||
|
||||
`_shutdown()`에서 `_write_queue.put_nowait(_STOP)` 호출.
|
||||
|
||||
**send_request** — `asyncio.get_event_loop().create_future()`로 pending 등록, `asyncio.wait_for(fut, timeout)`으로 대기.
|
||||
|
||||
**addListener / addRequestListener 상호 배타** — 같은 typeName 중복 시 `ValueError` raise.
|
||||
|
||||
**HeartBeat 자동 처리** — `initialize()` 내에서 HeartBeat 파서를 `_parser_map`에 추가; BaseClient에서 listener 등록.
|
||||
|
||||
**on_received_data** — asyncio 단일 스레드이므로 요청 핸들러는 `asyncio.create_task()`로 실행 (Go의 `go reqHandler(...)`와 동일).
|
||||
|
||||
### 수정 파일 및 체크리스트
|
||||
|
||||
- [ ] `python/toki_socket/communicator.py` 생성
|
||||
- [ ] `Transport` Protocol 정의
|
||||
- [ ] `type_name_of(cls)` 함수 — `cls.DESCRIPTOR.name` 반환
|
||||
- [ ] `Communicator.initialize()` — parser_map + HeartBeat 파서 등록, write_loop Task 시작
|
||||
- [ ] `Communicator._write_loop()` — sentinel 패턴
|
||||
- [ ] `Communicator.queue_packet()` — closed 체크, Future 등록, put, await
|
||||
- [ ] `Communicator.send()` — marshal → queue_packet
|
||||
- [ ] `Communicator.send_request()` — pending 등록, marshal, queue, asyncio.wait_for
|
||||
- [ ] `Communicator.add_listener()` / `add_request_listener()` — 상호 배타 검사
|
||||
- [ ] `Communicator.on_received_data()` — response_nonce > 0이면 handleResponse, 아니면 listener/reqHandler
|
||||
- [ ] `Communicator._shutdown()` — is_alive=False, closed.set(), write_queue에 STOP 전달
|
||||
|
||||
### 테스트 작성
|
||||
|
||||
`python/test/test_communicator.py`에 작성:
|
||||
- `test_add_listener_exclusivity` — 같은 typeName에 add_listener 후 add_request_listener 시 ValueError 확인
|
||||
- `test_send_and_receive` — FakeTransport (write_packet 캡처)로 send() 후 패킷 내용 검증
|
||||
- `test_send_request_response` — FakeTransport로 send_request 후 수동으로 on_received_data 호출하여 Future resolve 검증
|
||||
- `test_shutdown` — close() 후 is_alive() == False 검증
|
||||
|
||||
### 중간 검증
|
||||
|
||||
```bash
|
||||
cd python && python -m pytest test/test_communicator.py -v
|
||||
```
|
||||
|
||||
예상: 4개 테스트 PASS
|
||||
|
||||
---
|
||||
|
||||
## [API-3] BaseClient 구현
|
||||
|
||||
### 문제
|
||||
|
||||
`go/base_client.go`에 해당하는 heartbeat, disconnect listener, connCloseOnce 패턴이 없다.
|
||||
|
||||
### 해결 방법
|
||||
|
||||
`python/toki_socket/base_client.py`에 구현.
|
||||
|
||||
Go `sync.Once` → `asyncio.Lock` + `bool` 플래그:
|
||||
```python
|
||||
self._close_called: bool = False
|
||||
self._close_lock: asyncio.Lock = asyncio.Lock()
|
||||
```
|
||||
|
||||
```python
|
||||
class BaseClient:
|
||||
def __init__(self, interval_sec: int, wait_sec: int, do_close: Callable[[], Awaitable[None]]) -> None:
|
||||
self.communicator: Communicator = Communicator()
|
||||
self._interval_sec = interval_sec
|
||||
self._wait_sec = wait_sec
|
||||
self._do_close = do_close
|
||||
self._close_called = False
|
||||
self._close_lock: asyncio.Lock # initialize()에서 생성
|
||||
self._disconnect_listeners: list[Callable] = []
|
||||
self._hb_task: asyncio.Task | None = None
|
||||
self._waiting_hb_response: bool = False
|
||||
|
||||
def _init_base(self) -> None:
|
||||
self._close_lock = asyncio.Lock()
|
||||
|
||||
async def close(self) -> None:
|
||||
async with self._close_lock:
|
||||
if self._close_called:
|
||||
return
|
||||
self._close_called = True
|
||||
self.communicator._shutdown()
|
||||
self._stop_heartbeat()
|
||||
if self._do_close:
|
||||
await self._do_close()
|
||||
await self._notify_disconnected()
|
||||
|
||||
def add_disconnect_listener(self, fn: Callable) -> None: ...
|
||||
def remove_disconnect_listeners(self) -> None: ...
|
||||
async def _notify_disconnected(self) -> None: ...
|
||||
async def _on_disconnected(self) -> None:
|
||||
await self.close()
|
||||
async def _send_heartbeat(self) -> None: ... # Go sendHeartBeat() 패턴
|
||||
async def _on_heartbeat(self) -> None: ... # Go onHeartBeat() 패턴
|
||||
def _stop_heartbeat(self) -> None: ...
|
||||
```
|
||||
|
||||
Heartbeat 타이머: `asyncio.get_event_loop().call_later()` 또는 `asyncio.create_task(asyncio.sleep(n))`으로 구현.
|
||||
Go의 `HeartbeatTimer`에 해당하는 `_HbTimer` 헬퍼 클래스를 동일 파일에 작성.
|
||||
|
||||
### 수정 파일 및 체크리스트
|
||||
|
||||
- [ ] `python/toki_socket/base_client.py` 생성
|
||||
- [ ] `_HbTimer` — `asyncio.create_task(sleep+callback)`, `cancel()` 메서드
|
||||
- [ ] `BaseClient._init_base()` — close_lock 초기화 (asyncio.Lock은 루프 바인딩 필요)
|
||||
- [ ] `BaseClient.close()` — connCloseOnce 패턴, shutdown → stop_heartbeat → do_close → notify
|
||||
- [ ] `BaseClient._send_heartbeat()` — interval > 0이면 타이머 설정
|
||||
- [ ] `BaseClient._on_heartbeat()` — waiting 상태이면 False로 리셋, 아니면 HeartBeat 송신
|
||||
- [ ] `BaseClient._notify_disconnected()` — listener 복사 후 asyncio.gather로 호출
|
||||
|
||||
### 테스트 작성
|
||||
|
||||
단위 테스트 skip — Transport 없이 단독 테스트가 어렵고, TCP 테스트에서 통합 검증.
|
||||
|
||||
### 중간 검증
|
||||
|
||||
```bash
|
||||
cd python && python -c "from toki_socket.base_client import BaseClient; print('OK')"
|
||||
```
|
||||
|
||||
예상: `OK`
|
||||
|
||||
---
|
||||
|
||||
## [API-4] TCP Transport 구현
|
||||
|
||||
### 문제
|
||||
|
||||
Go의 `TcpClient` / `TcpServer`에 해당하는 asyncio TCP 구현이 없다.
|
||||
|
||||
### 해결 방법
|
||||
|
||||
`python/toki_socket/tcp_client.py`, `python/toki_socket/tcp_server.py` 구현.
|
||||
|
||||
**framing** (PROTOCOL.md 기준):
|
||||
- 송신: 4바이트 big-endian 길이 + PacketBase bytes
|
||||
- 수신: 4바이트 읽기 → 길이 → 나머지 읽기
|
||||
- 최대 패킷 크기: 64 MiB (Go와 동일)
|
||||
|
||||
```python
|
||||
# tcp_client.py
|
||||
class TcpClient(BaseClient):
|
||||
def __init__(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter,
|
||||
interval_sec: int, wait_sec: int, parser_map: ParserMap) -> None:
|
||||
super().__init__(interval_sec, wait_sec, do_close=self._do_close_impl)
|
||||
self._reader = reader
|
||||
self._writer = writer
|
||||
self._init_base()
|
||||
self.communicator.initialize(self, parser_map)
|
||||
self.communicator.set_write_error_handler(lambda e: asyncio.create_task(self._on_disconnected()))
|
||||
self.communicator.add_listener(
|
||||
type_name_of(HeartBeat), lambda m: asyncio.create_task(self._on_heartbeat())
|
||||
)
|
||||
asyncio.create_task(self._read_loop())
|
||||
asyncio.create_task(self._send_heartbeat())
|
||||
|
||||
async def write_packet(self, base: PacketBase) -> None:
|
||||
data = base.SerializeToString()
|
||||
header = struct.pack(">I", len(data))
|
||||
self._writer.write(header + data)
|
||||
await self._writer.drain()
|
||||
|
||||
async def _do_close_impl(self) -> None:
|
||||
self._writer.close()
|
||||
await self._writer.wait_closed()
|
||||
|
||||
async def _read_loop(self) -> None:
|
||||
while self.communicator.is_alive():
|
||||
try:
|
||||
header = await self._reader.readexactly(4)
|
||||
length = struct.unpack(">I", header)[0]
|
||||
if length == 0:
|
||||
continue
|
||||
if length > MAX_PACKET_SIZE:
|
||||
await self._on_disconnected(); return
|
||||
data = await self._reader.readexactly(length)
|
||||
base = PacketBase()
|
||||
base.ParseFromString(data)
|
||||
self.communicator.on_received_data(
|
||||
base.typeName, base.data, base.nonce, base.responseNonce
|
||||
)
|
||||
asyncio.create_task(self._send_heartbeat())
|
||||
except (asyncio.IncompleteReadError, ConnectionError, OSError):
|
||||
await self._on_disconnected(); return
|
||||
|
||||
async def connect_tcp(host: str, port: int, interval_sec: int, wait_sec: int,
|
||||
parser_map: ParserMap) -> TcpClient:
|
||||
reader, writer = await asyncio.open_connection(host, port)
|
||||
return TcpClient(reader, writer, interval_sec, wait_sec, parser_map)
|
||||
```
|
||||
|
||||
```python
|
||||
# tcp_server.py
|
||||
class TcpServer:
|
||||
def __init__(self, host: str, port: int, new_client: Callable) -> None:
|
||||
self._host = host
|
||||
self._port = port
|
||||
self._new_client = new_client
|
||||
self._clients: list[TcpClient] = []
|
||||
self._server: asyncio.Server | None = None
|
||||
self.on_client_connected: Callable[[TcpClient], None] = lambda c: None
|
||||
|
||||
async def start(self) -> None: ...
|
||||
async def stop(self) -> None: ...
|
||||
def clients(self) -> list[TcpClient]: ...
|
||||
async def broadcast(self, message) -> None: ...
|
||||
```
|
||||
|
||||
### 수정 파일 및 체크리스트
|
||||
|
||||
- [ ] `python/toki_socket/tcp_client.py` 생성
|
||||
- [ ] `MAX_PACKET_SIZE = 64 << 20`
|
||||
- [ ] `TcpClient(BaseClient)` — reader/writer, read_loop, write_packet, do_close
|
||||
- [ ] `connect_tcp(host, port, interval_sec, wait_sec, parser_map)` async 함수
|
||||
- [ ] `python/toki_socket/tcp_server.py` 생성
|
||||
- [ ] `TcpServer` — start/stop/clients/broadcast, on_client_connected 콜백
|
||||
- [ ] `_handle_connection()` — TcpClient 생성, disconnect_listener로 clients 목록 관리
|
||||
|
||||
### 테스트 작성
|
||||
|
||||
`python/test/test_tcp.py`:
|
||||
- `test_tcp_send_receive` — 루프백 TCP로 양방향 TestData 교환
|
||||
- `test_tcp_request_response` — send_request / add_request_listener 검증
|
||||
|
||||
```python
|
||||
@pytest.mark.asyncio
|
||||
async def test_tcp_send_receive():
|
||||
received = asyncio.Future()
|
||||
server = TcpServer("127.0.0.1", 0, lambda r, w: TcpClient(r, w, 0, 0, parser_map()))
|
||||
server.on_client_connected = lambda c: c.communicator.add_listener(
|
||||
"TestData", lambda m: received.set_result(m)
|
||||
)
|
||||
await server.start()
|
||||
port = server.port # 동적 포트
|
||||
client = await connect_tcp("127.0.0.1", port, 0, 0, parser_map())
|
||||
await client.communicator.send(TestData(index=1, message="hello"))
|
||||
result = await asyncio.wait_for(received, 2.0)
|
||||
assert result.index == 1 and result.message == "hello"
|
||||
await client.close()
|
||||
await server.stop()
|
||||
```
|
||||
|
||||
### 중간 검증
|
||||
|
||||
```bash
|
||||
cd python && python -m pytest test/test_tcp.py -v
|
||||
```
|
||||
|
||||
예상: 2개 테스트 PASS
|
||||
|
||||
---
|
||||
|
||||
## [API-5] WebSocket Transport 구현
|
||||
|
||||
### 문제
|
||||
|
||||
Go의 `WsClient` / `WsServer`에 해당하는 asyncio WebSocket 구현이 없다.
|
||||
|
||||
### 해결 방법
|
||||
|
||||
`websockets` 라이브러리 사용. WebSocket은 메시지 경계가 보장되므로 길이 헤더 없이 바이너리 프레임으로 전송.
|
||||
|
||||
```python
|
||||
# ws_client.py
|
||||
import websockets
|
||||
|
||||
class WsClient(BaseClient):
|
||||
def __init__(self, ws: websockets.WebSocketClientProtocol | websockets.WebSocketServerProtocol,
|
||||
interval_sec: int, wait_sec: int, parser_map: ParserMap) -> None:
|
||||
super().__init__(interval_sec, wait_sec, do_close=self._do_close_impl)
|
||||
self._ws = ws
|
||||
self._init_base()
|
||||
self.communicator.initialize(self, parser_map)
|
||||
self.communicator.set_write_error_handler(lambda e: asyncio.create_task(self._on_disconnected()))
|
||||
self.communicator.add_listener(
|
||||
type_name_of(HeartBeat), lambda m: asyncio.create_task(self._on_heartbeat())
|
||||
)
|
||||
asyncio.create_task(self._read_loop())
|
||||
asyncio.create_task(self._send_heartbeat())
|
||||
|
||||
async def write_packet(self, base: PacketBase) -> None:
|
||||
await self._ws.send(base.SerializeToString())
|
||||
|
||||
async def _do_close_impl(self) -> None:
|
||||
await self._ws.close()
|
||||
|
||||
async def _read_loop(self) -> None:
|
||||
try:
|
||||
async for message in self._ws:
|
||||
if not isinstance(message, bytes):
|
||||
continue
|
||||
base = PacketBase()
|
||||
base.ParseFromString(message)
|
||||
self.communicator.on_received_data(
|
||||
base.typeName, base.data, base.nonce, base.responseNonce
|
||||
)
|
||||
asyncio.create_task(self._send_heartbeat())
|
||||
except Exception:
|
||||
pass
|
||||
finally:
|
||||
await self._on_disconnected()
|
||||
|
||||
async def connect_ws(host: str, port: int, path: str,
|
||||
interval_sec: int, wait_sec: int, parser_map: ParserMap) -> WsClient:
|
||||
uri = f"ws://{host}:{port}{path}"
|
||||
ws = await websockets.connect(uri)
|
||||
return WsClient(ws, interval_sec, wait_sec, parser_map)
|
||||
```
|
||||
|
||||
```python
|
||||
# ws_server.py
|
||||
class WsServer:
|
||||
def __init__(self, host: str, port: int, path: str, new_client: Callable) -> None: ...
|
||||
async def start(self) -> None: ...
|
||||
async def stop(self) -> None: ...
|
||||
def clients(self) -> list[WsClient]: ...
|
||||
async def broadcast(self, message) -> None: ...
|
||||
```
|
||||
|
||||
### 수정 파일 및 체크리스트
|
||||
|
||||
- [ ] `python/toki_socket/ws_client.py` 생성
|
||||
- [ ] `WsClient(BaseClient)` — websocket conn, read_loop, write_packet, do_close
|
||||
- [ ] `connect_ws(host, port, path, interval_sec, wait_sec, parser_map)` async 함수
|
||||
- [ ] `python/toki_socket/ws_server.py` 생성
|
||||
- [ ] `WsServer` — start/stop/clients/broadcast, path 기반 핸들러
|
||||
- [ ] `_handler(ws, path)` — WsClient 생성, disconnect_listener로 목록 관리
|
||||
|
||||
### 테스트 작성
|
||||
|
||||
skip — TCP 테스트로 핵심 로직이 검증되며, WS는 크로스 테스트에서 검증.
|
||||
|
||||
### 중간 검증
|
||||
|
||||
```bash
|
||||
cd python && python -c "from toki_socket.ws_client import WsClient; from toki_socket.ws_server import WsServer; print('OK')"
|
||||
```
|
||||
|
||||
예상: `OK`
|
||||
|
||||
---
|
||||
|
||||
## [API-6] 단위 테스트
|
||||
|
||||
API-2, API-4 항목에서 이미 기술. 해당 항목 체크리스트 참조.
|
||||
|
||||
---
|
||||
|
||||
## [API-7] Python×Go 크로스 테스트
|
||||
|
||||
### 문제
|
||||
|
||||
Python이 Go 서버와 통신할 수 있는지, Go가 Python 서버와 통신할 수 있는지 검증이 없다.
|
||||
|
||||
### 해결 방법
|
||||
|
||||
포트 할당:
|
||||
- Go 서버 / Python 클라이언트: TCP `29390`, WS `29392`
|
||||
- Python 서버 / Go 클라이언트: TCP `29490`, WS `29492`
|
||||
|
||||
**Python 클라이언트 헬퍼** `python/crosstest/go_python_client.py`:
|
||||
- `--mode tcp|ws`, `--port N`, `--phase send-push|requests` 플래그
|
||||
- `INFO typeName python=TestData` 출력
|
||||
- 각 시나리오 결과를 `PASS scenario=N detail=...` / `FAIL scenario=N error=...` 출력
|
||||
- send-push: 시나리오 1(fire-and-forget), 2(server push 수신)
|
||||
- requests: 시나리오 3(단일 요청), 4(동시 5개 요청 nonce 검증)
|
||||
- Go crosstest 클라이언트(`go/crosstest/dart_go_client/main.go`)와 동일한 구조
|
||||
|
||||
**Go 오케스트레이터** `go/crosstest/go_python.go`:
|
||||
- Go 서버를 시작하고 `python go_python_client.py` 서브프로세스를 실행
|
||||
- 기존 `go_dart.go` 패턴과 동일 (`//go:build ignore`, `runPythonClient()` 헬퍼)
|
||||
- Python 패키지 디렉터리를 소스 위치 기반으로 찾는 `pythonPackageDir()` 함수
|
||||
|
||||
**Python 오케스트레이터** `python/crosstest/python_go.py`:
|
||||
- Python 서버를 시작하고 `go run ./crosstest/python_go_client` 서브프로세스를 실행
|
||||
- `go/crosstest/dart_go_client/main.go` 패턴을 참고하여 `--mode`, `--port`, `--phase` 플래그 지원
|
||||
- Go 패키지 디렉터리를 소스 위치 기반으로 찾는 헬퍼
|
||||
|
||||
**Go 클라이언트 헬퍼** `go/crosstest/python_go_client/main.go`:
|
||||
- 기존 `dart_go_client/main.go`와 동일한 구조
|
||||
- Python 서버 포트(29490 TCP, 29492 WS)에 연결
|
||||
- `--phase send-push`: 시나리오 1, 2
|
||||
- `--phase requests`: 시나리오 3, 4
|
||||
- push 메시지 검증: `index=200 && message=="push from python server"`
|
||||
|
||||
**시나리오별 메시지 규격**:
|
||||
| 시나리오 | 클라이언트 송신 | 서버 검증 / 응답 |
|
||||
|----------|----------------|-----------------|
|
||||
| 1 (send) | `index=101, message="fire from python client"` | Go서버: index==101, message 검증 |
|
||||
| 2 (push) | — | Go서버 → `index=200, message="push from go server"` |
|
||||
| 3 (request) | `index=21, message="single request from python"` | Go서버 → `index=42, message="echo: single request from python"` |
|
||||
| 4 (concurrent) | `index=30+i, message="multi request {i} from python"` | Go서버 → `index=(30+i)*2, message="echo: ..."` |
|
||||
|
||||
Python서버일 때는 go클라이언트 메시지 검증:
|
||||
- 1: `index=101, message="fire from go client"`
|
||||
- 2: Python서버 → `index=200, message="push from python server"`
|
||||
|
||||
### 수정 파일 및 체크리스트
|
||||
|
||||
- [ ] `python/crosstest/__init__.py` 생성 (빈 파일)
|
||||
- [ ] `python/crosstest/go_python_client.py` 생성
|
||||
- [ ] argparse로 `--mode`, `--port`, `--phase` 파싱
|
||||
- [ ] `INFO typeName python=TestData` 출력
|
||||
- [ ] `dial_with_retry(mode, port)` — 3초 내 재시도
|
||||
- [ ] `run_send_push(client)` — 시나리오 1, 2
|
||||
- [ ] `run_requests(client)` — 시나리오 3, 4 (asyncio.gather로 동시 요청)
|
||||
- [ ] 비동기 main → `asyncio.run(main())`
|
||||
- [ ] `go/crosstest/go_python.go` 생성
|
||||
- [ ] `//go:build ignore` 태그
|
||||
- [ ] `goPythonTCPPort = 29390`, `goPythonWSPort = 29392`
|
||||
- [ ] `runTCPSendPush/TCPRequests/WSSendPush/WSRequests()` 4개 함수
|
||||
- [ ] `runPythonClient(mode, port, phase, expected)` — `python go_python_client.py` 실행
|
||||
- [ ] `pythonPackageDir()` — 소스 위치 기반 탐색
|
||||
- [ ] `python/crosstest/python_go.py` 생성
|
||||
- [ ] `pythonGoTCPPort = 29490`, `pythonGoWSPort = 29492`
|
||||
- [ ] `run_tcp_send_push/tcp_requests/ws_send_push/ws_requests()` 4개 함수
|
||||
- [ ] `run_go_client(mode, port, phase, expected)` — `go run ./crosstest/python_go_client` 실행
|
||||
- [ ] `go_package_dir()` — 소스 위치 기반 탐색
|
||||
- [ ] 비동기 main → `asyncio.run(main())`
|
||||
- [ ] `go/crosstest/python_go_client/main.go` 생성
|
||||
- [ ] `dart_go_client/main.go`와 동일 구조
|
||||
- [ ] Python서버 포트(`--port N`)에 연결
|
||||
- [ ] push 메시지: `"push from python server"`
|
||||
|
||||
### 테스트 작성
|
||||
|
||||
크로스 테스트가 곧 테스트. 별도 단위 테스트 skip.
|
||||
|
||||
### 중간 검증
|
||||
|
||||
Go 서버 / Python 클라이언트:
|
||||
```bash
|
||||
cd go && go run ./crosstest/go_python.go
|
||||
```
|
||||
|
||||
Python 서버 / Go 클라이언트:
|
||||
```bash
|
||||
cd python && python crosstest/python_go.py
|
||||
```
|
||||
|
||||
예상: 각각 `PASS all ...` 출력
|
||||
|
||||
---
|
||||
|
||||
## 수정 파일 요약
|
||||
|
||||
| 파일 | 항목 |
|
||||
|------|------|
|
||||
| `python/pyproject.toml` | API-1 |
|
||||
| `python/toki_socket/__init__.py` | API-1 |
|
||||
| `python/toki_socket/packets/__init__.py` | API-1 |
|
||||
| `python/toki_socket/packets/message_common.proto` | API-1 |
|
||||
| `python/toki_socket/packets/message_common_pb2.py` | API-1 |
|
||||
| `python/toki_socket/packets/message_common_pb2.pyi` | API-1 |
|
||||
| `python/README.md` | API-1 |
|
||||
| `python/toki_socket/communicator.py` | API-2 |
|
||||
| `python/test/__init__.py` | API-2 |
|
||||
| `python/test/test_communicator.py` | API-2 |
|
||||
| `python/toki_socket/base_client.py` | API-3 |
|
||||
| `python/toki_socket/tcp_client.py` | API-4 |
|
||||
| `python/toki_socket/tcp_server.py` | API-4 |
|
||||
| `python/test/test_tcp.py` | API-4 |
|
||||
| `python/toki_socket/ws_client.py` | API-5 |
|
||||
| `python/toki_socket/ws_server.py` | API-5 |
|
||||
| `python/crosstest/__init__.py` | API-7 |
|
||||
| `python/crosstest/go_python_client.py` | API-7 |
|
||||
| `python/crosstest/python_go.py` | API-7 |
|
||||
| `go/crosstest/go_python.go` | API-7 |
|
||||
| `go/crosstest/python_go_client/main.go` | API-7 |
|
||||
|
||||
---
|
||||
|
||||
## 최종 검증
|
||||
|
||||
```bash
|
||||
# 1. Python 단위 테스트
|
||||
cd python && pip install -e ".[dev]" && python -m pytest test/ -v
|
||||
|
||||
# 2. Go 서버 / Python 클라이언트 크로스 테스트
|
||||
cd go && go run ./crosstest/go_python.go
|
||||
|
||||
# 3. Python 서버 / Go 클라이언트 크로스 테스트
|
||||
cd python && python crosstest/python_go.py
|
||||
|
||||
# 4. 기존 Go 크로스 테스트 회귀 확인
|
||||
cd go && go run ./crosstest/go_dart.go
|
||||
cd go && go run ./crosstest/go_kotlin.go
|
||||
```
|
||||
|
||||
예상 결과:
|
||||
1. `pytest`: 모든 테스트 PASS
|
||||
2. `go run ./crosstest/go_python.go`: `PASS all go-server/python-client crosstests passed`
|
||||
3. `python crosstest/python_go.py`: `PASS all python-server/go-client crosstests passed`
|
||||
4. 기존 크로스 테스트: 회귀 없음 (PASS)
|
||||
326
go/crosstest/go_python.go
Normal file
326
go/crosstest/go_python.go
Normal file
|
|
@ -0,0 +1,326 @@
|
|||
//go:build ignore
|
||||
|
||||
package main
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"runtime"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"google.golang.org/protobuf/proto"
|
||||
"nhooyr.io/websocket"
|
||||
|
||||
toki "toki-labs.com/toki_socket/go"
|
||||
"toki-labs.com/toki_socket/go/packets"
|
||||
)
|
||||
|
||||
const (
|
||||
host = "127.0.0.1"
|
||||
goPythonTCPPort = 29390
|
||||
goPythonWSPort = 29392
|
||||
wsPath = "/"
|
||||
processTimeout = 20 * time.Second
|
||||
serverObservationWindow = 200 * time.Millisecond
|
||||
)
|
||||
|
||||
func parserMap() toki.ParserMap {
|
||||
return toki.ParserMap{
|
||||
toki.TypeNameOf(&packets.TestData{}): func(b []byte) (proto.Message, error) {
|
||||
m := &packets.TestData{}
|
||||
return m, proto.Unmarshal(b, m)
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func main() {
|
||||
fmt.Printf("INFO typeName go=%s\n", toki.TypeNameOf(&packets.TestData{}))
|
||||
if err := run(); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "FAIL crosstest error=%v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
fmt.Println("PASS all go-server/python-client crosstests passed")
|
||||
}
|
||||
|
||||
func run() error {
|
||||
if err := runTCPSendPush(); err != nil {
|
||||
return err
|
||||
}
|
||||
time.Sleep(150 * time.Millisecond)
|
||||
if err := runTCPRequests(); err != nil {
|
||||
return err
|
||||
}
|
||||
time.Sleep(150 * time.Millisecond)
|
||||
if err := runWSSendPush(); err != nil {
|
||||
return err
|
||||
}
|
||||
time.Sleep(150 * time.Millisecond)
|
||||
return runWSRequests()
|
||||
}
|
||||
|
||||
func runTCPSendPush() error {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
received := make(chan bool, 1)
|
||||
server := toki.NewTcpServer(host, goPythonTCPPort, func(conn net.Conn) *toki.TcpClient {
|
||||
return toki.NewTcpClient(conn, 0, 0, parserMap())
|
||||
})
|
||||
server.OnClientConnected = func(client *toki.TcpClient) {
|
||||
toki.AddListenerTyped[*packets.TestData](&client.Communicator, func(data *packets.TestData) {
|
||||
fmt.Printf("SERVER_RECEIVED index=%d message=%s\n", data.GetIndex(), data.GetMessage())
|
||||
valid := data.GetIndex() == 101 && data.GetMessage() == "fire from python client"
|
||||
select {
|
||||
case received <- valid:
|
||||
default:
|
||||
}
|
||||
if valid {
|
||||
_ = client.Send(&packets.TestData{Index: 200, Message: "push from go server"})
|
||||
}
|
||||
})
|
||||
}
|
||||
if err := server.Start(ctx); err != nil {
|
||||
return err
|
||||
}
|
||||
defer server.Stop()
|
||||
|
||||
if err := runPythonClient("tcp", goPythonTCPPort, "send-push", map[string]bool{"1": true, "2": true}); err != nil {
|
||||
return err
|
||||
}
|
||||
select {
|
||||
case ok := <-received:
|
||||
if !ok {
|
||||
return errors.New("TCP send-push server received unexpected data")
|
||||
}
|
||||
case <-time.After(serverObservationWindow):
|
||||
return errors.New("TCP send-push server did not receive expected data")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func runTCPRequests() error {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
server := toki.NewTcpServer(host, goPythonTCPPort, func(conn net.Conn) *toki.TcpClient {
|
||||
return toki.NewTcpClient(conn, 0, 0, parserMap())
|
||||
})
|
||||
server.OnClientConnected = func(client *toki.TcpClient) {
|
||||
toki.AddRequestListenerTyped[*packets.TestData, *packets.TestData](&client.Communicator, func(req *packets.TestData) (*packets.TestData, error) {
|
||||
return &packets.TestData{
|
||||
Index: req.GetIndex() * 2,
|
||||
Message: "echo: " + req.GetMessage(),
|
||||
}, nil
|
||||
})
|
||||
}
|
||||
if err := server.Start(ctx); err != nil {
|
||||
return err
|
||||
}
|
||||
defer server.Stop()
|
||||
|
||||
return runPythonClient("tcp", goPythonTCPPort, "requests", map[string]bool{"3": true, "4": true})
|
||||
}
|
||||
|
||||
func runWSSendPush() error {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
received := make(chan bool, 1)
|
||||
server := toki.NewWsServer(host, goPythonWSPort, wsPath, func(conn *websocket.Conn) *toki.WsClient {
|
||||
return toki.NewWsClient(conn, 0, 0, parserMap())
|
||||
})
|
||||
server.OnClientConnected = func(client *toki.WsClient) {
|
||||
toki.AddListenerTyped[*packets.TestData](&client.Communicator, func(data *packets.TestData) {
|
||||
fmt.Printf("SERVER_RECEIVED index=%d message=%s\n", data.GetIndex(), data.GetMessage())
|
||||
valid := data.GetIndex() == 101 && data.GetMessage() == "fire from python client"
|
||||
select {
|
||||
case received <- valid:
|
||||
default:
|
||||
}
|
||||
if valid {
|
||||
_ = client.Send(&packets.TestData{Index: 200, Message: "push from go server"})
|
||||
}
|
||||
})
|
||||
}
|
||||
if err := server.Start(ctx); err != nil {
|
||||
return err
|
||||
}
|
||||
defer server.Stop()
|
||||
|
||||
if err := runPythonClient("ws", goPythonWSPort, "send-push", map[string]bool{"1": true, "2": true}); err != nil {
|
||||
return err
|
||||
}
|
||||
select {
|
||||
case ok := <-received:
|
||||
if !ok {
|
||||
return errors.New("WS send-push server received unexpected data")
|
||||
}
|
||||
case <-time.After(serverObservationWindow):
|
||||
return errors.New("WS send-push server did not receive expected data")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func runWSRequests() error {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
server := toki.NewWsServer(host, goPythonWSPort, wsPath, func(conn *websocket.Conn) *toki.WsClient {
|
||||
return toki.NewWsClient(conn, 0, 0, parserMap())
|
||||
})
|
||||
server.OnClientConnected = func(client *toki.WsClient) {
|
||||
toki.AddRequestListenerTyped[*packets.TestData, *packets.TestData](&client.Communicator, func(req *packets.TestData) (*packets.TestData, error) {
|
||||
return &packets.TestData{
|
||||
Index: req.GetIndex() * 2,
|
||||
Message: "echo: " + req.GetMessage(),
|
||||
}, nil
|
||||
})
|
||||
}
|
||||
if err := server.Start(ctx); err != nil {
|
||||
return err
|
||||
}
|
||||
defer server.Stop()
|
||||
|
||||
return runPythonClient("ws", goPythonWSPort, "requests", map[string]bool{"3": true, "4": true})
|
||||
}
|
||||
|
||||
func runPythonClient(mode string, port int, phase string, expected map[string]bool) error {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), processTimeout)
|
||||
defer cancel()
|
||||
|
||||
pythonDir, err := pythonPackageDir()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
cmd := exec.CommandContext(ctx, "python3", "crosstest/go_python_client.py",
|
||||
"--mode="+mode,
|
||||
fmt.Sprintf("--port=%d", port),
|
||||
"--phase="+phase,
|
||||
)
|
||||
cmd.Dir = pythonDir
|
||||
|
||||
stdoutPipe, err := cmd.StdoutPipe()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
stderrPipe, err := cmd.StderrPipe()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := cmd.Start(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
var mu sync.Mutex
|
||||
var resultLines []string
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(2)
|
||||
go scanLines(&wg, stdoutPipe, os.Stdout, &mu, &resultLines)
|
||||
go scanLines(&wg, stderrPipe, os.Stderr, nil, nil)
|
||||
|
||||
waitErr := cmd.Wait()
|
||||
wg.Wait()
|
||||
if ctx.Err() == context.DeadlineExceeded {
|
||||
return fmt.Errorf("python client %s/%s timed out", mode, phase)
|
||||
}
|
||||
return validateResultLines("python-client "+mode+"/"+phase, waitErr, resultLines, expected)
|
||||
}
|
||||
|
||||
func pythonPackageDir() (string, error) {
|
||||
candidates := make([]string, 0, 3)
|
||||
|
||||
_, filename, _, ok := runtime.Caller(0)
|
||||
if ok {
|
||||
repoRoot := filepath.Dir(filepath.Dir(filepath.Dir(filename)))
|
||||
candidates = append(candidates, filepath.Join(repoRoot, "python"))
|
||||
}
|
||||
if wd, err := os.Getwd(); err == nil {
|
||||
candidates = append(candidates, findPythonPackageCandidates(wd)...)
|
||||
}
|
||||
if executable, err := os.Executable(); err == nil {
|
||||
candidates = append(candidates, findPythonPackageCandidates(filepath.Dir(executable))...)
|
||||
}
|
||||
|
||||
for _, candidate := range candidates {
|
||||
if isPythonPackageDir(candidate) {
|
||||
return candidate, nil
|
||||
}
|
||||
}
|
||||
return "", fmt.Errorf("cannot resolve python package directory from candidates %v", candidates)
|
||||
}
|
||||
|
||||
func findPythonPackageCandidates(start string) []string {
|
||||
candidates := make([]string, 0)
|
||||
for dir := start; ; dir = filepath.Dir(dir) {
|
||||
candidates = append(candidates, filepath.Join(dir, "python"))
|
||||
if filepath.Base(dir) == "python" {
|
||||
candidates = append(candidates, dir)
|
||||
}
|
||||
parent := filepath.Dir(dir)
|
||||
if parent == dir {
|
||||
return candidates
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func isPythonPackageDir(dir string) bool {
|
||||
info, err := os.Stat(filepath.Join(dir, "pyproject.toml"))
|
||||
return err == nil && !info.IsDir()
|
||||
}
|
||||
|
||||
func scanLines(wg *sync.WaitGroup, reader io.Reader, writer *os.File, mu *sync.Mutex, resultLines *[]string) {
|
||||
defer wg.Done()
|
||||
scanner := bufio.NewScanner(reader)
|
||||
for scanner.Scan() {
|
||||
line := scanner.Text()
|
||||
fmt.Fprintln(writer, line)
|
||||
if mu != nil && startsWithResult(line) {
|
||||
mu.Lock()
|
||||
*resultLines = append(*resultLines, line)
|
||||
mu.Unlock()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func validateResultLines(label string, waitErr error, lines []string, expected map[string]bool) error {
|
||||
failed := make([]string, 0)
|
||||
passed := make(map[string]bool)
|
||||
re := regexp.MustCompile(`scenario=([^ ]+)`)
|
||||
for _, line := range lines {
|
||||
if len(line) >= 5 && line[:5] == "FAIL " {
|
||||
failed = append(failed, line)
|
||||
continue
|
||||
}
|
||||
if len(line) >= 5 && line[:5] == "PASS " {
|
||||
match := re.FindStringSubmatch(line)
|
||||
if len(match) == 2 {
|
||||
passed[match[1]] = true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
missing := make([]string, 0)
|
||||
for scenario := range expected {
|
||||
if !passed[scenario] {
|
||||
missing = append(missing, scenario)
|
||||
}
|
||||
}
|
||||
if waitErr != nil || len(failed) > 0 || len(missing) > 0 {
|
||||
return fmt.Errorf("%s failed waitErr=%v failed=%v missing=%v", label, waitErr, failed, missing)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func startsWithResult(line string) bool {
|
||||
return (len(line) >= 5 && line[:5] == "PASS ") || (len(line) >= 5 && line[:5] == "FAIL ")
|
||||
}
|
||||
205
go/crosstest/python_go_client/main.go
Normal file
205
go/crosstest/python_go_client/main.go
Normal file
|
|
@ -0,0 +1,205 @@
|
|||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"flag"
|
||||
"fmt"
|
||||
"os"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"google.golang.org/protobuf/proto"
|
||||
|
||||
toki "toki-labs.com/toki_socket/go"
|
||||
"toki-labs.com/toki_socket/go/packets"
|
||||
)
|
||||
|
||||
const (
|
||||
host = "127.0.0.1"
|
||||
wsPath = "/"
|
||||
connectWindow = 3 * time.Second
|
||||
requestWindow = 2 * time.Second
|
||||
)
|
||||
|
||||
type clientHandle struct {
|
||||
communicator *toki.Communicator
|
||||
send func(proto.Message) error
|
||||
close func() error
|
||||
}
|
||||
|
||||
func parserMap() toki.ParserMap {
|
||||
return toki.ParserMap{
|
||||
toki.TypeNameOf(&packets.TestData{}): func(b []byte) (proto.Message, error) {
|
||||
m := &packets.TestData{}
|
||||
return m, proto.Unmarshal(b, m)
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func main() {
|
||||
mode := flag.String("mode", "tcp", "transport mode: tcp or ws")
|
||||
port := flag.Int("port", 0, "server port")
|
||||
phase := flag.String("phase", "send-push", "test phase: send-push or requests")
|
||||
flag.Parse()
|
||||
|
||||
fmt.Printf("INFO typeName go=%s\n", toki.TypeNameOf(&packets.TestData{}))
|
||||
|
||||
if *port == 0 {
|
||||
fail("setup", "port is required")
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
client, err := dialWithRetry(*mode, *port)
|
||||
if err != nil {
|
||||
fail("setup", err.Error())
|
||||
os.Exit(1)
|
||||
}
|
||||
defer client.close()
|
||||
|
||||
var ok bool
|
||||
switch *phase {
|
||||
case "send-push":
|
||||
ok = runSendPush(client)
|
||||
case "requests":
|
||||
ok = runRequests(client)
|
||||
default:
|
||||
fail("setup", fmt.Sprintf("unknown phase %q", *phase))
|
||||
ok = false
|
||||
}
|
||||
if !ok {
|
||||
os.Exit(1)
|
||||
}
|
||||
}
|
||||
|
||||
func dialWithRetry(mode string, port int) (*clientHandle, error) {
|
||||
deadline := time.Now().Add(connectWindow)
|
||||
var lastErr error
|
||||
for time.Now().Before(deadline) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 300*time.Millisecond)
|
||||
handle, err := dial(ctx, mode, port)
|
||||
cancel()
|
||||
if err == nil {
|
||||
return handle, nil
|
||||
}
|
||||
lastErr = err
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
}
|
||||
return nil, fmt.Errorf("connect %s:%d timed out: %w", mode, port, lastErr)
|
||||
}
|
||||
|
||||
func dial(ctx context.Context, mode string, port int) (*clientHandle, error) {
|
||||
switch mode {
|
||||
case "tcp":
|
||||
client, err := toki.DialTcp(ctx, host, port, 0, 0, parserMap())
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &clientHandle{
|
||||
communicator: &client.Communicator,
|
||||
send: client.Send,
|
||||
close: client.Close,
|
||||
}, nil
|
||||
case "ws":
|
||||
client, err := toki.DialWsWithHeartbeat(ctx, host, port, wsPath, 0, 0, parserMap())
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &clientHandle{
|
||||
communicator: &client.Communicator,
|
||||
send: client.Send,
|
||||
close: client.Close,
|
||||
}, nil
|
||||
default:
|
||||
return nil, fmt.Errorf("unknown mode %q", mode)
|
||||
}
|
||||
}
|
||||
|
||||
func runSendPush(client *clientHandle) bool {
|
||||
pushCh := make(chan *packets.TestData, 1)
|
||||
toki.AddListenerTyped[*packets.TestData](client.communicator, func(msg *packets.TestData) {
|
||||
pushCh <- msg
|
||||
})
|
||||
|
||||
err := client.send(&packets.TestData{
|
||||
Index: 101,
|
||||
Message: "fire from go client",
|
||||
})
|
||||
if err != nil {
|
||||
fail("1", err.Error())
|
||||
return false
|
||||
}
|
||||
pass("1", "fire-and-forget sent")
|
||||
|
||||
select {
|
||||
case msg := <-pushCh:
|
||||
if msg.GetIndex() != 200 || msg.GetMessage() != "push from python server" {
|
||||
fail("2", fmt.Sprintf("unexpected push index=%d message=%q", msg.GetIndex(), msg.GetMessage()))
|
||||
return false
|
||||
}
|
||||
pass("2", "received push from python server")
|
||||
return true
|
||||
case <-time.After(requestWindow):
|
||||
fail("2", "timeout waiting for server push")
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func runRequests(client *clientHandle) bool {
|
||||
res, err := toki.SendRequestTyped[*packets.TestData, *packets.TestData](
|
||||
client.communicator,
|
||||
&packets.TestData{Index: 21, Message: "single request from go"},
|
||||
requestWindow,
|
||||
)
|
||||
if err != nil {
|
||||
fail("3", err.Error())
|
||||
return false
|
||||
}
|
||||
if res.GetIndex() != 42 || res.GetMessage() != "echo: single request from go" {
|
||||
fail("3", fmt.Sprintf("unexpected response index=%d message=%q", res.GetIndex(), res.GetMessage()))
|
||||
return false
|
||||
}
|
||||
pass("3", "single request response matched")
|
||||
|
||||
const count = 5
|
||||
var wg sync.WaitGroup
|
||||
errCh := make(chan error, count)
|
||||
for i := 0; i < count; i++ {
|
||||
i := i
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
index := int32(30 + i)
|
||||
message := fmt.Sprintf("multi request %d from go", i)
|
||||
res, err := toki.SendRequestTyped[*packets.TestData, *packets.TestData](
|
||||
client.communicator,
|
||||
&packets.TestData{Index: index, Message: message},
|
||||
requestWindow,
|
||||
)
|
||||
if err != nil {
|
||||
errCh <- err
|
||||
return
|
||||
}
|
||||
if res.GetIndex() != index*2 || res.GetMessage() != "echo: "+message {
|
||||
errCh <- fmt.Errorf("request %d got index=%d message=%q", i, res.GetIndex(), res.GetMessage())
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
close(errCh)
|
||||
for err := range errCh {
|
||||
if err != nil {
|
||||
fail("4", err.Error())
|
||||
return false
|
||||
}
|
||||
}
|
||||
pass("4", "concurrent request responses matched")
|
||||
return true
|
||||
}
|
||||
|
||||
func pass(scenario, detail string) {
|
||||
fmt.Printf("PASS scenario=%s detail=%s\n", scenario, detail)
|
||||
}
|
||||
|
||||
func fail(scenario, detail string) {
|
||||
fmt.Printf("FAIL scenario=%s error=%s\n", scenario, detail)
|
||||
}
|
||||
35
python/README.md
Normal file
35
python/README.md
Normal file
|
|
@ -0,0 +1,35 @@
|
|||
# Toki Socket Python
|
||||
|
||||
Python 3.11+ implementation of the Toki Socket protocol.
|
||||
|
||||
## Supported Transports
|
||||
|
||||
| Transport | Status |
|
||||
|-----------|--------|
|
||||
| TCP | Supported |
|
||||
| WebSocket | Supported |
|
||||
| TLS+TCP | Not supported yet |
|
||||
| WSS | Not supported yet |
|
||||
|
||||
TLS and WSS are intentionally out of scope for this first Python port.
|
||||
|
||||
## Development
|
||||
|
||||
Install the package with test dependencies:
|
||||
|
||||
```bash
|
||||
python -m pip install -e ".[dev]"
|
||||
```
|
||||
|
||||
Run same-language tests:
|
||||
|
||||
```bash
|
||||
python -m pytest test/ -v
|
||||
```
|
||||
|
||||
Run Go/Python crosstests:
|
||||
|
||||
```bash
|
||||
python crosstest/python_go.py
|
||||
```
|
||||
|
||||
1
python/crosstest/__init__.py
Normal file
1
python/crosstest/__init__.py
Normal file
|
|
@ -0,0 +1 @@
|
|||
|
||||
140
python/crosstest/go_python_client.py
Normal file
140
python/crosstest/go_python_client.py
Normal file
|
|
@ -0,0 +1,140 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||
|
||||
from toki_socket.communicator import type_name_of
|
||||
from toki_socket.packets.message_common_pb2 import TestData
|
||||
from toki_socket.tcp_client import connect_tcp
|
||||
from toki_socket.ws_client import connect_ws
|
||||
|
||||
HOST = "127.0.0.1"
|
||||
WS_PATH = "/"
|
||||
CONNECT_WINDOW = 3.0
|
||||
REQUEST_WINDOW = 2.0
|
||||
|
||||
|
||||
def parser_map():
|
||||
return {type_name_of(TestData): TestData.FromString}
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--mode", choices=["tcp", "ws"], default="tcp")
|
||||
parser.add_argument("--port", type=int, required=True)
|
||||
parser.add_argument("--phase", choices=["send-push", "requests"], default="send-push")
|
||||
args = parser.parse_args()
|
||||
|
||||
print(f"INFO typeName python={type_name_of(TestData)}")
|
||||
|
||||
try:
|
||||
client = await dial_with_retry(args.mode, args.port)
|
||||
except Exception as exc:
|
||||
fail("setup", str(exc))
|
||||
raise SystemExit(1) from exc
|
||||
|
||||
try:
|
||||
if args.phase == "send-push":
|
||||
ok = await run_send_push(client)
|
||||
else:
|
||||
ok = await run_requests(client)
|
||||
finally:
|
||||
await client.close()
|
||||
|
||||
if not ok:
|
||||
raise SystemExit(1)
|
||||
|
||||
|
||||
async def dial_with_retry(mode: str, port: int):
|
||||
deadline = asyncio.get_running_loop().time() + CONNECT_WINDOW
|
||||
last_error: Exception | None = None
|
||||
while asyncio.get_running_loop().time() < deadline:
|
||||
try:
|
||||
return await asyncio.wait_for(dial(mode, port), 0.3)
|
||||
except Exception as exc:
|
||||
last_error = exc
|
||||
await asyncio.sleep(0.1)
|
||||
raise TimeoutError(f"connect {mode}:{port} timed out: {last_error}")
|
||||
|
||||
|
||||
async def dial(mode: str, port: int):
|
||||
if mode == "tcp":
|
||||
return await connect_tcp(HOST, port, 0, 0, parser_map())
|
||||
if mode == "ws":
|
||||
return await connect_ws(HOST, port, WS_PATH, 0, 0, parser_map())
|
||||
raise ValueError(f"unknown mode {mode!r}")
|
||||
|
||||
|
||||
async def run_send_push(client) -> bool:
|
||||
loop = asyncio.get_running_loop()
|
||||
push = loop.create_future()
|
||||
client.communicator.add_listener(
|
||||
type_name_of(TestData),
|
||||
lambda message: push.set_result(message) if not push.done() else None,
|
||||
)
|
||||
|
||||
try:
|
||||
await client.communicator.send(TestData(index=101, message="fire from python client"))
|
||||
passed("1", "fire-and-forget sent")
|
||||
|
||||
message = await asyncio.wait_for(push, REQUEST_WINDOW)
|
||||
if message.index != 200 or message.message != "push from go server":
|
||||
fail("2", f"unexpected push index={message.index} message={message.message!r}")
|
||||
return False
|
||||
passed("2", "received push from go server")
|
||||
return True
|
||||
except Exception as exc:
|
||||
fail("2", str(exc))
|
||||
return False
|
||||
|
||||
|
||||
async def run_requests(client) -> bool:
|
||||
try:
|
||||
single = await client.communicator.send_request(
|
||||
TestData(index=21, message="single request from python"),
|
||||
TestData,
|
||||
timeout=REQUEST_WINDOW,
|
||||
)
|
||||
if single.index != 42 or single.message != "echo: single request from python":
|
||||
fail("3", f"unexpected response index={single.index} message={single.message!r}")
|
||||
return False
|
||||
passed("3", "single request response matched")
|
||||
except Exception as exc:
|
||||
fail("3", str(exc))
|
||||
return False
|
||||
|
||||
async def request_one(i: int) -> None:
|
||||
index = 30 + i
|
||||
message = f"multi request {i} from python"
|
||||
res = await client.communicator.send_request(
|
||||
TestData(index=index, message=message),
|
||||
TestData,
|
||||
timeout=REQUEST_WINDOW,
|
||||
)
|
||||
if res.index != index * 2 or res.message != "echo: " + message:
|
||||
raise AssertionError(f"request {i} got index={res.index} message={res.message!r}")
|
||||
|
||||
try:
|
||||
await asyncio.gather(*(request_one(i) for i in range(5)))
|
||||
passed("4", "concurrent request responses matched")
|
||||
return True
|
||||
except Exception as exc:
|
||||
fail("4", str(exc))
|
||||
return False
|
||||
|
||||
|
||||
def passed(scenario: str, detail: str) -> None:
|
||||
print(f"PASS scenario={scenario} detail={detail}")
|
||||
|
||||
|
||||
def fail(scenario: str, error: str) -> None:
|
||||
print(f"FAIL scenario={scenario} error={error}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
|
||||
179
python/crosstest/python_go.py
Normal file
179
python/crosstest/python_go.py
Normal file
|
|
@ -0,0 +1,179 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import re
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||
|
||||
from toki_socket.communicator import type_name_of
|
||||
from toki_socket.packets.message_common_pb2 import TestData
|
||||
from toki_socket.tcp_server import TcpServer
|
||||
from toki_socket.ws_server import WsServer
|
||||
|
||||
HOST = "127.0.0.1"
|
||||
PYTHON_GO_TCP_PORT = 29490
|
||||
PYTHON_GO_WS_PORT = 29492
|
||||
WS_PATH = "/"
|
||||
PROCESS_TIMEOUT = 20.0
|
||||
SERVER_OBSERVATION_WINDOW = 0.2
|
||||
|
||||
|
||||
def parser_map():
|
||||
return {type_name_of(TestData): TestData.FromString}
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
print(f"INFO typeName python={type_name_of(TestData)}")
|
||||
try:
|
||||
await run_tcp_send_push()
|
||||
await asyncio.sleep(0.15)
|
||||
await run_tcp_requests()
|
||||
await asyncio.sleep(0.15)
|
||||
await run_ws_send_push()
|
||||
await asyncio.sleep(0.15)
|
||||
await run_ws_requests()
|
||||
except Exception as exc:
|
||||
print(f"FAIL crosstest error={exc}", file=sys.stderr)
|
||||
raise SystemExit(1) from exc
|
||||
print("PASS all python-server/go-client crosstests passed")
|
||||
|
||||
|
||||
async def run_tcp_send_push() -> None:
|
||||
received = asyncio.get_running_loop().create_future()
|
||||
server = TcpServer(HOST, PYTHON_GO_TCP_PORT, interval_sec=0, wait_sec=0, parser_map=parser_map())
|
||||
|
||||
def on_connected(client):
|
||||
def on_message(data):
|
||||
print(f"SERVER_RECEIVED index={data.index} message={data.message}")
|
||||
valid = data.index == 101 and data.message == "fire from go client"
|
||||
if not received.done():
|
||||
received.set_result(valid)
|
||||
if valid:
|
||||
asyncio.create_task(
|
||||
client.communicator.send(TestData(index=200, message="push from python server"))
|
||||
)
|
||||
|
||||
client.communicator.add_listener(type_name_of(TestData), on_message)
|
||||
|
||||
server.on_client_connected = on_connected
|
||||
await server.start()
|
||||
try:
|
||||
await run_go_client("tcp", PYTHON_GO_TCP_PORT, "send-push", {"1", "2"})
|
||||
ok = await asyncio.wait_for(received, SERVER_OBSERVATION_WINDOW)
|
||||
if not ok:
|
||||
raise RuntimeError("TCP send-push server received unexpected data")
|
||||
finally:
|
||||
await server.stop()
|
||||
|
||||
|
||||
async def run_tcp_requests() -> None:
|
||||
server = TcpServer(HOST, PYTHON_GO_TCP_PORT, interval_sec=0, wait_sec=0, parser_map=parser_map())
|
||||
server.on_client_connected = lambda client: client.communicator.add_request_listener(
|
||||
type_name_of(TestData),
|
||||
lambda req: TestData(index=req.index * 2, message="echo: " + req.message),
|
||||
)
|
||||
await server.start()
|
||||
try:
|
||||
await run_go_client("tcp", PYTHON_GO_TCP_PORT, "requests", {"3", "4"})
|
||||
finally:
|
||||
await server.stop()
|
||||
|
||||
|
||||
async def run_ws_send_push() -> None:
|
||||
received = asyncio.get_running_loop().create_future()
|
||||
server = WsServer(HOST, PYTHON_GO_WS_PORT, WS_PATH, interval_sec=0, wait_sec=0, parser_map=parser_map())
|
||||
|
||||
def on_connected(client):
|
||||
def on_message(data):
|
||||
print(f"SERVER_RECEIVED index={data.index} message={data.message}")
|
||||
valid = data.index == 101 and data.message == "fire from go client"
|
||||
if not received.done():
|
||||
received.set_result(valid)
|
||||
if valid:
|
||||
asyncio.create_task(
|
||||
client.communicator.send(TestData(index=200, message="push from python server"))
|
||||
)
|
||||
|
||||
client.communicator.add_listener(type_name_of(TestData), on_message)
|
||||
|
||||
server.on_client_connected = on_connected
|
||||
await server.start()
|
||||
try:
|
||||
await run_go_client("ws", PYTHON_GO_WS_PORT, "send-push", {"1", "2"})
|
||||
ok = await asyncio.wait_for(received, SERVER_OBSERVATION_WINDOW)
|
||||
if not ok:
|
||||
raise RuntimeError("WS send-push server received unexpected data")
|
||||
finally:
|
||||
await server.stop()
|
||||
|
||||
|
||||
async def run_ws_requests() -> None:
|
||||
server = WsServer(HOST, PYTHON_GO_WS_PORT, WS_PATH, interval_sec=0, wait_sec=0, parser_map=parser_map())
|
||||
server.on_client_connected = lambda client: client.communicator.add_request_listener(
|
||||
type_name_of(TestData),
|
||||
lambda req: TestData(index=req.index * 2, message="echo: " + req.message),
|
||||
)
|
||||
await server.start()
|
||||
try:
|
||||
await run_go_client("ws", PYTHON_GO_WS_PORT, "requests", {"3", "4"})
|
||||
finally:
|
||||
await server.stop()
|
||||
|
||||
|
||||
async def run_go_client(mode: str, port: int, phase: str, expected: set[str]) -> None:
|
||||
go_dir = go_package_dir()
|
||||
process = await asyncio.create_subprocess_exec(
|
||||
"go",
|
||||
"run",
|
||||
"./crosstest/python_go_client",
|
||||
f"--mode={mode}",
|
||||
f"--port={port}",
|
||||
f"--phase={phase}",
|
||||
cwd=go_dir,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
)
|
||||
try:
|
||||
stdout, stderr = await asyncio.wait_for(process.communicate(), PROCESS_TIMEOUT)
|
||||
except TimeoutError as exc:
|
||||
process.kill()
|
||||
await process.wait()
|
||||
raise TimeoutError(f"go client {mode}/{phase} timed out") from exc
|
||||
|
||||
result_lines: list[str] = []
|
||||
for line in stdout.decode().splitlines():
|
||||
print(line)
|
||||
if line.startswith(("PASS ", "FAIL ")):
|
||||
result_lines.append(line)
|
||||
for line in stderr.decode().splitlines():
|
||||
print(line, file=sys.stderr)
|
||||
|
||||
validate_result_lines(f"go-client {mode}/{phase}", process.returncode, result_lines, expected)
|
||||
|
||||
|
||||
def validate_result_lines(label: str, return_code: int | None, lines: list[str], expected: set[str]) -> None:
|
||||
failed = [line for line in lines if line.startswith("FAIL ")]
|
||||
passed: set[str] = set()
|
||||
for line in lines:
|
||||
if line.startswith("PASS "):
|
||||
match = re.search(r"scenario=([^ ]+)", line)
|
||||
if match:
|
||||
passed.add(match.group(1))
|
||||
missing = expected - passed
|
||||
if return_code or failed or missing:
|
||||
raise RuntimeError(f"{label} failed returnCode={return_code} failed={failed} missing={sorted(missing)}")
|
||||
|
||||
|
||||
def go_package_dir() -> Path:
|
||||
repo_root = Path(__file__).resolve().parents[2]
|
||||
candidate = repo_root / "go"
|
||||
if (candidate / "go.mod").is_file():
|
||||
return candidate
|
||||
raise RuntimeError(f"cannot resolve go package directory from {candidate}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
|
||||
28
python/pyproject.toml
Normal file
28
python/pyproject.toml
Normal file
|
|
@ -0,0 +1,28 @@
|
|||
[build-system]
|
||||
requires = ["setuptools>=68"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "toki-socket"
|
||||
version = "0.1.0"
|
||||
description = "Python implementation of the Toki Socket binary socket protocol."
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
dependencies = [
|
||||
"protobuf>=4.25",
|
||||
"websockets>=12.0",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
dev = [
|
||||
"pytest>=8.0",
|
||||
"pytest-asyncio>=0.23",
|
||||
]
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
where = ["."]
|
||||
include = ["toki_socket*"]
|
||||
|
||||
[tool.setuptools.package-data]
|
||||
"toki_socket.packets" = ["*.proto", "*.pyi"]
|
||||
|
||||
1
python/test/__init__.py
Normal file
1
python/test/__init__.py
Normal file
|
|
@ -0,0 +1 @@
|
|||
|
||||
101
python/test/test_communicator.py
Normal file
101
python/test/test_communicator.py
Normal file
|
|
@ -0,0 +1,101 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
|
||||
from toki_socket.communicator import Communicator, type_name_of
|
||||
from toki_socket.packets.message_common_pb2 import PacketBase, TestData as ProtoTestData
|
||||
|
||||
|
||||
class FakeTransport:
|
||||
def __init__(self) -> None:
|
||||
self.packets: list[PacketBase] = []
|
||||
self.closed = False
|
||||
|
||||
async def write_packet(self, base: PacketBase) -> None:
|
||||
self.packets.append(base)
|
||||
|
||||
async def close(self) -> None:
|
||||
self.closed = True
|
||||
|
||||
|
||||
def parser_map():
|
||||
return {type_name_of(ProtoTestData): ProtoTestData.FromString}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_listener_exclusivity():
|
||||
transport = FakeTransport()
|
||||
communicator = Communicator()
|
||||
communicator.initialize(transport, parser_map())
|
||||
communicator.add_listener(type_name_of(ProtoTestData), lambda _: None)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
communicator.add_request_listener(type_name_of(ProtoTestData), lambda req: req)
|
||||
|
||||
await communicator.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_and_receive():
|
||||
transport = FakeTransport()
|
||||
communicator = Communicator()
|
||||
communicator.initialize(transport, parser_map())
|
||||
received = asyncio.get_running_loop().create_future()
|
||||
communicator.add_listener(type_name_of(ProtoTestData), received.set_result)
|
||||
|
||||
await communicator.send(ProtoTestData(index=7, message="hello"))
|
||||
assert len(transport.packets) == 1
|
||||
packet = transport.packets[0]
|
||||
assert packet.typeName == "TestData"
|
||||
assert packet.nonce == 1
|
||||
|
||||
communicator.on_received_data(packet.typeName, packet.data, packet.nonce, 0)
|
||||
result = await asyncio.wait_for(received, 1)
|
||||
assert result.index == 7
|
||||
assert result.message == "hello"
|
||||
|
||||
await communicator.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_request_response():
|
||||
transport = FakeTransport()
|
||||
communicator = Communicator()
|
||||
communicator.initialize(transport, parser_map())
|
||||
|
||||
task = asyncio.create_task(
|
||||
communicator.send_request(ProtoTestData(index=5, message="request"), ProtoTestData, timeout=1)
|
||||
)
|
||||
for _ in range(20):
|
||||
if transport.packets:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
assert len(transport.packets) == 1
|
||||
request = transport.packets[0]
|
||||
response = ProtoTestData(index=10, message="echo: request")
|
||||
communicator.on_received_data(
|
||||
type_name_of(ProtoTestData),
|
||||
response.SerializeToString(),
|
||||
nonce=2,
|
||||
response_nonce=request.nonce,
|
||||
)
|
||||
|
||||
result = await task
|
||||
assert result.index == 10
|
||||
assert result.message == "echo: request"
|
||||
|
||||
await communicator.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_shutdown():
|
||||
transport = FakeTransport()
|
||||
communicator = Communicator()
|
||||
communicator.initialize(transport, parser_map())
|
||||
await communicator.close()
|
||||
|
||||
assert not communicator.is_alive()
|
||||
assert transport.closed
|
||||
66
python/test/test_tcp.py
Normal file
66
python/test/test_tcp.py
Normal file
|
|
@ -0,0 +1,66 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
|
||||
from toki_socket.communicator import type_name_of
|
||||
from toki_socket.packets.message_common_pb2 import TestData as ProtoTestData
|
||||
from toki_socket.tcp_client import TcpClient, connect_tcp
|
||||
from toki_socket.tcp_server import TcpServer
|
||||
|
||||
|
||||
def parser_map():
|
||||
return {type_name_of(ProtoTestData): ProtoTestData.FromString}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tcp_send_receive():
|
||||
received = asyncio.get_running_loop().create_future()
|
||||
server = TcpServer(
|
||||
"127.0.0.1",
|
||||
0,
|
||||
lambda reader, writer: TcpClient(reader, writer, 0, 0, parser_map()),
|
||||
)
|
||||
server.on_client_connected = lambda client: client.communicator.add_listener(
|
||||
type_name_of(ProtoTestData),
|
||||
lambda message: received.set_result(message) if not received.done() else None,
|
||||
)
|
||||
await server.start()
|
||||
client = await connect_tcp("127.0.0.1", server.port, 0, 0, parser_map())
|
||||
|
||||
await client.communicator.send(ProtoTestData(index=1, message="hello"))
|
||||
result = await asyncio.wait_for(received, 2)
|
||||
|
||||
assert result.index == 1
|
||||
assert result.message == "hello"
|
||||
|
||||
await client.close()
|
||||
await server.stop()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tcp_request_response():
|
||||
server = TcpServer(
|
||||
"127.0.0.1",
|
||||
0,
|
||||
lambda reader, writer: TcpClient(reader, writer, 0, 0, parser_map()),
|
||||
)
|
||||
server.on_client_connected = lambda client: client.communicator.add_request_listener(
|
||||
type_name_of(ProtoTestData),
|
||||
lambda req: ProtoTestData(index=req.index * 2, message="echo: " + req.message),
|
||||
)
|
||||
await server.start()
|
||||
client = await connect_tcp("127.0.0.1", server.port, 0, 0, parser_map())
|
||||
|
||||
result = await client.communicator.send_request(
|
||||
ProtoTestData(index=21, message="single request"),
|
||||
ProtoTestData,
|
||||
timeout=2,
|
||||
)
|
||||
|
||||
assert result.index == 42
|
||||
assert result.message == "echo: single request"
|
||||
|
||||
await client.close()
|
||||
await server.stop()
|
||||
29
python/toki_socket/__init__.py
Normal file
29
python/toki_socket/__init__.py
Normal file
|
|
@ -0,0 +1,29 @@
|
|||
from toki_socket.base_client import BaseClient
|
||||
from toki_socket.communicator import (
|
||||
Communicator,
|
||||
NotConnectedError,
|
||||
ParserMap,
|
||||
Transport,
|
||||
type_name_of,
|
||||
)
|
||||
from toki_socket.tcp_client import MAX_PACKET_SIZE, TcpClient, connect_tcp
|
||||
from toki_socket.tcp_server import TcpServer
|
||||
from toki_socket.ws_client import WsClient, connect_ws
|
||||
from toki_socket.ws_server import WsServer
|
||||
|
||||
__all__ = [
|
||||
"BaseClient",
|
||||
"Communicator",
|
||||
"MAX_PACKET_SIZE",
|
||||
"NotConnectedError",
|
||||
"ParserMap",
|
||||
"TcpClient",
|
||||
"TcpServer",
|
||||
"Transport",
|
||||
"WsClient",
|
||||
"WsServer",
|
||||
"connect_tcp",
|
||||
"connect_ws",
|
||||
"type_name_of",
|
||||
]
|
||||
|
||||
104
python/toki_socket/base_client.py
Normal file
104
python/toki_socket/base_client.py
Normal file
|
|
@ -0,0 +1,104 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import inspect
|
||||
from collections.abc import Awaitable, Callable
|
||||
from contextlib import suppress
|
||||
from typing import Any
|
||||
|
||||
from toki_socket.communicator import Communicator
|
||||
from toki_socket.packets.message_common_pb2 import HeartBeat
|
||||
|
||||
|
||||
class _HbTimer:
|
||||
def __init__(self, delay: float, callback: Callable[[], Awaitable[None]]) -> None:
|
||||
self._task = asyncio.create_task(self._run(delay, callback))
|
||||
|
||||
async def _run(self, delay: float, callback: Callable[[], Awaitable[None]]) -> None:
|
||||
await asyncio.sleep(delay)
|
||||
await callback()
|
||||
|
||||
def cancel(self) -> None:
|
||||
self._task.cancel()
|
||||
|
||||
|
||||
class BaseClient:
|
||||
def __init__(
|
||||
self,
|
||||
interval_sec: int,
|
||||
wait_sec: int,
|
||||
do_close: Callable[[], Awaitable[None]],
|
||||
) -> None:
|
||||
self.communicator = Communicator()
|
||||
self._interval_sec = interval_sec
|
||||
self._wait_sec = wait_sec
|
||||
self._do_close = do_close
|
||||
self._close_called = False
|
||||
self._close_lock = asyncio.Lock()
|
||||
self._disconnect_listeners: list[Callable[[Any], Any]] = []
|
||||
self._hb_timer: _HbTimer | None = None
|
||||
self._waiting_hb_response = False
|
||||
|
||||
def _init_base(self) -> None:
|
||||
self._close_lock = asyncio.Lock()
|
||||
|
||||
async def close(self) -> None:
|
||||
async with self._close_lock:
|
||||
if self._close_called:
|
||||
return
|
||||
self._close_called = True
|
||||
|
||||
self.communicator._shutdown()
|
||||
self._stop_heartbeat()
|
||||
await self._do_close()
|
||||
await self._notify_disconnected()
|
||||
|
||||
def add_disconnect_listener(self, fn: Callable[[Any], Any]) -> None:
|
||||
self._disconnect_listeners.append(fn)
|
||||
|
||||
def remove_disconnect_listeners(self) -> None:
|
||||
self._disconnect_listeners.clear()
|
||||
|
||||
async def _notify_disconnected(self) -> None:
|
||||
listeners = list(self._disconnect_listeners)
|
||||
self._disconnect_listeners.clear()
|
||||
for listener in listeners:
|
||||
result = listener(self)
|
||||
if inspect.isawaitable(result):
|
||||
await result
|
||||
|
||||
async def _on_disconnected(self) -> None:
|
||||
await self.close()
|
||||
|
||||
async def _send_heartbeat(self) -> None:
|
||||
if not self.communicator.is_alive() or self._interval_sec <= 0:
|
||||
return
|
||||
self._stop_heartbeat()
|
||||
self._hb_timer = _HbTimer(self._interval_sec, self._heartbeat_interval_elapsed)
|
||||
|
||||
async def _heartbeat_interval_elapsed(self) -> None:
|
||||
if not self.communicator.is_alive():
|
||||
return
|
||||
self._waiting_hb_response = True
|
||||
with suppress(Exception):
|
||||
await self.communicator.send(HeartBeat())
|
||||
self._stop_heartbeat()
|
||||
if self._wait_sec > 0:
|
||||
self._hb_timer = _HbTimer(self._wait_sec, self._heartbeat_wait_elapsed)
|
||||
|
||||
async def _heartbeat_wait_elapsed(self) -> None:
|
||||
if self.communicator.is_alive():
|
||||
await self._on_disconnected()
|
||||
|
||||
async def _on_heartbeat(self) -> None:
|
||||
if self._waiting_hb_response:
|
||||
self._waiting_hb_response = False
|
||||
return
|
||||
with suppress(Exception):
|
||||
await self.communicator.send(HeartBeat())
|
||||
|
||||
def _stop_heartbeat(self) -> None:
|
||||
if self._hb_timer is not None:
|
||||
self._hb_timer.cancel()
|
||||
self._hb_timer = None
|
||||
|
||||
293
python/toki_socket/communicator.py
Normal file
293
python/toki_socket/communicator.py
Normal file
|
|
@ -0,0 +1,293 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import inspect
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Any, Protocol, TypeAlias, TypeVar, runtime_checkable
|
||||
|
||||
from google.protobuf.message import Message
|
||||
|
||||
from toki_socket.packets.message_common_pb2 import HeartBeat, PacketBase
|
||||
|
||||
T = TypeVar("T", bound=Message)
|
||||
ParserMap: TypeAlias = dict[str, Callable[[bytes], Message]]
|
||||
RequestHandler: TypeAlias = Callable[..., Message | Awaitable[Message | None] | None]
|
||||
|
||||
_STOP = object()
|
||||
|
||||
|
||||
class NotConnectedError(RuntimeError):
|
||||
def __init__(self) -> None:
|
||||
super().__init__("not connected")
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class Transport(Protocol):
|
||||
async def write_packet(self, base: PacketBase) -> None: ...
|
||||
|
||||
async def close(self) -> None: ...
|
||||
|
||||
|
||||
class Communicator:
|
||||
def __init__(self) -> None:
|
||||
self._nonce = 0
|
||||
self._is_alive = False
|
||||
self._parser_map: ParserMap = {}
|
||||
self._handlers: dict[str, list[Callable[[Message], Any]]] = {}
|
||||
self._req_handlers: dict[str, RequestHandler] = {}
|
||||
self._pending_requests: dict[int, tuple[str, asyncio.Future[Message]]] = {}
|
||||
self._write_queue: asyncio.Queue[Any] = asyncio.Queue(maxsize=64)
|
||||
self._closed = asyncio.Event()
|
||||
self._transport: Transport | None = None
|
||||
self._write_loop_task: asyncio.Task[None] | None = None
|
||||
self._write_error_handler: Callable[[Exception], Any] | None = None
|
||||
|
||||
def initialize(self, transport: Transport, parser_map: ParserMap) -> None:
|
||||
self._transport = transport
|
||||
self._parser_map = dict(parser_map)
|
||||
self._parser_map[type_name_of(HeartBeat)] = HeartBeat.FromString
|
||||
self._handlers = {}
|
||||
self._req_handlers = {}
|
||||
self._pending_requests = {}
|
||||
self._write_queue = asyncio.Queue(maxsize=64)
|
||||
self._closed = asyncio.Event()
|
||||
self._is_alive = True
|
||||
self._write_loop_task = asyncio.create_task(self._write_loop())
|
||||
|
||||
def is_alive(self) -> bool:
|
||||
return self._is_alive
|
||||
|
||||
def set_write_error_handler(self, fn: Callable[[Exception], Any]) -> None:
|
||||
self._write_error_handler = fn
|
||||
|
||||
def next_nonce(self) -> int:
|
||||
self._nonce += 1
|
||||
return self._nonce
|
||||
|
||||
def _shutdown(self) -> None:
|
||||
if not self._is_alive and self._closed.is_set():
|
||||
return
|
||||
self._is_alive = False
|
||||
self._closed.set()
|
||||
for _, pending in list(self._pending_requests.values()):
|
||||
if not pending.done():
|
||||
pending.set_exception(NotConnectedError())
|
||||
self._pending_requests.clear()
|
||||
try:
|
||||
self._write_queue.put_nowait(_STOP)
|
||||
except asyncio.QueueFull:
|
||||
if self._write_loop_task is not None:
|
||||
self._write_loop_task.cancel()
|
||||
|
||||
async def close(self) -> None:
|
||||
self._shutdown()
|
||||
if self._transport is not None:
|
||||
await self._transport.close()
|
||||
|
||||
async def wait_closed(self) -> None:
|
||||
await self._closed.wait()
|
||||
|
||||
async def _write_loop(self) -> None:
|
||||
while True:
|
||||
item = await self._write_queue.get()
|
||||
if item is _STOP:
|
||||
return
|
||||
base, done = item
|
||||
try:
|
||||
if self._transport is None:
|
||||
raise NotConnectedError()
|
||||
await self._transport.write_packet(base)
|
||||
if not done.done():
|
||||
done.set_result(None)
|
||||
except Exception as exc:
|
||||
if not done.done():
|
||||
done.set_exception(exc)
|
||||
if self._write_error_handler is not None:
|
||||
self._write_error_handler(exc)
|
||||
|
||||
async def queue_packet(self, base: PacketBase) -> None:
|
||||
if not self.is_alive():
|
||||
raise NotConnectedError()
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
done: asyncio.Future[None] = loop.create_future()
|
||||
item = (base, done)
|
||||
try:
|
||||
self._write_queue.put_nowait(item)
|
||||
except asyncio.QueueFull:
|
||||
closed_task = asyncio.create_task(self._closed.wait())
|
||||
put_task = asyncio.create_task(self._write_queue.put(item))
|
||||
finished, pending = await asyncio.wait(
|
||||
{closed_task, put_task},
|
||||
return_when=asyncio.FIRST_COMPLETED,
|
||||
)
|
||||
for task in pending:
|
||||
task.cancel()
|
||||
if closed_task in finished:
|
||||
raise NotConnectedError()
|
||||
|
||||
closed_task = asyncio.create_task(self._closed.wait())
|
||||
finished, pending = await asyncio.wait(
|
||||
{done, closed_task},
|
||||
return_when=asyncio.FIRST_COMPLETED,
|
||||
)
|
||||
for task in pending:
|
||||
task.cancel()
|
||||
if closed_task in finished and not done.done():
|
||||
raise NotConnectedError()
|
||||
await done
|
||||
|
||||
async def send(self, message: Message) -> None:
|
||||
if not self.is_alive():
|
||||
raise NotConnectedError()
|
||||
await self.queue_packet(
|
||||
PacketBase(
|
||||
typeName=type_name_of(message),
|
||||
nonce=self.next_nonce(),
|
||||
data=message.SerializeToString(),
|
||||
)
|
||||
)
|
||||
|
||||
async def send_request(
|
||||
self,
|
||||
req: Message,
|
||||
res_type: type[T] | T,
|
||||
timeout: float = 30.0,
|
||||
) -> T:
|
||||
if not self.is_alive():
|
||||
raise NotConnectedError()
|
||||
request_nonce = self.next_nonce()
|
||||
expected_type_name = type_name_of(res_type)
|
||||
loop = asyncio.get_running_loop()
|
||||
pending: asyncio.Future[Message] = loop.create_future()
|
||||
self._pending_requests[request_nonce] = (expected_type_name, pending)
|
||||
|
||||
try:
|
||||
await self.queue_packet(
|
||||
PacketBase(
|
||||
typeName=type_name_of(req),
|
||||
nonce=request_nonce,
|
||||
data=req.SerializeToString(),
|
||||
)
|
||||
)
|
||||
timeout = timeout if timeout > 0 else 30.0
|
||||
return await asyncio.wait_for(pending, timeout)
|
||||
except TimeoutError as exc:
|
||||
raise TimeoutError(f"request timeout for nonce {request_nonce}") from exc
|
||||
finally:
|
||||
self._pending_requests.pop(request_nonce, None)
|
||||
|
||||
def add_listener(self, type_name: str, fn: Callable[[Message], Any]) -> None:
|
||||
if type_name in self._req_handlers:
|
||||
raise ValueError(f"type {type_name} is already registered with add_request_listener")
|
||||
self._handlers.setdefault(type_name, []).append(fn)
|
||||
|
||||
def remove_listeners(self, type_name: str) -> None:
|
||||
self._handlers.pop(type_name, None)
|
||||
|
||||
def add_request_listener(self, type_name: str, fn: RequestHandler) -> None:
|
||||
if self._handlers.get(type_name):
|
||||
raise ValueError(f"type {type_name} is already registered with add_listener")
|
||||
if type_name in self._req_handlers:
|
||||
raise ValueError(f"type {type_name} is already registered with add_request_listener")
|
||||
self._req_handlers[type_name] = fn
|
||||
|
||||
def on_received_data(
|
||||
self,
|
||||
type_name: str,
|
||||
data: bytes,
|
||||
nonce: int = 0,
|
||||
response_nonce: int = 0,
|
||||
) -> None:
|
||||
if response_nonce > 0:
|
||||
self._handle_response(type_name, data, response_nonce)
|
||||
return
|
||||
|
||||
req_handler = self._req_handlers.get(type_name)
|
||||
listeners = list(self._handlers.get(type_name, []))
|
||||
|
||||
if req_handler is not None:
|
||||
try:
|
||||
message = self._parse(type_name, data)
|
||||
except Exception:
|
||||
return
|
||||
asyncio.create_task(self._run_request_handler(req_handler, message, nonce))
|
||||
return
|
||||
|
||||
if not listeners:
|
||||
return
|
||||
try:
|
||||
message = self._parse(type_name, data)
|
||||
except Exception:
|
||||
return
|
||||
for listener in listeners:
|
||||
listener(message)
|
||||
|
||||
async def _run_request_handler(
|
||||
self,
|
||||
handler: RequestHandler,
|
||||
message: Message,
|
||||
request_nonce: int,
|
||||
) -> None:
|
||||
try:
|
||||
result = handler(message, request_nonce) if _accepts_nonce(handler) else handler(message)
|
||||
if inspect.isawaitable(result):
|
||||
result = await result
|
||||
if result is None or not self.is_alive():
|
||||
return
|
||||
await self.queue_packet(
|
||||
PacketBase(
|
||||
typeName=type_name_of(result),
|
||||
nonce=self.next_nonce(),
|
||||
responseNonce=request_nonce,
|
||||
data=result.SerializeToString(),
|
||||
)
|
||||
)
|
||||
except Exception:
|
||||
return
|
||||
|
||||
def _handle_response(self, type_name: str, data: bytes, response_nonce: int) -> None:
|
||||
item = self._pending_requests.pop(response_nonce, None)
|
||||
if item is None:
|
||||
return
|
||||
expected_type_name, pending = item
|
||||
if type_name != expected_type_name:
|
||||
pending.set_exception(
|
||||
ValueError(
|
||||
f"response type mismatch for nonce {response_nonce}: "
|
||||
f"expected {expected_type_name}, got {type_name}"
|
||||
)
|
||||
)
|
||||
return
|
||||
try:
|
||||
pending.set_result(self._parse(type_name, data))
|
||||
except Exception as exc:
|
||||
pending.set_exception(exc)
|
||||
|
||||
def _parse(self, type_name: str, data: bytes) -> Message:
|
||||
parser = self._parser_map.get(type_name)
|
||||
if parser is None:
|
||||
raise ValueError(f"protobuf parser is not registered for type {type_name}")
|
||||
return parser(data)
|
||||
|
||||
|
||||
def type_name_of(message_or_class: Message | type[Message]) -> str:
|
||||
descriptor = getattr(message_or_class, "DESCRIPTOR", None)
|
||||
if descriptor is None:
|
||||
descriptor = message_or_class.__class__.DESCRIPTOR
|
||||
return descriptor.name
|
||||
|
||||
|
||||
def _accepts_nonce(handler: RequestHandler) -> bool:
|
||||
try:
|
||||
signature = inspect.signature(handler)
|
||||
except (TypeError, ValueError):
|
||||
return False
|
||||
positional = [
|
||||
param
|
||||
for param in signature.parameters.values()
|
||||
if param.kind
|
||||
in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD)
|
||||
]
|
||||
return len(positional) >= 2
|
||||
|
||||
4
python/toki_socket/packets/__init__.py
Normal file
4
python/toki_socket/packets/__init__.py
Normal file
|
|
@ -0,0 +1,4 @@
|
|||
from toki_socket.packets.message_common_pb2 import HeartBeat, PacketBase, TestData
|
||||
|
||||
__all__ = ["HeartBeat", "PacketBase", "TestData"]
|
||||
|
||||
16
python/toki_socket/packets/message_common.proto
Normal file
16
python/toki_socket/packets/message_common.proto
Normal file
|
|
@ -0,0 +1,16 @@
|
|||
syntax = "proto3";
|
||||
|
||||
message PacketBase {
|
||||
string typeName = 1;
|
||||
int32 nonce = 2;
|
||||
bytes data = 3;
|
||||
int32 responseNonce = 4;
|
||||
}
|
||||
|
||||
message HeartBeat {}
|
||||
|
||||
message TestData {
|
||||
int32 index = 1;
|
||||
string message = 2;
|
||||
}
|
||||
|
||||
40
python/toki_socket/packets/message_common_pb2.py
Normal file
40
python/toki_socket/packets/message_common_pb2.py
Normal file
|
|
@ -0,0 +1,40 @@
|
|||
# -*- coding: utf-8 -*-
|
||||
# Generated by the protocol buffer compiler. DO NOT EDIT!
|
||||
# NO CHECKED-IN PROTOBUF GENCODE
|
||||
# source: message_common.proto
|
||||
# Protobuf Python Version: 5.29.3
|
||||
"""Generated protocol buffer code."""
|
||||
from google.protobuf import descriptor as _descriptor
|
||||
from google.protobuf import descriptor_pool as _descriptor_pool
|
||||
from google.protobuf import runtime_version as _runtime_version
|
||||
from google.protobuf import symbol_database as _symbol_database
|
||||
from google.protobuf.internal import builder as _builder
|
||||
_runtime_version.ValidateProtobufRuntimeVersion(
|
||||
_runtime_version.Domain.PUBLIC,
|
||||
5,
|
||||
29,
|
||||
3,
|
||||
'',
|
||||
'message_common.proto'
|
||||
)
|
||||
# @@protoc_insertion_point(imports)
|
||||
|
||||
_sym_db = _symbol_database.Default()
|
||||
|
||||
|
||||
|
||||
|
||||
DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x14message_common.proto\"R\n\nPacketBase\x12\x10\n\x08typeName\x18\x01 \x01(\t\x12\r\n\x05nonce\x18\x02 \x01(\x05\x12\x0c\n\x04\x64\x61ta\x18\x03 \x01(\x0c\x12\x15\n\rresponseNonce\x18\x04 \x01(\x05\"\x0b\n\tHeartBeat\"*\n\x08TestData\x12\r\n\x05index\x18\x01 \x01(\x05\x12\x0f\n\x07message\x18\x02 \x01(\tb\x06proto3')
|
||||
|
||||
_globals = globals()
|
||||
_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
|
||||
_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'message_common_pb2', _globals)
|
||||
if not _descriptor._USE_C_DESCRIPTORS:
|
||||
DESCRIPTOR._loaded_options = None
|
||||
_globals['_PACKETBASE']._serialized_start=24
|
||||
_globals['_PACKETBASE']._serialized_end=106
|
||||
_globals['_HEARTBEAT']._serialized_start=108
|
||||
_globals['_HEARTBEAT']._serialized_end=119
|
||||
_globals['_TESTDATA']._serialized_start=121
|
||||
_globals['_TESTDATA']._serialized_end=163
|
||||
# @@protoc_insertion_point(module_scope)
|
||||
29
python/toki_socket/packets/message_common_pb2.pyi
Normal file
29
python/toki_socket/packets/message_common_pb2.pyi
Normal file
|
|
@ -0,0 +1,29 @@
|
|||
from google.protobuf import descriptor as _descriptor
|
||||
from google.protobuf import message as _message
|
||||
from typing import ClassVar as _ClassVar, Optional as _Optional
|
||||
|
||||
DESCRIPTOR: _descriptor.FileDescriptor
|
||||
|
||||
class PacketBase(_message.Message):
|
||||
__slots__ = ("typeName", "nonce", "data", "responseNonce")
|
||||
TYPENAME_FIELD_NUMBER: _ClassVar[int]
|
||||
NONCE_FIELD_NUMBER: _ClassVar[int]
|
||||
DATA_FIELD_NUMBER: _ClassVar[int]
|
||||
RESPONSENONCE_FIELD_NUMBER: _ClassVar[int]
|
||||
typeName: str
|
||||
nonce: int
|
||||
data: bytes
|
||||
responseNonce: int
|
||||
def __init__(self, typeName: _Optional[str] = ..., nonce: _Optional[int] = ..., data: _Optional[bytes] = ..., responseNonce: _Optional[int] = ...) -> None: ...
|
||||
|
||||
class HeartBeat(_message.Message):
|
||||
__slots__ = ()
|
||||
def __init__(self) -> None: ...
|
||||
|
||||
class TestData(_message.Message):
|
||||
__slots__ = ("index", "message")
|
||||
INDEX_FIELD_NUMBER: _ClassVar[int]
|
||||
MESSAGE_FIELD_NUMBER: _ClassVar[int]
|
||||
index: int
|
||||
message: str
|
||||
def __init__(self, index: _Optional[int] = ..., message: _Optional[str] = ...) -> None: ...
|
||||
84
python/toki_socket/tcp_client.py
Normal file
84
python/toki_socket/tcp_client.py
Normal file
|
|
@ -0,0 +1,84 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import struct
|
||||
from contextlib import suppress
|
||||
|
||||
from google.protobuf.message import DecodeError
|
||||
|
||||
from toki_socket.base_client import BaseClient
|
||||
from toki_socket.communicator import ParserMap, type_name_of
|
||||
from toki_socket.packets.message_common_pb2 import HeartBeat, PacketBase
|
||||
|
||||
MAX_PACKET_SIZE = 64 << 20
|
||||
|
||||
|
||||
class TcpClient(BaseClient):
|
||||
def __init__(
|
||||
self,
|
||||
reader: asyncio.StreamReader,
|
||||
writer: asyncio.StreamWriter,
|
||||
interval_sec: int,
|
||||
wait_sec: int,
|
||||
parser_map: ParserMap,
|
||||
) -> None:
|
||||
super().__init__(interval_sec, wait_sec, self._do_close_impl)
|
||||
self._reader = reader
|
||||
self._writer = writer
|
||||
self._init_base()
|
||||
self.communicator.initialize(self, parser_map)
|
||||
self.communicator.set_write_error_handler(lambda _: asyncio.create_task(self._on_disconnected()))
|
||||
self.communicator.add_listener(
|
||||
type_name_of(HeartBeat),
|
||||
lambda _: asyncio.create_task(self._on_heartbeat()),
|
||||
)
|
||||
self._read_task = asyncio.create_task(self._read_loop())
|
||||
asyncio.create_task(self._send_heartbeat())
|
||||
|
||||
async def write_packet(self, base: PacketBase) -> None:
|
||||
data = base.SerializeToString()
|
||||
if len(data) > MAX_PACKET_SIZE:
|
||||
raise ValueError("packet exceeds maximum size")
|
||||
self._writer.write(struct.pack(">I", len(data)) + data)
|
||||
await self._writer.drain()
|
||||
|
||||
async def _do_close_impl(self) -> None:
|
||||
self._writer.close()
|
||||
with suppress(Exception):
|
||||
await self._writer.wait_closed()
|
||||
|
||||
async def _read_loop(self) -> None:
|
||||
while self.communicator.is_alive():
|
||||
try:
|
||||
header = await self._reader.readexactly(4)
|
||||
length = struct.unpack(">I", header)[0]
|
||||
if length == 0:
|
||||
continue
|
||||
if length > MAX_PACKET_SIZE:
|
||||
await self._on_disconnected()
|
||||
return
|
||||
data = await self._reader.readexactly(length)
|
||||
base = PacketBase()
|
||||
base.ParseFromString(data)
|
||||
self.communicator.on_received_data(
|
||||
base.typeName,
|
||||
base.data,
|
||||
base.nonce,
|
||||
base.responseNonce,
|
||||
)
|
||||
asyncio.create_task(self._send_heartbeat())
|
||||
except (asyncio.IncompleteReadError, ConnectionError, OSError, DecodeError):
|
||||
await self._on_disconnected()
|
||||
return
|
||||
|
||||
|
||||
async def connect_tcp(
|
||||
host: str,
|
||||
port: int,
|
||||
interval_sec: int,
|
||||
wait_sec: int,
|
||||
parser_map: ParserMap,
|
||||
) -> TcpClient:
|
||||
reader, writer = await asyncio.open_connection(host, port)
|
||||
return TcpClient(reader, writer, interval_sec, wait_sec, parser_map)
|
||||
|
||||
80
python/toki_socket/tcp_server.py
Normal file
80
python/toki_socket/tcp_server.py
Normal file
|
|
@ -0,0 +1,80 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import inspect
|
||||
from collections.abc import Callable
|
||||
|
||||
from google.protobuf.message import Message
|
||||
|
||||
from toki_socket.communicator import ParserMap
|
||||
from toki_socket.tcp_client import TcpClient
|
||||
|
||||
|
||||
class TcpServer:
|
||||
def __init__(
|
||||
self,
|
||||
host: str,
|
||||
port: int,
|
||||
new_client: Callable[[asyncio.StreamReader, asyncio.StreamWriter], TcpClient] | None = None,
|
||||
interval_sec: int = 0,
|
||||
wait_sec: int = 0,
|
||||
parser_map: ParserMap | None = None,
|
||||
) -> None:
|
||||
self._host = host
|
||||
self._port = port
|
||||
self._new_client = new_client
|
||||
self._interval_sec = interval_sec
|
||||
self._wait_sec = wait_sec
|
||||
self._parser_map = parser_map
|
||||
self._clients: list[TcpClient] = []
|
||||
self._server: asyncio.Server | None = None
|
||||
self.on_client_connected: Callable[[TcpClient], object] = lambda _: None
|
||||
|
||||
@property
|
||||
def port(self) -> int:
|
||||
if self._server is None or not self._server.sockets:
|
||||
return self._port
|
||||
return int(self._server.sockets[0].getsockname()[1])
|
||||
|
||||
def clients(self) -> list[TcpClient]:
|
||||
return list(self._clients)
|
||||
|
||||
async def start(self) -> None:
|
||||
self._server = await asyncio.start_server(self._handle_connection, self._host, self._port)
|
||||
|
||||
async def stop(self) -> None:
|
||||
if self._server is not None:
|
||||
self._server.close()
|
||||
await self._server.wait_closed()
|
||||
self._server = None
|
||||
clients = list(self._clients)
|
||||
self._clients.clear()
|
||||
for client in clients:
|
||||
await client.close()
|
||||
|
||||
async def broadcast(self, message: Message) -> None:
|
||||
for client in self.clients():
|
||||
await client.communicator.send(message)
|
||||
|
||||
async def _handle_connection(
|
||||
self,
|
||||
reader: asyncio.StreamReader,
|
||||
writer: asyncio.StreamWriter,
|
||||
) -> None:
|
||||
if self._new_client is not None:
|
||||
client = self._new_client(reader, writer)
|
||||
else:
|
||||
if self._parser_map is None:
|
||||
raise ValueError("parser_map is required when new_client is not provided")
|
||||
client = TcpClient(reader, writer, self._interval_sec, self._wait_sec, self._parser_map)
|
||||
|
||||
self._clients.append(client)
|
||||
client.add_disconnect_listener(lambda c: self._remove_client(c))
|
||||
result = self.on_client_connected(client)
|
||||
if inspect.isawaitable(result):
|
||||
await result
|
||||
|
||||
def _remove_client(self, client: TcpClient) -> None:
|
||||
if client in self._clients:
|
||||
self._clients.remove(client)
|
||||
|
||||
80
python/toki_socket/ws_client.py
Normal file
80
python/toki_socket/ws_client.py
Normal file
|
|
@ -0,0 +1,80 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from contextlib import suppress
|
||||
from typing import Any
|
||||
|
||||
from google.protobuf.message import DecodeError
|
||||
from websockets.exceptions import ConnectionClosed
|
||||
|
||||
try:
|
||||
from websockets.asyncio.client import connect as ws_connect
|
||||
except ImportError: # websockets 12.x legacy API
|
||||
from websockets import connect as ws_connect
|
||||
|
||||
from toki_socket.base_client import BaseClient
|
||||
from toki_socket.communicator import ParserMap, type_name_of
|
||||
from toki_socket.packets.message_common_pb2 import HeartBeat, PacketBase
|
||||
|
||||
|
||||
class WsClient(BaseClient):
|
||||
def __init__(
|
||||
self,
|
||||
ws: Any,
|
||||
interval_sec: int,
|
||||
wait_sec: int,
|
||||
parser_map: ParserMap,
|
||||
) -> None:
|
||||
super().__init__(interval_sec, wait_sec, self._do_close_impl)
|
||||
self._ws = ws
|
||||
self._init_base()
|
||||
self.communicator.initialize(self, parser_map)
|
||||
self.communicator.set_write_error_handler(lambda _: asyncio.create_task(self._on_disconnected()))
|
||||
self.communicator.add_listener(
|
||||
type_name_of(HeartBeat),
|
||||
lambda _: asyncio.create_task(self._on_heartbeat()),
|
||||
)
|
||||
self._read_task = asyncio.create_task(self._read_loop())
|
||||
asyncio.create_task(self._send_heartbeat())
|
||||
|
||||
async def write_packet(self, base: PacketBase) -> None:
|
||||
await self._ws.send(base.SerializeToString())
|
||||
|
||||
async def wait_closed(self) -> None:
|
||||
await self.communicator.wait_closed()
|
||||
|
||||
async def _do_close_impl(self) -> None:
|
||||
with suppress(Exception):
|
||||
await self._ws.close()
|
||||
|
||||
async def _read_loop(self) -> None:
|
||||
try:
|
||||
async for message in self._ws:
|
||||
if not isinstance(message, bytes):
|
||||
continue
|
||||
base = PacketBase()
|
||||
base.ParseFromString(message)
|
||||
self.communicator.on_received_data(
|
||||
base.typeName,
|
||||
base.data,
|
||||
base.nonce,
|
||||
base.responseNonce,
|
||||
)
|
||||
asyncio.create_task(self._send_heartbeat())
|
||||
except (ConnectionClosed, ConnectionError, OSError, DecodeError):
|
||||
pass
|
||||
finally:
|
||||
await self._on_disconnected()
|
||||
|
||||
|
||||
async def connect_ws(
|
||||
host: str,
|
||||
port: int,
|
||||
path: str,
|
||||
interval_sec: int,
|
||||
wait_sec: int,
|
||||
parser_map: ParserMap,
|
||||
) -> WsClient:
|
||||
uri = f"ws://{host}:{port}{path}"
|
||||
ws = await ws_connect(uri)
|
||||
return WsClient(ws, interval_sec, wait_sec, parser_map)
|
||||
91
python/toki_socket/ws_server.py
Normal file
91
python/toki_socket/ws_server.py
Normal file
|
|
@ -0,0 +1,91 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
from google.protobuf.message import Message
|
||||
|
||||
try:
|
||||
from websockets.asyncio.server import serve as ws_serve
|
||||
|
||||
_NEW_WEBSOCKETS_API = True
|
||||
except ImportError: # websockets 12.x legacy API
|
||||
from websockets import serve as ws_serve
|
||||
|
||||
_NEW_WEBSOCKETS_API = False
|
||||
|
||||
from toki_socket.communicator import ParserMap
|
||||
from toki_socket.ws_client import WsClient
|
||||
|
||||
|
||||
class WsServer:
|
||||
def __init__(
|
||||
self,
|
||||
host: str,
|
||||
port: int,
|
||||
path: str,
|
||||
new_client: Callable[[Any], WsClient] | None = None,
|
||||
interval_sec: int = 0,
|
||||
wait_sec: int = 0,
|
||||
parser_map: ParserMap | None = None,
|
||||
) -> None:
|
||||
self._host = host
|
||||
self._port = port
|
||||
self._path = path
|
||||
self._new_client = new_client
|
||||
self._interval_sec = interval_sec
|
||||
self._wait_sec = wait_sec
|
||||
self._parser_map = parser_map
|
||||
self._clients: list[WsClient] = []
|
||||
self._server: Any | None = None
|
||||
self.on_client_connected: Callable[[WsClient], object] = lambda _: None
|
||||
|
||||
@property
|
||||
def port(self) -> int:
|
||||
if self._server is None or not self._server.sockets:
|
||||
return self._port
|
||||
return int(self._server.sockets[0].getsockname()[1])
|
||||
|
||||
def clients(self) -> list[WsClient]:
|
||||
return list(self._clients)
|
||||
|
||||
async def start(self) -> None:
|
||||
handler = self._handle_websocket if _NEW_WEBSOCKETS_API else self._handle_legacy_websocket
|
||||
self._server = await ws_serve(handler, self._host, self._port)
|
||||
|
||||
async def stop(self) -> None:
|
||||
if self._server is not None:
|
||||
self._server.close()
|
||||
await self._server.wait_closed()
|
||||
self._server = None
|
||||
clients = list(self._clients)
|
||||
self._clients.clear()
|
||||
for client in clients:
|
||||
await client.close()
|
||||
|
||||
async def broadcast(self, message: Message) -> None:
|
||||
for client in self.clients():
|
||||
await client.communicator.send(message)
|
||||
|
||||
async def _handle_websocket(self, ws: Any) -> None:
|
||||
if self._new_client is not None:
|
||||
client = self._new_client(ws)
|
||||
else:
|
||||
if self._parser_map is None:
|
||||
raise ValueError("parser_map is required when new_client is not provided")
|
||||
client = WsClient(ws, self._interval_sec, self._wait_sec, self._parser_map)
|
||||
|
||||
self._clients.append(client)
|
||||
client.add_disconnect_listener(lambda c: self._remove_client(c))
|
||||
result = self.on_client_connected(client)
|
||||
if inspect.isawaitable(result):
|
||||
await result
|
||||
await client.wait_closed()
|
||||
|
||||
def _remove_client(self, client: WsClient) -> None:
|
||||
if client in self._clients:
|
||||
self._clients.remove(client)
|
||||
|
||||
async def _handle_legacy_websocket(self, ws: Any, _path: str) -> None:
|
||||
await self._handle_websocket(ws)
|
||||
Loading…
Reference in a new issue