ws-scale 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- ws_scale-0.1.0/LICENSE +19 -0
- ws_scale-0.1.0/PKG-INFO +149 -0
- ws_scale-0.1.0/README.md +135 -0
- ws_scale-0.1.0/pyproject.toml +33 -0
- ws_scale-0.1.0/setup.cfg +4 -0
- ws_scale-0.1.0/src/ws_scale.egg-info/PKG-INFO +149 -0
- ws_scale-0.1.0/src/ws_scale.egg-info/SOURCES.txt +16 -0
- ws_scale-0.1.0/src/ws_scale.egg-info/dependency_links.txt +1 -0
- ws_scale-0.1.0/src/ws_scale.egg-info/requires.txt +3 -0
- ws_scale-0.1.0/src/ws_scale.egg-info/top_level.txt +1 -0
- ws_scale-0.1.0/src/wsscale/__init__.py +5 -0
- ws_scale-0.1.0/src/wsscale/cache/__init__.py +0 -0
- ws_scale-0.1.0/src/wsscale/cache/redis.py +21 -0
- ws_scale-0.1.0/src/wsscale/ids/__init__.py +0 -0
- ws_scale-0.1.0/src/wsscale/ids/master.py +201 -0
- ws_scale-0.1.0/src/wsscale/ids/node.py +122 -0
- ws_scale-0.1.0/src/wsscale/ws/__init__.py +0 -0
- ws_scale-0.1.0/src/wsscale/ws/server.py +80 -0
ws_scale-0.1.0/LICENSE
ADDED
|
@@ -0,0 +1,19 @@
|
|
|
1
|
+
Copyright (c) 2018 The Python Packaging Authority
|
|
2
|
+
|
|
3
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
4
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
5
|
+
in the Software without restriction, including without limitation the rights
|
|
6
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
7
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
8
|
+
furnished to do so, subject to the following conditions:
|
|
9
|
+
|
|
10
|
+
The above copyright notice and this permission notice shall be included in all
|
|
11
|
+
copies or substantial portions of the Software.
|
|
12
|
+
|
|
13
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
14
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
15
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
16
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
17
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
18
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
19
|
+
SOFTWARE.
|
ws_scale-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,149 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: ws-scale
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: A websocket server application
|
|
5
|
+
Project-URL: Homepage, https://github.com/stevinho29/ws-scale
|
|
6
|
+
Project-URL: Issues, https://github.com/stevinho29/ws-scale/issues
|
|
7
|
+
Requires-Python: >=3.13
|
|
8
|
+
Description-Content-Type: text/markdown
|
|
9
|
+
License-File: LICENSE
|
|
10
|
+
Requires-Dist: aiohttp>=3.14.1
|
|
11
|
+
Requires-Dist: redis>=8.0.1
|
|
12
|
+
Requires-Dist: websockets>=16.0
|
|
13
|
+
Dynamic: license-file
|
|
14
|
+
|
|
15
|
+
# ws-scale
|
|
16
|
+
|
|
17
|
+
A Python library for building scalable WebSocket servers with distributed Snowflake-style ID generation. Each WebSocket node registers with a central master, receives a unique node ID, and generates collision-free 64-bit IDs across the cluster using Redis for coordination.
|
|
18
|
+
|
|
19
|
+
## Architecture
|
|
20
|
+
|
|
21
|
+
```
|
|
22
|
+
┌─────────────────────────┐
|
|
23
|
+
│ IDGeneratorMaster │ HTTP :9000
|
|
24
|
+
│ (aiohttp, one per DC) │◄──── register / heartbeat
|
|
25
|
+
└────────────┬────────────┘
|
|
26
|
+
│ Redis :6379
|
|
27
|
+
┌────────────▼────────────┐ ┌──────────────────────────┐
|
|
28
|
+
│ WebsocketServer #1 │ │ WebsocketServer #2 │
|
|
29
|
+
│ + IDGeneratorNode │ │ + IDGeneratorNode │
|
|
30
|
+
│ (websockets, :8080) │ │ (websockets, :8081) │
|
|
31
|
+
└─────────────────────────┘ └──────────────────────────┘
|
|
32
|
+
```
|
|
33
|
+
|
|
34
|
+
- **IDGeneratorMaster** — HTTP server that allocates node IDs and evicts stale nodes (no heartbeat in 2 min). Supports up to 32 nodes per datacenter.
|
|
35
|
+
- **IDGeneratorNode** — Registers with the master on startup, sends heartbeats every 60 s, and generates Snowflake IDs (41-bit timestamp | 5-bit datacenter | 5-bit node | 12-bit sequence).
|
|
36
|
+
- **WebsocketServer** — Wraps `websockets.serve`, bootstraps an `IDGeneratorNode`, and routes connections to a user-supplied async handler.
|
|
37
|
+
|
|
38
|
+
## Requirements
|
|
39
|
+
|
|
40
|
+
- Python 3.13+
|
|
41
|
+
- Redis
|
|
42
|
+
|
|
43
|
+
## Installation
|
|
44
|
+
|
|
45
|
+
```bash
|
|
46
|
+
pip install ws-scale
|
|
47
|
+
```
|
|
48
|
+
|
|
49
|
+
Or with [uv](https://github.com/astral-sh/uv):
|
|
50
|
+
|
|
51
|
+
```bash
|
|
52
|
+
uv add ws-scale
|
|
53
|
+
```
|
|
54
|
+
|
|
55
|
+
## Configuration
|
|
56
|
+
|
|
57
|
+
Settings can be provided via a Python module or environment variables.
|
|
58
|
+
|
|
59
|
+
| Setting | Default | Description |
|
|
60
|
+
|---|---|---|
|
|
61
|
+
| `WS_REDIS_HOST` | `localhost` | Redis hostname |
|
|
62
|
+
| `WS_REDIS_PORT` | `6379` | Redis port |
|
|
63
|
+
| `WS_MASTER_HOST` | `http://0.0.0.0` | ID master hostname |
|
|
64
|
+
| `WS_MASTER_PORT` | `9000` | ID master port |
|
|
65
|
+
|
|
66
|
+
**Via settings module** (`settings.py`):
|
|
67
|
+
|
|
68
|
+
```python
|
|
69
|
+
WS_REDIS_HOST = "localhost"
|
|
70
|
+
WS_REDIS_PORT = 6379
|
|
71
|
+
|
|
72
|
+
WS_MASTER_HOST = "http://0.0.0.0"
|
|
73
|
+
WS_MASTER_PORT = 9000
|
|
74
|
+
```
|
|
75
|
+
|
|
76
|
+
**Via environment variables** (pass `settings_path=None`):
|
|
77
|
+
|
|
78
|
+
```bash
|
|
79
|
+
export WS_REDIS_HOST=localhost
|
|
80
|
+
export WS_REDIS_PORT=6379
|
|
81
|
+
export WS_MASTER_HOST=http://0.0.0.0
|
|
82
|
+
export WS_MASTER_PORT=9000
|
|
83
|
+
```
|
|
84
|
+
|
|
85
|
+
## Quick Start
|
|
86
|
+
|
|
87
|
+
```python
|
|
88
|
+
import asyncio
|
|
89
|
+
from wsscale import IDGeneratorMaster, WebsocketServer
|
|
90
|
+
|
|
91
|
+
async def handle(conn):
|
|
92
|
+
async for message in conn:
|
|
93
|
+
await conn.send(f"echo: {message}")
|
|
94
|
+
|
|
95
|
+
async def main():
|
|
96
|
+
master = IDGeneratorMaster(datacenter_id=0, settings_path="settings")
|
|
97
|
+
server = WebsocketServer(datacenter_id=0, port=8080, settings_path="settings")
|
|
98
|
+
|
|
99
|
+
await master.cleanup() # clear stale Redis state on fresh start
|
|
100
|
+
await master.serve() # start HTTP master on WS_MASTER_PORT
|
|
101
|
+
await server.bootstrap(handler=handle) # register node, start heartbeat, serve WS
|
|
102
|
+
|
|
103
|
+
asyncio.run(main())
|
|
104
|
+
```
|
|
105
|
+
|
|
106
|
+
Run the bundled example from the command line:
|
|
107
|
+
|
|
108
|
+
```bash
|
|
109
|
+
python src/example.py --port 8080 --datacenter 0 --generator True
|
|
110
|
+
```
|
|
111
|
+
|
|
112
|
+
## API Reference
|
|
113
|
+
|
|
114
|
+
### `WebsocketServer(datacenter_id, port, settings_path=None)`
|
|
115
|
+
|
|
116
|
+
| Parameter | Type | Description |
|
|
117
|
+
|---|---|---|
|
|
118
|
+
| `datacenter_id` | `int` | 5-bit datacenter identifier (0–31) |
|
|
119
|
+
| `port` | `int` | TCP port for the WebSocket server |
|
|
120
|
+
| `settings_path` | `str \| None` | Dotted module path to a settings file, or `None` to use env vars |
|
|
121
|
+
|
|
122
|
+
**`await bootstrap(handler)`** — Registers the ID node, starts the heartbeat loop, and runs the WebSocket server.
|
|
123
|
+
|
|
124
|
+
---
|
|
125
|
+
|
|
126
|
+
### `IDGeneratorMaster(datacenter_id, settings_path=None)`
|
|
127
|
+
|
|
128
|
+
| Parameter | Type | Description |
|
|
129
|
+
|---|---|---|
|
|
130
|
+
| `datacenter_id` | `str` | Logical datacenter identifier |
|
|
131
|
+
| `settings_path` | `str \| None` | Dotted module path to a settings file, or `None` to use env vars |
|
|
132
|
+
|
|
133
|
+
**`await serve()`** — Starts the aiohttp HTTP server exposing `POST /register` and `POST /heartbeat`.
|
|
134
|
+
**`await cleanup()`** — Removes all node entries from Redis (call before `serve()` on a fresh start).
|
|
135
|
+
**`await monitor_nodes()`** — Evicts nodes that have missed heartbeats for more than 2 minutes.
|
|
136
|
+
|
|
137
|
+
---
|
|
138
|
+
|
|
139
|
+
### `IDGeneratorNode(server_host, server_port, data_center_id, redis_host, redis_port)`
|
|
140
|
+
|
|
141
|
+
Typically instantiated automatically by `WebsocketServer`, but can be used standalone.
|
|
142
|
+
|
|
143
|
+
**`await generate_id(datacenter_id, node_id) → int`** — Returns a 64-bit Snowflake-style unique ID.
|
|
144
|
+
**`await register()`** — Contacts the master and obtains a `node_id`.
|
|
145
|
+
**`await heart_beat_loop(interval_seconds)`** — Sends periodic heartbeats to the master forever.
|
|
146
|
+
|
|
147
|
+
## License
|
|
148
|
+
|
|
149
|
+
MIT — see [LICENSE](LICENSE).
|
ws_scale-0.1.0/README.md
ADDED
|
@@ -0,0 +1,135 @@
|
|
|
1
|
+
# ws-scale
|
|
2
|
+
|
|
3
|
+
A Python library for building scalable WebSocket servers with distributed Snowflake-style ID generation. Each WebSocket node registers with a central master, receives a unique node ID, and generates collision-free 64-bit IDs across the cluster using Redis for coordination.
|
|
4
|
+
|
|
5
|
+
## Architecture
|
|
6
|
+
|
|
7
|
+
```
|
|
8
|
+
┌─────────────────────────┐
|
|
9
|
+
│ IDGeneratorMaster │ HTTP :9000
|
|
10
|
+
│ (aiohttp, one per DC) │◄──── register / heartbeat
|
|
11
|
+
└────────────┬────────────┘
|
|
12
|
+
│ Redis :6379
|
|
13
|
+
┌────────────▼────────────┐ ┌──────────────────────────┐
|
|
14
|
+
│ WebsocketServer #1 │ │ WebsocketServer #2 │
|
|
15
|
+
│ + IDGeneratorNode │ │ + IDGeneratorNode │
|
|
16
|
+
│ (websockets, :8080) │ │ (websockets, :8081) │
|
|
17
|
+
└─────────────────────────┘ └──────────────────────────┘
|
|
18
|
+
```
|
|
19
|
+
|
|
20
|
+
- **IDGeneratorMaster** — HTTP server that allocates node IDs and evicts stale nodes (no heartbeat in 2 min). Supports up to 32 nodes per datacenter.
|
|
21
|
+
- **IDGeneratorNode** — Registers with the master on startup, sends heartbeats every 60 s, and generates Snowflake IDs (41-bit timestamp | 5-bit datacenter | 5-bit node | 12-bit sequence).
|
|
22
|
+
- **WebsocketServer** — Wraps `websockets.serve`, bootstraps an `IDGeneratorNode`, and routes connections to a user-supplied async handler.
|
|
23
|
+
|
|
24
|
+
## Requirements
|
|
25
|
+
|
|
26
|
+
- Python 3.13+
|
|
27
|
+
- Redis
|
|
28
|
+
|
|
29
|
+
## Installation
|
|
30
|
+
|
|
31
|
+
```bash
|
|
32
|
+
pip install ws-scale
|
|
33
|
+
```
|
|
34
|
+
|
|
35
|
+
Or with [uv](https://github.com/astral-sh/uv):
|
|
36
|
+
|
|
37
|
+
```bash
|
|
38
|
+
uv add ws-scale
|
|
39
|
+
```
|
|
40
|
+
|
|
41
|
+
## Configuration
|
|
42
|
+
|
|
43
|
+
Settings can be provided via a Python module or environment variables.
|
|
44
|
+
|
|
45
|
+
| Setting | Default | Description |
|
|
46
|
+
|---|---|---|
|
|
47
|
+
| `WS_REDIS_HOST` | `localhost` | Redis hostname |
|
|
48
|
+
| `WS_REDIS_PORT` | `6379` | Redis port |
|
|
49
|
+
| `WS_MASTER_HOST` | `http://0.0.0.0` | ID master hostname |
|
|
50
|
+
| `WS_MASTER_PORT` | `9000` | ID master port |
|
|
51
|
+
|
|
52
|
+
**Via settings module** (`settings.py`):
|
|
53
|
+
|
|
54
|
+
```python
|
|
55
|
+
WS_REDIS_HOST = "localhost"
|
|
56
|
+
WS_REDIS_PORT = 6379
|
|
57
|
+
|
|
58
|
+
WS_MASTER_HOST = "http://0.0.0.0"
|
|
59
|
+
WS_MASTER_PORT = 9000
|
|
60
|
+
```
|
|
61
|
+
|
|
62
|
+
**Via environment variables** (pass `settings_path=None`):
|
|
63
|
+
|
|
64
|
+
```bash
|
|
65
|
+
export WS_REDIS_HOST=localhost
|
|
66
|
+
export WS_REDIS_PORT=6379
|
|
67
|
+
export WS_MASTER_HOST=http://0.0.0.0
|
|
68
|
+
export WS_MASTER_PORT=9000
|
|
69
|
+
```
|
|
70
|
+
|
|
71
|
+
## Quick Start
|
|
72
|
+
|
|
73
|
+
```python
|
|
74
|
+
import asyncio
|
|
75
|
+
from wsscale import IDGeneratorMaster, WebsocketServer
|
|
76
|
+
|
|
77
|
+
async def handle(conn):
|
|
78
|
+
async for message in conn:
|
|
79
|
+
await conn.send(f"echo: {message}")
|
|
80
|
+
|
|
81
|
+
async def main():
|
|
82
|
+
master = IDGeneratorMaster(datacenter_id=0, settings_path="settings")
|
|
83
|
+
server = WebsocketServer(datacenter_id=0, port=8080, settings_path="settings")
|
|
84
|
+
|
|
85
|
+
await master.cleanup() # clear stale Redis state on fresh start
|
|
86
|
+
await master.serve() # start HTTP master on WS_MASTER_PORT
|
|
87
|
+
await server.bootstrap(handler=handle) # register node, start heartbeat, serve WS
|
|
88
|
+
|
|
89
|
+
asyncio.run(main())
|
|
90
|
+
```
|
|
91
|
+
|
|
92
|
+
Run the bundled example from the command line:
|
|
93
|
+
|
|
94
|
+
```bash
|
|
95
|
+
python src/example.py --port 8080 --datacenter 0 --generator True
|
|
96
|
+
```
|
|
97
|
+
|
|
98
|
+
## API Reference
|
|
99
|
+
|
|
100
|
+
### `WebsocketServer(datacenter_id, port, settings_path=None)`
|
|
101
|
+
|
|
102
|
+
| Parameter | Type | Description |
|
|
103
|
+
|---|---|---|
|
|
104
|
+
| `datacenter_id` | `int` | 5-bit datacenter identifier (0–31) |
|
|
105
|
+
| `port` | `int` | TCP port for the WebSocket server |
|
|
106
|
+
| `settings_path` | `str \| None` | Dotted module path to a settings file, or `None` to use env vars |
|
|
107
|
+
|
|
108
|
+
**`await bootstrap(handler)`** — Registers the ID node, starts the heartbeat loop, and runs the WebSocket server.
|
|
109
|
+
|
|
110
|
+
---
|
|
111
|
+
|
|
112
|
+
### `IDGeneratorMaster(datacenter_id, settings_path=None)`
|
|
113
|
+
|
|
114
|
+
| Parameter | Type | Description |
|
|
115
|
+
|---|---|---|
|
|
116
|
+
| `datacenter_id` | `str` | Logical datacenter identifier |
|
|
117
|
+
| `settings_path` | `str \| None` | Dotted module path to a settings file, or `None` to use env vars |
|
|
118
|
+
|
|
119
|
+
**`await serve()`** — Starts the aiohttp HTTP server exposing `POST /register` and `POST /heartbeat`.
|
|
120
|
+
**`await cleanup()`** — Removes all node entries from Redis (call before `serve()` on a fresh start).
|
|
121
|
+
**`await monitor_nodes()`** — Evicts nodes that have missed heartbeats for more than 2 minutes.
|
|
122
|
+
|
|
123
|
+
---
|
|
124
|
+
|
|
125
|
+
### `IDGeneratorNode(server_host, server_port, data_center_id, redis_host, redis_port)`
|
|
126
|
+
|
|
127
|
+
Typically instantiated automatically by `WebsocketServer`, but can be used standalone.
|
|
128
|
+
|
|
129
|
+
**`await generate_id(datacenter_id, node_id) → int`** — Returns a 64-bit Snowflake-style unique ID.
|
|
130
|
+
**`await register()`** — Contacts the master and obtains a `node_id`.
|
|
131
|
+
**`await heart_beat_loop(interval_seconds)`** — Sends periodic heartbeats to the master forever.
|
|
132
|
+
|
|
133
|
+
## License
|
|
134
|
+
|
|
135
|
+
MIT — see [LICENSE](LICENSE).
|
|
@@ -0,0 +1,33 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "ws-scale"
|
|
3
|
+
version = "0.1.0"
|
|
4
|
+
description = "A websocket server application"
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
requires-python = ">=3.13"
|
|
7
|
+
dependencies = [
|
|
8
|
+
"aiohttp>=3.14.1",
|
|
9
|
+
"redis>=8.0.1",
|
|
10
|
+
"websockets>=16.0",
|
|
11
|
+
]
|
|
12
|
+
|
|
13
|
+
[build-system]
|
|
14
|
+
requires = ["setuptools>=61.0"]
|
|
15
|
+
build-backend = "setuptools.build_meta"
|
|
16
|
+
|
|
17
|
+
authors = [
|
|
18
|
+
{ name="Ndemanou steve", email="ndemanousteve@gmail.com" },
|
|
19
|
+
]
|
|
20
|
+
|
|
21
|
+
classifiers = [
|
|
22
|
+
"Programming Language :: Python :: 3",
|
|
23
|
+
"Operating System :: OS Independent",
|
|
24
|
+
]
|
|
25
|
+
license = "MIT"
|
|
26
|
+
license-files = ["LICEN[CS]E*"]
|
|
27
|
+
|
|
28
|
+
[project.urls]
|
|
29
|
+
Homepage = "https://github.com/stevinho29/ws-scale"
|
|
30
|
+
Issues = "https://github.com/stevinho29/ws-scale/issues"
|
|
31
|
+
|
|
32
|
+
[tool.setuptools.packages.find]
|
|
33
|
+
where = ["src"]
|
ws_scale-0.1.0/setup.cfg
ADDED
|
@@ -0,0 +1,149 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: ws-scale
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: A websocket server application
|
|
5
|
+
Project-URL: Homepage, https://github.com/stevinho29/ws-scale
|
|
6
|
+
Project-URL: Issues, https://github.com/stevinho29/ws-scale/issues
|
|
7
|
+
Requires-Python: >=3.13
|
|
8
|
+
Description-Content-Type: text/markdown
|
|
9
|
+
License-File: LICENSE
|
|
10
|
+
Requires-Dist: aiohttp>=3.14.1
|
|
11
|
+
Requires-Dist: redis>=8.0.1
|
|
12
|
+
Requires-Dist: websockets>=16.0
|
|
13
|
+
Dynamic: license-file
|
|
14
|
+
|
|
15
|
+
# ws-scale
|
|
16
|
+
|
|
17
|
+
A Python library for building scalable WebSocket servers with distributed Snowflake-style ID generation. Each WebSocket node registers with a central master, receives a unique node ID, and generates collision-free 64-bit IDs across the cluster using Redis for coordination.
|
|
18
|
+
|
|
19
|
+
## Architecture
|
|
20
|
+
|
|
21
|
+
```
|
|
22
|
+
┌─────────────────────────┐
|
|
23
|
+
│ IDGeneratorMaster │ HTTP :9000
|
|
24
|
+
│ (aiohttp, one per DC) │◄──── register / heartbeat
|
|
25
|
+
└────────────┬────────────┘
|
|
26
|
+
│ Redis :6379
|
|
27
|
+
┌────────────▼────────────┐ ┌──────────────────────────┐
|
|
28
|
+
│ WebsocketServer #1 │ │ WebsocketServer #2 │
|
|
29
|
+
│ + IDGeneratorNode │ │ + IDGeneratorNode │
|
|
30
|
+
│ (websockets, :8080) │ │ (websockets, :8081) │
|
|
31
|
+
└─────────────────────────┘ └──────────────────────────┘
|
|
32
|
+
```
|
|
33
|
+
|
|
34
|
+
- **IDGeneratorMaster** — HTTP server that allocates node IDs and evicts stale nodes (no heartbeat in 2 min). Supports up to 32 nodes per datacenter.
|
|
35
|
+
- **IDGeneratorNode** — Registers with the master on startup, sends heartbeats every 60 s, and generates Snowflake IDs (41-bit timestamp | 5-bit datacenter | 5-bit node | 12-bit sequence).
|
|
36
|
+
- **WebsocketServer** — Wraps `websockets.serve`, bootstraps an `IDGeneratorNode`, and routes connections to a user-supplied async handler.
|
|
37
|
+
|
|
38
|
+
## Requirements
|
|
39
|
+
|
|
40
|
+
- Python 3.13+
|
|
41
|
+
- Redis
|
|
42
|
+
|
|
43
|
+
## Installation
|
|
44
|
+
|
|
45
|
+
```bash
|
|
46
|
+
pip install ws-scale
|
|
47
|
+
```
|
|
48
|
+
|
|
49
|
+
Or with [uv](https://github.com/astral-sh/uv):
|
|
50
|
+
|
|
51
|
+
```bash
|
|
52
|
+
uv add ws-scale
|
|
53
|
+
```
|
|
54
|
+
|
|
55
|
+
## Configuration
|
|
56
|
+
|
|
57
|
+
Settings can be provided via a Python module or environment variables.
|
|
58
|
+
|
|
59
|
+
| Setting | Default | Description |
|
|
60
|
+
|---|---|---|
|
|
61
|
+
| `WS_REDIS_HOST` | `localhost` | Redis hostname |
|
|
62
|
+
| `WS_REDIS_PORT` | `6379` | Redis port |
|
|
63
|
+
| `WS_MASTER_HOST` | `http://0.0.0.0` | ID master hostname |
|
|
64
|
+
| `WS_MASTER_PORT` | `9000` | ID master port |
|
|
65
|
+
|
|
66
|
+
**Via settings module** (`settings.py`):
|
|
67
|
+
|
|
68
|
+
```python
|
|
69
|
+
WS_REDIS_HOST = "localhost"
|
|
70
|
+
WS_REDIS_PORT = 6379
|
|
71
|
+
|
|
72
|
+
WS_MASTER_HOST = "http://0.0.0.0"
|
|
73
|
+
WS_MASTER_PORT = 9000
|
|
74
|
+
```
|
|
75
|
+
|
|
76
|
+
**Via environment variables** (pass `settings_path=None`):
|
|
77
|
+
|
|
78
|
+
```bash
|
|
79
|
+
export WS_REDIS_HOST=localhost
|
|
80
|
+
export WS_REDIS_PORT=6379
|
|
81
|
+
export WS_MASTER_HOST=http://0.0.0.0
|
|
82
|
+
export WS_MASTER_PORT=9000
|
|
83
|
+
```
|
|
84
|
+
|
|
85
|
+
## Quick Start
|
|
86
|
+
|
|
87
|
+
```python
|
|
88
|
+
import asyncio
|
|
89
|
+
from wsscale import IDGeneratorMaster, WebsocketServer
|
|
90
|
+
|
|
91
|
+
async def handle(conn):
|
|
92
|
+
async for message in conn:
|
|
93
|
+
await conn.send(f"echo: {message}")
|
|
94
|
+
|
|
95
|
+
async def main():
|
|
96
|
+
master = IDGeneratorMaster(datacenter_id=0, settings_path="settings")
|
|
97
|
+
server = WebsocketServer(datacenter_id=0, port=8080, settings_path="settings")
|
|
98
|
+
|
|
99
|
+
await master.cleanup() # clear stale Redis state on fresh start
|
|
100
|
+
await master.serve() # start HTTP master on WS_MASTER_PORT
|
|
101
|
+
await server.bootstrap(handler=handle) # register node, start heartbeat, serve WS
|
|
102
|
+
|
|
103
|
+
asyncio.run(main())
|
|
104
|
+
```
|
|
105
|
+
|
|
106
|
+
Run the bundled example from the command line:
|
|
107
|
+
|
|
108
|
+
```bash
|
|
109
|
+
python src/example.py --port 8080 --datacenter 0 --generator True
|
|
110
|
+
```
|
|
111
|
+
|
|
112
|
+
## API Reference
|
|
113
|
+
|
|
114
|
+
### `WebsocketServer(datacenter_id, port, settings_path=None)`
|
|
115
|
+
|
|
116
|
+
| Parameter | Type | Description |
|
|
117
|
+
|---|---|---|
|
|
118
|
+
| `datacenter_id` | `int` | 5-bit datacenter identifier (0–31) |
|
|
119
|
+
| `port` | `int` | TCP port for the WebSocket server |
|
|
120
|
+
| `settings_path` | `str \| None` | Dotted module path to a settings file, or `None` to use env vars |
|
|
121
|
+
|
|
122
|
+
**`await bootstrap(handler)`** — Registers the ID node, starts the heartbeat loop, and runs the WebSocket server.
|
|
123
|
+
|
|
124
|
+
---
|
|
125
|
+
|
|
126
|
+
### `IDGeneratorMaster(datacenter_id, settings_path=None)`
|
|
127
|
+
|
|
128
|
+
| Parameter | Type | Description |
|
|
129
|
+
|---|---|---|
|
|
130
|
+
| `datacenter_id` | `str` | Logical datacenter identifier |
|
|
131
|
+
| `settings_path` | `str \| None` | Dotted module path to a settings file, or `None` to use env vars |
|
|
132
|
+
|
|
133
|
+
**`await serve()`** — Starts the aiohttp HTTP server exposing `POST /register` and `POST /heartbeat`.
|
|
134
|
+
**`await cleanup()`** — Removes all node entries from Redis (call before `serve()` on a fresh start).
|
|
135
|
+
**`await monitor_nodes()`** — Evicts nodes that have missed heartbeats for more than 2 minutes.
|
|
136
|
+
|
|
137
|
+
---
|
|
138
|
+
|
|
139
|
+
### `IDGeneratorNode(server_host, server_port, data_center_id, redis_host, redis_port)`
|
|
140
|
+
|
|
141
|
+
Typically instantiated automatically by `WebsocketServer`, but can be used standalone.
|
|
142
|
+
|
|
143
|
+
**`await generate_id(datacenter_id, node_id) → int`** — Returns a 64-bit Snowflake-style unique ID.
|
|
144
|
+
**`await register()`** — Contacts the master and obtains a `node_id`.
|
|
145
|
+
**`await heart_beat_loop(interval_seconds)`** — Sends periodic heartbeats to the master forever.
|
|
146
|
+
|
|
147
|
+
## License
|
|
148
|
+
|
|
149
|
+
MIT — see [LICENSE](LICENSE).
|
|
@@ -0,0 +1,16 @@
|
|
|
1
|
+
LICENSE
|
|
2
|
+
README.md
|
|
3
|
+
pyproject.toml
|
|
4
|
+
src/ws_scale.egg-info/PKG-INFO
|
|
5
|
+
src/ws_scale.egg-info/SOURCES.txt
|
|
6
|
+
src/ws_scale.egg-info/dependency_links.txt
|
|
7
|
+
src/ws_scale.egg-info/requires.txt
|
|
8
|
+
src/ws_scale.egg-info/top_level.txt
|
|
9
|
+
src/wsscale/__init__.py
|
|
10
|
+
src/wsscale/cache/__init__.py
|
|
11
|
+
src/wsscale/cache/redis.py
|
|
12
|
+
src/wsscale/ids/__init__.py
|
|
13
|
+
src/wsscale/ids/master.py
|
|
14
|
+
src/wsscale/ids/node.py
|
|
15
|
+
src/wsscale/ws/__init__.py
|
|
16
|
+
src/wsscale/ws/server.py
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
wsscale
|
|
File without changes
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
import redis.asyncio as redis
|
|
2
|
+
from redis import Redis
|
|
3
|
+
|
|
4
|
+
class RedisInstance:
|
|
5
|
+
"""Manages a single async Redis client connection."""
|
|
6
|
+
|
|
7
|
+
def __init__(self, redis_host:str, redis_port:int):
|
|
8
|
+
"""
|
|
9
|
+
Args:
|
|
10
|
+
redis_host: Hostname or IP address of the Redis server.
|
|
11
|
+
redis_port: Port number the Redis server is listening on.
|
|
12
|
+
"""
|
|
13
|
+
self.client:Redis = redis.Redis(host=redis_host, port=redis_port)
|
|
14
|
+
|
|
15
|
+
def get_instance(self):
|
|
16
|
+
"""Return the underlying async Redis client."""
|
|
17
|
+
return self.client
|
|
18
|
+
|
|
19
|
+
async def __del__(self):
|
|
20
|
+
"""Close the Redis connection when the instance is garbage-collected."""
|
|
21
|
+
await self.client.aclose()
|
|
File without changes
|
|
@@ -0,0 +1,201 @@
|
|
|
1
|
+
from datetime import datetime, timezone
|
|
2
|
+
from typing import Optional
|
|
3
|
+
from aiohttp import web
|
|
4
|
+
from importlib import import_module
|
|
5
|
+
import os
|
|
6
|
+
import json
|
|
7
|
+
import logging
|
|
8
|
+
|
|
9
|
+
from wsscale.cache.redis import RedisInstance
|
|
10
|
+
logger = logging.getLogger("ws-scale.ids.master")
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
class IDGeneratorMaster:
|
|
14
|
+
"""
|
|
15
|
+
HTTP server that manages node registration and liveness for distributed Snowflake ID generation.
|
|
16
|
+
|
|
17
|
+
Nodes register themselves to receive a unique node ID, then send periodic heartbeats.
|
|
18
|
+
Nodes that miss heartbeats for more than two minutes are evicted via ``monitor_nodes``.
|
|
19
|
+
"""
|
|
20
|
+
|
|
21
|
+
MAX_NODES = 32
|
|
22
|
+
|
|
23
|
+
def __init__(self, datacenter_id: str, settings_path:str=None):
|
|
24
|
+
"""
|
|
25
|
+
Args:
|
|
26
|
+
datacenter_id: Logical datacenter identifier for this master instance.
|
|
27
|
+
settings_path: Dotted module path to a settings file (e.g. ``"app.settings"``).
|
|
28
|
+
Falls back to environment variables when omitted.
|
|
29
|
+
"""
|
|
30
|
+
self.datacenter_id = datacenter_id
|
|
31
|
+
if settings_path:
|
|
32
|
+
try:
|
|
33
|
+
settings = import_module(name=settings_path)
|
|
34
|
+
redis_host = getattr(settings, "WS_REDIS_HOST")
|
|
35
|
+
redis_port = getattr(settings, "WS_REDIS_PORT")
|
|
36
|
+
id_master_host = getattr(settings, "WS_MASTER_HOST")
|
|
37
|
+
id_master_port = getattr(settings, "WS_MASTER_PORT")
|
|
38
|
+
except ModuleNotFoundError:
|
|
39
|
+
logger.error(f"Settings module not found, invalid path {settings_path}")
|
|
40
|
+
raise
|
|
41
|
+
else:
|
|
42
|
+
redis_host = os.environ.get("WS_REDIS_HOST")
|
|
43
|
+
redis_port = os.environ.get("WS_REDIS_PORT")
|
|
44
|
+
id_master_host = os.environ.get("WS_MASTER_HOST")
|
|
45
|
+
id_master_port = os.environ.get("WS_MASTER_PORT")
|
|
46
|
+
|
|
47
|
+
self.port = id_master_port
|
|
48
|
+
self.host = id_master_host
|
|
49
|
+
self.redis = RedisInstance(redis_host=redis_host, redis_port=redis_port).get_instance()
|
|
50
|
+
self.log_header = "[IDGeneratorMaster]"
|
|
51
|
+
|
|
52
|
+
def _get_now(self) -> float:
|
|
53
|
+
"""Return the current UTC time as a millisecond-precision Unix timestamp."""
|
|
54
|
+
return datetime.now(tz=timezone.utc).timestamp() * 1000
|
|
55
|
+
|
|
56
|
+
async def _generate_node_id(self) -> Optional[int]:
|
|
57
|
+
"""Atomically allocate the next node ID from Redis.
|
|
58
|
+
|
|
59
|
+
Returns:
|
|
60
|
+
The new node ID (1-based), or ``None`` if ``MAX_NODES`` has been reached.
|
|
61
|
+
"""
|
|
62
|
+
node_key = "node_count"
|
|
63
|
+
count_bytes = await self.redis.get(node_key)
|
|
64
|
+
count = int(count_bytes.decode()) if count_bytes else 0
|
|
65
|
+
if count < self.MAX_NODES:
|
|
66
|
+
return await self.redis.incr(node_key, 1)
|
|
67
|
+
return None
|
|
68
|
+
|
|
69
|
+
async def _cache_node(self, node_id: int) -> None:
|
|
70
|
+
"""Store a node's initial last-seen timestamp in Redis.
|
|
71
|
+
|
|
72
|
+
Args:
|
|
73
|
+
node_id: The node to cache.
|
|
74
|
+
"""
|
|
75
|
+
await self.redis.set(name=f"node_{node_id}", value=self._get_now())
|
|
76
|
+
|
|
77
|
+
async def _update_node(self, node_id: int, heart_beat: float) -> None:
|
|
78
|
+
"""Overwrite a node's last-seen timestamp with a new heartbeat value.
|
|
79
|
+
|
|
80
|
+
Args:
|
|
81
|
+
node_id: The node whose timestamp should be updated.
|
|
82
|
+
heart_beat: New timestamp in milliseconds.
|
|
83
|
+
"""
|
|
84
|
+
await self.redis.set(name=f"node_{node_id}", value=heart_beat)
|
|
85
|
+
|
|
86
|
+
async def _delete_nodes(self, node_ids: list[int]) -> None:
|
|
87
|
+
"""Remove one or more nodes from Redis.
|
|
88
|
+
|
|
89
|
+
Args:
|
|
90
|
+
node_ids: List of node IDs to delete. No-op if the list is empty.
|
|
91
|
+
"""
|
|
92
|
+
if node_ids:
|
|
93
|
+
await self.redis.delete(*[f"node_{id}" for id in node_ids])
|
|
94
|
+
|
|
95
|
+
async def _get_node(self, node_id: int):
|
|
96
|
+
"""Fetch a single node's last-seen timestamp from Redis.
|
|
97
|
+
|
|
98
|
+
Args:
|
|
99
|
+
node_id: The node to look up.
|
|
100
|
+
|
|
101
|
+
Returns:
|
|
102
|
+
The raw bytes value stored in Redis, or ``None`` if the key does not exist.
|
|
103
|
+
"""
|
|
104
|
+
return await self.redis.get(name=f"node_{node_id}")
|
|
105
|
+
|
|
106
|
+
async def _get_nodes(self, node_ids: list[int]):
|
|
107
|
+
"""Fetch last-seen timestamps for multiple nodes in one round-trip.
|
|
108
|
+
|
|
109
|
+
Args:
|
|
110
|
+
node_ids: List of node IDs to look up.
|
|
111
|
+
|
|
112
|
+
Returns:
|
|
113
|
+
A list of raw bytes values (or ``None`` for missing keys), in the same order as ``node_ids``.
|
|
114
|
+
"""
|
|
115
|
+
return await self.redis.mget([f"node_{id}" for id in node_ids])
|
|
116
|
+
|
|
117
|
+
async def register(self, request: web.Request) -> web.Response:
|
|
118
|
+
"""HTTP POST /register — allocate a node ID and cache it.
|
|
119
|
+
|
|
120
|
+
Returns:
|
|
121
|
+
201 with ``{"node_id": <int>}`` on success, or 400 if the node limit is reached.
|
|
122
|
+
"""
|
|
123
|
+
node_id = await self._generate_node_id()
|
|
124
|
+
if not node_id:
|
|
125
|
+
return web.Response(
|
|
126
|
+
status=400,
|
|
127
|
+
content_type="application/json",
|
|
128
|
+
body=json.dumps({"message": "maximum number of nodes has been reached"}).encode(),
|
|
129
|
+
)
|
|
130
|
+
await self._cache_node(node_id=node_id)
|
|
131
|
+
return web.Response(
|
|
132
|
+
status=201,
|
|
133
|
+
content_type="application/json",
|
|
134
|
+
body=json.dumps({"message": f"node registered with success: {node_id}", "node_id": node_id}).encode(),
|
|
135
|
+
)
|
|
136
|
+
|
|
137
|
+
async def heartbeat(self, request: web.Request) -> web.Response:
|
|
138
|
+
"""HTTP POST /heartbeat — refresh a node's last-seen timestamp.
|
|
139
|
+
|
|
140
|
+
Expects a JSON body with ``{"node_id": <int>}``.
|
|
141
|
+
|
|
142
|
+
Returns:
|
|
143
|
+
200 on success, 400 if the body is invalid or the node is unknown.
|
|
144
|
+
"""
|
|
145
|
+
log_header = f"{self.log_header}[heartbeat]"
|
|
146
|
+
try:
|
|
147
|
+
body = await request.json()
|
|
148
|
+
except json.JSONDecodeError:
|
|
149
|
+
return web.Response(
|
|
150
|
+
status=400,
|
|
151
|
+
content_type="application/json",
|
|
152
|
+
body=json.dumps({"message": "invalid or empty body"}).encode()
|
|
153
|
+
)
|
|
154
|
+
node_id = body.get("node_id")
|
|
155
|
+
|
|
156
|
+
node = await self._get_node(node_id=node_id)
|
|
157
|
+
if node:
|
|
158
|
+
await self._update_node(node_id=node_id, heart_beat=self._get_now())
|
|
159
|
+
logger.info(f"{log_header} last seen date updated with success for node: {node_id}")
|
|
160
|
+
return web.Response(
|
|
161
|
+
status=200,
|
|
162
|
+
content_type="application/json",
|
|
163
|
+
body=json.dumps({"message": "last seen date updated with success"}).encode(),
|
|
164
|
+
)
|
|
165
|
+
logger.info(f"{log_header} unknown node with id {node_id}")
|
|
166
|
+
return web.Response(
|
|
167
|
+
status=400,
|
|
168
|
+
content_type="application/json",
|
|
169
|
+
body=json.dumps({"message": f"unknown node with id {node_id}"}).encode(),
|
|
170
|
+
)
|
|
171
|
+
|
|
172
|
+
async def monitor_nodes(self) -> None:
|
|
173
|
+
"""Evict nodes that have not sent a heartbeat in the last two minutes."""
|
|
174
|
+
log_header = f"{self.log_header}[monitor_nodes]"
|
|
175
|
+
two_minutes_ago = self._get_now() - 60 * 2 * 1000
|
|
176
|
+
node_heart_beats = await self._get_nodes(list(range(self.MAX_NODES)))
|
|
177
|
+
to_remove = []
|
|
178
|
+
for i, heart_beat in enumerate(node_heart_beats):
|
|
179
|
+
if heart_beat:
|
|
180
|
+
hb_value = float(heart_beat.decode() if isinstance(heart_beat, bytes) else heart_beat)
|
|
181
|
+
if hb_value < two_minutes_ago:
|
|
182
|
+
to_remove.append(i)
|
|
183
|
+
await self._delete_nodes(to_remove)
|
|
184
|
+
logger.info(f"{log_header} nodes {to_remove} removed")
|
|
185
|
+
|
|
186
|
+
async def serve(self) -> None:
|
|
187
|
+
"""Start the aiohttp HTTP server and block until it is stopped."""
|
|
188
|
+
app = web.Application()
|
|
189
|
+
app.router.add_post("/register", self.register)
|
|
190
|
+
app.router.add_post("/heartbeat", self.heartbeat)
|
|
191
|
+
runner = web.AppRunner(app)
|
|
192
|
+
await runner.setup()
|
|
193
|
+
site = web.TCPSite(runner, None, self.port)
|
|
194
|
+
await site.start()
|
|
195
|
+
logger.info(f"{self.log_header} id master server listening on localhost:{self.port}")
|
|
196
|
+
|
|
197
|
+
async def cleanup(self):
|
|
198
|
+
"""Delete all node entries from Redis (intended for graceful shutdown)."""
|
|
199
|
+
log_header = f"{self.log_header}[cleanup]"
|
|
200
|
+
await self._delete_nodes(list(range(self.MAX_NODES)))
|
|
201
|
+
logger.info(f"{log_header} redis database cleaned up with success")
|
|
@@ -0,0 +1,122 @@
|
|
|
1
|
+
import logging
|
|
2
|
+
import asyncio
|
|
3
|
+
from aiohttp import ClientSession
|
|
4
|
+
from datetime import datetime, timezone
|
|
5
|
+
|
|
6
|
+
from wsscale.cache.redis import RedisInstance
|
|
7
|
+
|
|
8
|
+
logger = logging.getLogger("ws-scale.ids.node")
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
class IDGeneratorNode:
|
|
12
|
+
"""
|
|
13
|
+
Snowflake-style distributed ID generator node.
|
|
14
|
+
|
|
15
|
+
Each node registers with the ``IDGeneratorMaster`` to obtain a unique ``node_id``,
|
|
16
|
+
then uses that ID together with a datacenter ID, a millisecond timestamp, and a
|
|
17
|
+
per-millisecond sequence counter to produce 64-bit sortable unique IDs.
|
|
18
|
+
"""
|
|
19
|
+
|
|
20
|
+
def __init__(self,
|
|
21
|
+
server_host:str,
|
|
22
|
+
server_port:int,
|
|
23
|
+
data_center_id:int,
|
|
24
|
+
redis_host:str,
|
|
25
|
+
redis_port:int
|
|
26
|
+
):
|
|
27
|
+
"""
|
|
28
|
+
Args:
|
|
29
|
+
server_host: Hostname or URL of the ``IDGeneratorMaster`` server.
|
|
30
|
+
server_port: Port the master server is listening on.
|
|
31
|
+
data_center_id: Logical datacenter identifier embedded in generated IDs.
|
|
32
|
+
redis_host: Hostname of the Redis server used for sequence counters.
|
|
33
|
+
redis_port: Port of the Redis server.
|
|
34
|
+
"""
|
|
35
|
+
self.epoch = datetime.now(tz=timezone.utc).replace(hour=0, minute=0, second=0, microsecond=0)
|
|
36
|
+
self.redis = RedisInstance(redis_host=redis_host, redis_port=redis_port).get_instance()
|
|
37
|
+
self.server_url = f"{server_host}:{server_port}"
|
|
38
|
+
self.data_center_id = data_center_id
|
|
39
|
+
self.node_id = None
|
|
40
|
+
self.registered = False
|
|
41
|
+
self.log_header = "[IDGeneratorNode]"
|
|
42
|
+
|
|
43
|
+
def get_elapsed_time_since_custom_epoch(self):
|
|
44
|
+
"""Return milliseconds elapsed since the custom epoch (midnight UTC of the start day)."""
|
|
45
|
+
now = datetime.now(tz=timezone.utc)
|
|
46
|
+
delta = (now - self.epoch).total_seconds()
|
|
47
|
+
return delta * 1000
|
|
48
|
+
|
|
49
|
+
async def get_sequence_number(self) -> int:
|
|
50
|
+
"""Atomically retrieve and increment the per-millisecond sequence counter from Redis.
|
|
51
|
+
|
|
52
|
+
The key expires after 1 ms so the counter resets each millisecond.
|
|
53
|
+
|
|
54
|
+
Returns:
|
|
55
|
+
Current sequence number for this millisecond.
|
|
56
|
+
"""
|
|
57
|
+
sequence_key = f'{self.node_id}_sequence_number'
|
|
58
|
+
if await self.redis.get(sequence_key):
|
|
59
|
+
sequence = await self.redis.incr(name=sequence_key, amount=1)
|
|
60
|
+
else:
|
|
61
|
+
sequence = await self.redis.set(name=sequence_key, value=0, px=1)
|
|
62
|
+
return sequence
|
|
63
|
+
|
|
64
|
+
|
|
65
|
+
async def generate_id(self, datacenter_id:int, node_id:int) -> int:
|
|
66
|
+
"""Produce a 64-bit Snowflake-style unique ID.
|
|
67
|
+
|
|
68
|
+
Bit layout: 41-bit timestamp | 5-bit datacenter ID | 5-bit node ID | 12-bit sequence.
|
|
69
|
+
|
|
70
|
+
Args:
|
|
71
|
+
datacenter_id: 5-bit datacenter identifier (0–31).
|
|
72
|
+
node_id: 5-bit node identifier (0–31).
|
|
73
|
+
|
|
74
|
+
Returns:
|
|
75
|
+
A globally unique 64-bit integer.
|
|
76
|
+
"""
|
|
77
|
+
timestamp = self.get_elapsed_time_since_custom_epoch()
|
|
78
|
+
sequence_number = await self.get_sequence_number()
|
|
79
|
+
_id = (int(timestamp) << 22) | (datacenter_id << 17) | (node_id << 12) | sequence_number
|
|
80
|
+
return int(_id)
|
|
81
|
+
|
|
82
|
+
async def register(self):
|
|
83
|
+
"""Register this node with the master server to obtain a unique ``node_id``."""
|
|
84
|
+
log_header = f"{self.log_header}[register]"
|
|
85
|
+
async with ClientSession() as client:
|
|
86
|
+
logger.info(f"{log_header} start registering id generator node to master")
|
|
87
|
+
async with client.post(url=self.server_url + '/register', timeout=30) as req:
|
|
88
|
+
if req.status == 201:
|
|
89
|
+
body = await req.json()
|
|
90
|
+
self.node_id = body.get("node_id")
|
|
91
|
+
self.registered = True
|
|
92
|
+
logger.info(f"{log_header} node registered with success, node_id: {self.node_id}")
|
|
93
|
+
else:
|
|
94
|
+
logger.error(f"{log_header} request to register ended with status code {req.status}")
|
|
95
|
+
logger.error(f"{log_header} {await req.text()}")
|
|
96
|
+
|
|
97
|
+
async def heart_beat_loop(self, interval_seconds):
|
|
98
|
+
"""Run ``heart_beat`` on a fixed interval forever.
|
|
99
|
+
|
|
100
|
+
Args:
|
|
101
|
+
interval_seconds: Seconds to wait between each heartbeat.
|
|
102
|
+
"""
|
|
103
|
+
log_header = f"{self.log_header}[heart_beat_loop]"
|
|
104
|
+
while True:
|
|
105
|
+
await asyncio.sleep(delay=interval_seconds)
|
|
106
|
+
logger.info(f"{log_header} heart beat called after {interval_seconds}")
|
|
107
|
+
await self.heart_beat()
|
|
108
|
+
|
|
109
|
+
async def heart_beat(self):
|
|
110
|
+
"""Send a single heartbeat POST to the master server to renew this node's liveness."""
|
|
111
|
+
log_header = f"{self.log_header}[heart_beat]"
|
|
112
|
+
logger.info(f"{log_header} about to send heart beat for node: {self.node_id}")
|
|
113
|
+
async with ClientSession() as client:
|
|
114
|
+
async with client.post(
|
|
115
|
+
url=self.server_url + "/heartbeat",
|
|
116
|
+
json={"node_id": self.node_id},
|
|
117
|
+
timeout=30) as req:
|
|
118
|
+
if req.status == 200:
|
|
119
|
+
logger.info(f"{log_header} heart beat sent with success")
|
|
120
|
+
else:
|
|
121
|
+
logger.error(f"{log_header} request to send heart beat ended with status code {req.status}")
|
|
122
|
+
logger.error(f"{log_header} {await req.text()}")
|
|
File without changes
|
|
@@ -0,0 +1,80 @@
|
|
|
1
|
+
|
|
2
|
+
import asyncio
|
|
3
|
+
import logging
|
|
4
|
+
import os
|
|
5
|
+
from importlib import import_module
|
|
6
|
+
from typing import Callable, Awaitable, Optional, Any
|
|
7
|
+
from websockets.asyncio.connection import Connection
|
|
8
|
+
from websockets.asyncio.server import serve, Server
|
|
9
|
+
from wsscale.ids.node import IDGeneratorNode
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
Handler = Awaitable[Callable[[Connection], Any]]
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
logger = logging.getLogger("ws-scale.server")
|
|
16
|
+
|
|
17
|
+
class WebsocketServer:
|
|
18
|
+
"""
|
|
19
|
+
Main entry point for the WebSocket server application.
|
|
20
|
+
|
|
21
|
+
Initializes an ``IDGeneratorNode``, registers it with the master on startup,
|
|
22
|
+
and drives a ``websockets`` server that routes incoming connections to the
|
|
23
|
+
caller-supplied handler.
|
|
24
|
+
"""
|
|
25
|
+
|
|
26
|
+
def __init__(self, datacenter_id:int, port:int, settings_path:Optional[str]=None):
|
|
27
|
+
"""
|
|
28
|
+
Args:
|
|
29
|
+
datacenter_id: 5-bit datacenter identifier embedded in generated IDs.
|
|
30
|
+
port: Local TCP port the WebSocket server will listen on.
|
|
31
|
+
settings_path: Dotted module path to a settings file (e.g. ``"app.settings"``).
|
|
32
|
+
Falls back to environment variables when omitted.
|
|
33
|
+
"""
|
|
34
|
+
self.port = port
|
|
35
|
+
self.server: Server = None
|
|
36
|
+
self.log_header = "[WebsocketServer]"
|
|
37
|
+
if settings_path:
|
|
38
|
+
try:
|
|
39
|
+
settings = import_module(name=settings_path)
|
|
40
|
+
redis_host = getattr(settings, "WS_REDIS_HOST")
|
|
41
|
+
redis_port = getattr(settings, "WS_REDIS_PORT")
|
|
42
|
+
id_master_host = getattr(settings, "WS_MASTER_HOST")
|
|
43
|
+
id_master_port = getattr(settings, "WS_MASTER_PORT")
|
|
44
|
+
except ModuleNotFoundError:
|
|
45
|
+
logger.error(f"Settings module not found, invalid path {settings_path}")
|
|
46
|
+
raise
|
|
47
|
+
|
|
48
|
+
else:
|
|
49
|
+
redis_host = os.environ.get("WS_REDIS_HOST")
|
|
50
|
+
redis_port = os.environ.get("WS_REDIS_PORT")
|
|
51
|
+
id_master_host = os.environ.get("WS_MASTER_HOST")
|
|
52
|
+
id_master_port = os.environ.get("WS_MASTER_PORT")
|
|
53
|
+
|
|
54
|
+
try:
|
|
55
|
+
self.node = IDGeneratorNode(
|
|
56
|
+
server_host=id_master_host,
|
|
57
|
+
server_port=id_master_port,
|
|
58
|
+
data_center_id=datacenter_id,
|
|
59
|
+
redis_host=redis_host,
|
|
60
|
+
redis_port=redis_port
|
|
61
|
+
)
|
|
62
|
+
except Exception as e:
|
|
63
|
+
logger.warning(e)
|
|
64
|
+
logger.error("Unable to start id generator node, please verify your configuration")
|
|
65
|
+
|
|
66
|
+
async def bootstrap(self, handler: Handler):
|
|
67
|
+
"""Register the ID node, start the heartbeat task, and run the WebSocket server.
|
|
68
|
+
|
|
69
|
+
Args:
|
|
70
|
+
handler: Async callable that handles each incoming WebSocket connection.
|
|
71
|
+
"""
|
|
72
|
+
log_header = f"{self.log_header}[bootstrap]"
|
|
73
|
+
await self.node.register()
|
|
74
|
+
asyncio.create_task(self.node.heart_beat_loop(interval_seconds=60))
|
|
75
|
+
logger.info(f"{log_header} websockets server listening on localhost:{self.port}")
|
|
76
|
+
async with serve(handler=handler, host="", port=self.port) as server:
|
|
77
|
+
self.server = server
|
|
78
|
+
await server.serve_forever()
|
|
79
|
+
|
|
80
|
+
|