CheeseAPI 2.0.8b6__tar.gz → 2.0.8b8__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.
- {cheeseapi-2.0.8b6 → cheeseapi-2.0.8b8}/CheeseAPI/__init__.py +1 -1
- {cheeseapi-2.0.8b6 → cheeseapi-2.0.8b8}/CheeseAPI/app.py +73 -17
- {cheeseapi-2.0.8b6 → cheeseapi-2.0.8b8}/CheeseAPI/cors.py +3 -2
- {cheeseapi-2.0.8b6 → cheeseapi-2.0.8b8}/CheeseAPI/file.py +2 -1
- {cheeseapi-2.0.8b6 → cheeseapi-2.0.8b8}/CheeseAPI/printer.py +19 -5
- {cheeseapi-2.0.8b6 → cheeseapi-2.0.8b8}/CheeseAPI/request.py +18 -10
- {cheeseapi-2.0.8b6 → cheeseapi-2.0.8b8}/CheeseAPI/response.py +63 -49
- {cheeseapi-2.0.8b6 → cheeseapi-2.0.8b8}/CheeseAPI/scheduler.py +83 -33
- {cheeseapi-2.0.8b6 → cheeseapi-2.0.8b8}/CheeseAPI/websocket.py +140 -30
- {cheeseapi-2.0.8b6 → cheeseapi-2.0.8b8}/PKG-INFO +2 -2
- {cheeseapi-2.0.8b6 → cheeseapi-2.0.8b8}/pyproject.toml +1 -1
- {cheeseapi-2.0.8b6 → cheeseapi-2.0.8b8}/.gitignore +0 -0
- {cheeseapi-2.0.8b6 → cheeseapi-2.0.8b8}/CheeseAPI/route.py +0 -0
- {cheeseapi-2.0.8b6 → cheeseapi-2.0.8b8}/CheeseAPI/signal.py +0 -0
- {cheeseapi-2.0.8b6 → cheeseapi-2.0.8b8}/CheeseAPI/static.py +0 -0
- {cheeseapi-2.0.8b6 → cheeseapi-2.0.8b8}/CheeseAPI/validator.py +0 -0
- {cheeseapi-2.0.8b6 → cheeseapi-2.0.8b8}/LICENSE +0 -0
- {cheeseapi-2.0.8b6 → cheeseapi-2.0.8b8}/README.md +0 -0
|
@@ -1,7 +1,7 @@
|
|
|
1
1
|
from CheeseAPI.app import CheeseAPI, AppProxy
|
|
2
2
|
from CheeseAPI.printer import Printer
|
|
3
3
|
from CheeseAPI.request import Request, RequestProxy
|
|
4
|
-
from CheeseAPI.response import Response, ResponseProxy, FileResponse
|
|
4
|
+
from CheeseAPI.response import Response, ResponseProxy, FileResponse, RedirectResponse
|
|
5
5
|
from CheeseAPI.websocket import Websocket
|
|
6
6
|
from CheeseAPI.file import File
|
|
7
7
|
from CheeseAPI.route import Route, RouteProxy
|
|
@@ -1,4 +1,4 @@
|
|
|
1
|
-
import os, pathlib, multiprocessing, ssl, socket, asyncio, concurrent.futures, inspect, importlib.util
|
|
1
|
+
import os, re, pathlib, multiprocessing, ssl, socket, asyncio, concurrent.futures, inspect, importlib.util
|
|
2
2
|
from typing import Type, Literal, Callable, AsyncIterable, TYPE_CHECKING
|
|
3
3
|
|
|
4
4
|
import signal, redis
|
|
@@ -44,8 +44,6 @@ class AppProxy:
|
|
|
44
44
|
|
|
45
45
|
self.app._is_running = True
|
|
46
46
|
|
|
47
|
-
self.app.signal.before_server_start.send()
|
|
48
|
-
|
|
49
47
|
self.server_start()
|
|
50
48
|
|
|
51
49
|
waiting_list = self.worker_start()
|
|
@@ -215,6 +213,7 @@ class AppProxy:
|
|
|
215
213
|
|
|
216
214
|
processes = []
|
|
217
215
|
if workers == 1:
|
|
216
|
+
''' 单进程模式下工作进程复用当前进程,after_workers_start 在 worker 内、开始接收请求前触发 '''
|
|
218
217
|
self.worker_running(True)
|
|
219
218
|
else:
|
|
220
219
|
processes: list[multiprocessing.Process] = []
|
|
@@ -223,10 +222,10 @@ class AppProxy:
|
|
|
223
222
|
process.start()
|
|
224
223
|
processes.append(process)
|
|
225
224
|
|
|
226
|
-
|
|
227
|
-
|
|
228
|
-
|
|
229
|
-
|
|
225
|
+
self.after_workers_start(workers)
|
|
226
|
+
self.app.signal.after_workers_start.send(kwargs = {
|
|
227
|
+
'workers': workers
|
|
228
|
+
})
|
|
230
229
|
return processes
|
|
231
230
|
|
|
232
231
|
def before_workers_start(self, workers: int) -> int:
|
|
@@ -272,6 +271,12 @@ class AppProxy:
|
|
|
272
271
|
'is_first': is_first
|
|
273
272
|
})
|
|
274
273
|
|
|
274
|
+
if is_first and self.app.workers == 1:
|
|
275
|
+
self.after_workers_start(self.app.workers)
|
|
276
|
+
self.app.signal.after_workers_start.send(kwargs = {
|
|
277
|
+
'workers': self.app.workers
|
|
278
|
+
})
|
|
279
|
+
|
|
275
280
|
while True:
|
|
276
281
|
client_socket, addr = await loop.sock_accept(self.server_socket)
|
|
277
282
|
loop.create_task(self.client_socket_process(client_socket, addr))
|
|
@@ -289,6 +294,7 @@ class AppProxy:
|
|
|
289
294
|
self.app.printer.app_error(e)
|
|
290
295
|
|
|
291
296
|
async def client_socket_process(self, client_socket: socket.socket, addr: tuple[str, int]):
|
|
297
|
+
request = None
|
|
292
298
|
try:
|
|
293
299
|
loop = asyncio.get_event_loop()
|
|
294
300
|
|
|
@@ -298,6 +304,8 @@ class AppProxy:
|
|
|
298
304
|
async for request, response in self.get_request(client_socket, addr):
|
|
299
305
|
if not response:
|
|
300
306
|
response = await self.get_response(request)
|
|
307
|
+
if inspect.isfunction(request.fn):
|
|
308
|
+
await self.attach_cors_headers(request, response)
|
|
301
309
|
|
|
302
310
|
if not response._proxy:
|
|
303
311
|
response = self.app.ResponseProxy_Class(self.app, response).response
|
|
@@ -313,7 +321,11 @@ class AppProxy:
|
|
|
313
321
|
if client_socket._closed is False:
|
|
314
322
|
client_socket.close()
|
|
315
323
|
except Exception as e:
|
|
316
|
-
|
|
324
|
+
''' 请求尚未创建时无法定位请求上下文,直接上报原始异常,而不是引用未赋值的 request '''
|
|
325
|
+
if request is None:
|
|
326
|
+
self.app.printer.app_error(e)
|
|
327
|
+
else:
|
|
328
|
+
self.app.printer.fn_error(e, request)
|
|
317
329
|
|
|
318
330
|
if client_socket._closed is False:
|
|
319
331
|
client_socket.close()
|
|
@@ -331,6 +343,25 @@ class AppProxy:
|
|
|
331
343
|
'response': response
|
|
332
344
|
})
|
|
333
345
|
|
|
346
|
+
async def attach_cors_headers(self, request: Request, response: Response):
|
|
347
|
+
''' 简单请求:把 CORS 响应头并入路由响应;来源不在白名单内时不附加任何 CORS 头 '''
|
|
348
|
+
if not request.headers.get('origin'):
|
|
349
|
+
return
|
|
350
|
+
|
|
351
|
+
cors_response = self.get_request_cors(request).get_response(request)
|
|
352
|
+
if cors_response.status != 204:
|
|
353
|
+
return
|
|
354
|
+
|
|
355
|
+
for key, value in cors_response.headers.items():
|
|
356
|
+
response.headers.setdefault(key, value)
|
|
357
|
+
|
|
358
|
+
def get_request_cors(self, request: Request) -> CORS:
|
|
359
|
+
''' 取请求命中路由的 CORS 配置,未配置时回退应用级 '''
|
|
360
|
+
route = self.app.route._proxy.get_route(request.method, request.path)
|
|
361
|
+
if route != 404 and route != 405 and route[0]['cors'] is not None:
|
|
362
|
+
return route[0]['cors']
|
|
363
|
+
return self.app.cors
|
|
364
|
+
|
|
334
365
|
async def get_response(self, request: Request) -> Response:
|
|
335
366
|
if inspect.isfunction(request.fn):
|
|
336
367
|
try:
|
|
@@ -352,7 +383,7 @@ class AppProxy:
|
|
|
352
383
|
async def get_request(self, client_socket: socket.socket, addr: tuple[str, int]) -> AsyncIterable[tuple[Request, Response | None]]:
|
|
353
384
|
keep_alive_max_requests = 0
|
|
354
385
|
request = None
|
|
355
|
-
while
|
|
386
|
+
while client_socket._closed is False:
|
|
356
387
|
if request:
|
|
357
388
|
keep_alive_max_requests += 1
|
|
358
389
|
|
|
@@ -364,6 +395,9 @@ class AppProxy:
|
|
|
364
395
|
if request._proxy.protocol is None or not self.app.keep_alive or (request._proxy.protocol == 'HTTP/1.0' and request.headers.get('connection') != 'keep-alive') or (request._proxy.protocol == 'HTTP/1.1' and request.headers.get('connection') == 'close'):
|
|
365
396
|
break
|
|
366
397
|
|
|
398
|
+
if keep_alive_max_requests >= self.app.keep_alive_max_requests:
|
|
399
|
+
break
|
|
400
|
+
|
|
367
401
|
client_socket, addr = await self.before_request(client_socket, addr)
|
|
368
402
|
await self.app.signal.before_request.async_send(kwargs = {
|
|
369
403
|
'client_socket': client_socket,
|
|
@@ -426,7 +460,9 @@ class AppProxy:
|
|
|
426
460
|
relative_path = relative_path[1:]
|
|
427
461
|
|
|
428
462
|
path = os.path.join(_path, relative_path)
|
|
429
|
-
|
|
463
|
+
static_root = os.path.abspath(_path)
|
|
464
|
+
target_path = os.path.abspath(path)
|
|
465
|
+
if target_path != static_root and target_path.startswith(f'{static_root}{os.sep}') is False:
|
|
430
466
|
return Response(status = 403)
|
|
431
467
|
|
|
432
468
|
if os.path.exists(path):
|
|
@@ -440,13 +476,33 @@ class AppProxy:
|
|
|
440
476
|
return Response(status = 404)
|
|
441
477
|
|
|
442
478
|
async def get_cors_response(self, request: Request) -> Response:
|
|
443
|
-
|
|
444
|
-
if
|
|
445
|
-
|
|
446
|
-
|
|
447
|
-
|
|
448
|
-
|
|
449
|
-
|
|
479
|
+
''' 预检请求(OPTIONS + Origin + Access-Control-Request-Method)返回 CORS 响应,其余 405 场景保持裸 405 '''
|
|
480
|
+
if request.method != 'OPTIONS' or not request.headers.get('origin') or not request.headers.get('access-control-request-method'):
|
|
481
|
+
return Response(status = 405)
|
|
482
|
+
|
|
483
|
+
cors = self.get_route_cors(request.path) or self.app.cors
|
|
484
|
+
return cors.get_response(request)
|
|
485
|
+
|
|
486
|
+
def get_route_cors(self, path: str) -> CORS | None:
|
|
487
|
+
''' 按路径取路由级 CORS 配置(与请求方法无关),供预检请求使用 '''
|
|
488
|
+
routes = self.app.route.routes.get(path)
|
|
489
|
+
if routes is None:
|
|
490
|
+
for _path, dynamic_routes in self.app.route._proxy.dynamic_routes.items():
|
|
491
|
+
if re.match(_path, path) is not None:
|
|
492
|
+
routes = dynamic_routes
|
|
493
|
+
break
|
|
494
|
+
|
|
495
|
+
if not routes:
|
|
496
|
+
return None
|
|
497
|
+
|
|
498
|
+
if routes.get('OPTIONS') and routes['OPTIONS']['cors'] is not None:
|
|
499
|
+
return routes['OPTIONS']['cors']
|
|
500
|
+
|
|
501
|
+
for route in routes.values():
|
|
502
|
+
if route['cors'] is not None:
|
|
503
|
+
return route['cors']
|
|
504
|
+
|
|
505
|
+
return None
|
|
450
506
|
|
|
451
507
|
async def before_request(self, client_socket: socket.socket, addr: tuple[str, int]) -> tuple[socket.socket, tuple]:
|
|
452
508
|
return client_socket, addr
|
|
@@ -27,8 +27,9 @@ class CORS:
|
|
|
27
27
|
|
|
28
28
|
if origin in self.allow_origins:
|
|
29
29
|
headers['access-control-allow-origin'] = origin
|
|
30
|
-
elif '*' in self.allow_origins
|
|
31
|
-
|
|
30
|
+
elif '*' in self.allow_origins:
|
|
31
|
+
''' 通配来源与凭证不能同用,携带凭证时回显具体 origin '''
|
|
32
|
+
headers['access-control-allow-origin'] = origin if self.allow_credentials else '*'
|
|
32
33
|
|
|
33
34
|
headers['access-control-allow-methods'] = ', '.join(self.allow_methods)
|
|
34
35
|
|
|
@@ -29,6 +29,7 @@ class File:
|
|
|
29
29
|
elif len(args) == 2:
|
|
30
30
|
self._name = args[0]
|
|
31
31
|
self._data = args[1]
|
|
32
|
+
self._data_in_file = False
|
|
32
33
|
|
|
33
34
|
def save(self, path: str, update_path: bool = False, data_in_file: bool = False):
|
|
34
35
|
'''
|
|
@@ -65,7 +66,7 @@ class File:
|
|
|
65
66
|
|
|
66
67
|
@property
|
|
67
68
|
def data(self) -> bytes:
|
|
68
|
-
if self._data:
|
|
69
|
+
if self._data is not None:
|
|
69
70
|
return self._data
|
|
70
71
|
else:
|
|
71
72
|
with open(self._path, 'rb') as f:
|
|
@@ -12,6 +12,9 @@ if TYPE_CHECKING:
|
|
|
12
12
|
|
|
13
13
|
HTTP_STATUS_COLOR: tuple[str] = ('blue', 'green', 'cyan', 'yellow', 'red')
|
|
14
14
|
|
|
15
|
+
MAX_ERROR_LINES: int = 50
|
|
16
|
+
''' 单条错误日志中堆栈的最大行数;超出部分会被截断,避免异常堆栈无限增长 '''
|
|
17
|
+
|
|
15
18
|
class Printer:
|
|
16
19
|
__slots__ = ('_app', '_progress_bar')
|
|
17
20
|
|
|
@@ -33,8 +36,19 @@ class Printer:
|
|
|
33
36
|
def server_start(self):
|
|
34
37
|
self.app.logger.print('START', f'CheeseAPI is running on {self.app.host}:{self.app.port}', f'CheeseAPI is running on <cyan>{self.app.host}:{self.app.port}</cyan>')
|
|
35
38
|
|
|
39
|
+
def _error(self) -> str:
|
|
40
|
+
'''
|
|
41
|
+
格式化当前异常的堆栈
|
|
42
|
+
|
|
43
|
+
限制最大行数,避免反复报错时堆栈越来越长把日志写爆
|
|
44
|
+
'''
|
|
45
|
+
lines = traceback.format_exc()[:-1].splitlines()
|
|
46
|
+
if len(lines) > MAX_ERROR_LINES:
|
|
47
|
+
lines = lines[:MAX_ERROR_LINES] + [f'... truncated {len(lines) - MAX_ERROR_LINES} lines ...']
|
|
48
|
+
return '\n'.join(lines).replace('\n', '\n ')
|
|
49
|
+
|
|
36
50
|
def app_error(self, e: Exception):
|
|
37
|
-
error =
|
|
51
|
+
error = self._error()
|
|
38
52
|
self.app.logger.error(f'An error occurred causing the server to stop:\n {error}', f'An error occurred causing the server to stop:\n {self.app.logger.encode(error)}')
|
|
39
53
|
|
|
40
54
|
def server_stop(self):
|
|
@@ -44,11 +58,11 @@ class Printer:
|
|
|
44
58
|
self.app.logger.print('STOP', 'CheeseAPI has stopped')
|
|
45
59
|
|
|
46
60
|
def fn_error(self, e: Exception, request: 'Request'):
|
|
47
|
-
error =
|
|
61
|
+
error = self._error()
|
|
48
62
|
self.app.logger.danger(f'An error occurred causing the {request.ip} visited {request.method} {request.full_path}:\n {error}', f'An error occurred causing the <cyan>{request.ip}</cyan> visited <cyan>{request.method} {self.app.logger.encode(request.full_path)}</cyan>:\n {self.app.logger.encode(error)}')
|
|
49
63
|
|
|
50
64
|
def websocket_error(self, e: Exception, websocket: 'Websocket'):
|
|
51
|
-
error =
|
|
65
|
+
error = self._error()
|
|
52
66
|
self.app.logger.danger(f'An error occurred causing the {websocket.request.ip} disconnected {websocket.request.method} {websocket.request.full_path}:\n {error}', f'An error occurred causing the <cyan>{websocket.request.ip}</cyan> disconnected <cyan>{websocket.request.method} {self.app.logger.encode(websocket.request.full_path)}</cyan>:\n {self.app.logger.encode(error)}')
|
|
53
67
|
|
|
54
68
|
def response(self, request: 'Request', response: 'Response'):
|
|
@@ -65,11 +79,11 @@ class Printer:
|
|
|
65
79
|
self.app.logger.print('WEBSOCKET', f'The {websocket.request.ip} disconnected {websocket.request.method} {websocket.request.full_path}', f'The <cyan>{websocket.request.ip}</cyan> disconnected <cyan>{websocket.request.method} {self.app.logger.encode(websocket.request.full_path)}</cyan>')
|
|
66
80
|
|
|
67
81
|
def websocket_message_error(self, e: Exception, websocket: 'Websocket'):
|
|
68
|
-
error =
|
|
82
|
+
error = self._error()
|
|
69
83
|
self.app.logger.danger(f'An error occurred causing the {websocket.request.ip} received a message to {websocket.request.method} {websocket.request.full_path}:\n {error}', f'An error occurred causing the <cyan>{websocket.request.ip}</cyan> receive a message to <cyan>{websocket.request.method} {self.app.logger.encode(websocket.request.full_path)}</cyan>:\n {self.app.logger.encode(error)}')
|
|
70
84
|
|
|
71
85
|
def scheduler_error(self, e: Exception, task: 'Task'):
|
|
72
|
-
error =
|
|
86
|
+
error = self._error()
|
|
73
87
|
self.app.logger.danger(f'An error occurred in the scheduled task {task.key} running:\n {error}', f'An error occurred in the scheduled task running <green>{self.app.logger.encode(task.key)}</green>:\n {self.app.logger.encode(error)}')
|
|
74
88
|
|
|
75
89
|
@property
|
|
@@ -169,7 +169,12 @@ class RequestProxy:
|
|
|
169
169
|
range_part = range_part.strip()
|
|
170
170
|
if '-' in range_part:
|
|
171
171
|
start, end = range_part.split('-', 1)
|
|
172
|
-
|
|
172
|
+
if start:
|
|
173
|
+
self.request.ranges.append((int(start), int(end) if end else None))
|
|
174
|
+
elif end:
|
|
175
|
+
self.request.ranges.append((-int(end), None)) # 后缀 Range(`bytes=-50`):负起点表示从末尾倒数,末尾 50 字节
|
|
176
|
+
else:
|
|
177
|
+
self.request.ranges.append((0, None))
|
|
173
178
|
|
|
174
179
|
if 'upgrade' in self.request.headers and self.request.headers['upgrade'] == 'websocket':
|
|
175
180
|
self.request._method = 'WEBSOCKET'
|
|
@@ -246,17 +251,20 @@ class RequestProxy:
|
|
|
246
251
|
if self.request.body is None:
|
|
247
252
|
return
|
|
248
253
|
|
|
254
|
+
body = self.request.body
|
|
249
255
|
content_type = self.request.headers.get('content-type')
|
|
250
|
-
|
|
251
|
-
|
|
252
|
-
|
|
253
|
-
self.request.
|
|
254
|
-
elif
|
|
256
|
+
media_type = (content_type or '').split(';')[0].strip().lower() or None # 忽略 `; charset=utf-8` 等参数
|
|
257
|
+
|
|
258
|
+
if media_type == 'text/plain' or media_type is None:
|
|
259
|
+
self.request._body = body.decode()
|
|
260
|
+
elif media_type == 'application/json':
|
|
261
|
+
self.request._json = json.loads(body)
|
|
262
|
+
elif media_type == 'application/x-www-form-urlencoded':
|
|
255
263
|
self.request._form = {
|
|
256
|
-
key: value[0] for key, value in urllib.parse.parse_qs(
|
|
264
|
+
key: value[0] for key, value in urllib.parse.parse_qs(body.decode()).items()
|
|
257
265
|
}
|
|
258
|
-
elif
|
|
259
|
-
for part in
|
|
266
|
+
elif media_type == 'multipart/form-data':
|
|
267
|
+
for part in body.split(f'--{content_type.split("boundary=")[1].strip()}'.encode()):
|
|
260
268
|
if part == b'' or part == b'--\r\n':
|
|
261
269
|
continue
|
|
262
270
|
|
|
@@ -290,4 +298,4 @@ class RequestProxy:
|
|
|
290
298
|
if self.request.headers.get('content-disposition'):
|
|
291
299
|
match = re.search(r'filename="([^"]*)"', self.request.headers['content-disposition'])
|
|
292
300
|
if match:
|
|
293
|
-
self.request._file = File(match.group(1),
|
|
301
|
+
self.request._file = File(match.group(1), body) # 始终用原始 bytes,避免 text/plain 分支已解码成 str
|
|
@@ -1,4 +1,4 @@
|
|
|
1
|
-
import socket, asyncio, datetime, json, zlib, gzip, os, mimetypes, uuid
|
|
1
|
+
import socket, asyncio, datetime, json, zlib, gzip, os, mimetypes, uuid
|
|
2
2
|
from typing import TYPE_CHECKING, TypedDict, Literal, AsyncIterable
|
|
3
3
|
|
|
4
4
|
import brotli, zstandard
|
|
@@ -127,7 +127,7 @@ class RedirectResponse(Response):
|
|
|
127
127
|
headers = {}
|
|
128
128
|
headers['location'] = location
|
|
129
129
|
|
|
130
|
-
super().__init__(
|
|
130
|
+
super().__init__(body, status, headers)
|
|
131
131
|
|
|
132
132
|
class FileResponse(Response):
|
|
133
133
|
__slots__ = ('file', 'preview', 'transmission_type', 'chunked_size')
|
|
@@ -169,21 +169,34 @@ class ResponseProxy:
|
|
|
169
169
|
status, headers, data = await anext(gen)
|
|
170
170
|
|
|
171
171
|
bytes = [f'HTTP/1.1 {status} {HTTP_STATUS[status]}']
|
|
172
|
-
|
|
172
|
+
for key, value in headers.items():
|
|
173
|
+
if isinstance(value, list): # 某些响应头(如 set-cookie)需要写成多行
|
|
174
|
+
bytes.extend(f'{key}: {_value}' for _value in value)
|
|
175
|
+
else:
|
|
176
|
+
bytes.append(f'{key}: {value}')
|
|
173
177
|
bytes.extend(['', ''])
|
|
174
178
|
bytes = '\r\n'.join(bytes).encode()
|
|
175
179
|
if not no_body:
|
|
176
|
-
if
|
|
177
|
-
bytes +=
|
|
180
|
+
if headers.get('transfer-encoding') == 'chunked':
|
|
181
|
+
bytes += f'{len(data):x}'.encode() + b'\r\n' + data + b'\r\n'
|
|
178
182
|
else:
|
|
179
183
|
bytes += data
|
|
180
184
|
await loop.sock_sendall(client_socket, bytes)
|
|
181
185
|
|
|
182
186
|
if not no_body:
|
|
183
|
-
if
|
|
184
|
-
|
|
185
|
-
|
|
186
|
-
|
|
187
|
+
if headers.get('transfer-encoding') == 'chunked':
|
|
188
|
+
'''
|
|
189
|
+
分块传输:无论生成器是否抛异常,都必须补发 0 长度块结束分块流,
|
|
190
|
+
否则客户端会认为响应被截断
|
|
191
|
+
'''
|
|
192
|
+
try:
|
|
193
|
+
async for _, _, data in gen:
|
|
194
|
+
await loop.sock_sendall(client_socket, f'{len(data):x}'.encode() + b'\r\n' + data + b'\r\n')
|
|
195
|
+
finally:
|
|
196
|
+
try:
|
|
197
|
+
await loop.sock_sendall(client_socket, b'0\r\n\r\n')
|
|
198
|
+
except Exception:
|
|
199
|
+
...
|
|
187
200
|
elif self.request.ranges:
|
|
188
201
|
async for _, _, data in gen:
|
|
189
202
|
await loop.sock_sendall(client_socket, data)
|
|
@@ -192,19 +205,14 @@ class ResponseProxy:
|
|
|
192
205
|
self.app.printer.response(self.request, self.response)
|
|
193
206
|
|
|
194
207
|
async def get_status(self, status: int, headers: dict[str, str], body: dict | list | str | bytes | None) -> tuple[int, dict[str, str], dict | list | str | bytes | None]:
|
|
195
|
-
if isinstance(self.response, FileResponse):
|
|
196
|
-
|
|
197
|
-
|
|
198
|
-
|
|
199
|
-
|
|
200
|
-
|
|
201
|
-
|
|
202
|
-
|
|
203
|
-
if len(self.response.file.data) < max_range:
|
|
204
|
-
status = 416
|
|
205
|
-
else:
|
|
206
|
-
if os.path.getsize(self.response.file.path) < max_range:
|
|
207
|
-
status = 416
|
|
208
|
+
if isinstance(self.response, FileResponse) and self.request.ranges:
|
|
209
|
+
max_range = -1
|
|
210
|
+
for range in self.request.ranges:
|
|
211
|
+
if range[1] is not None:
|
|
212
|
+
max_range = max(max_range, range[1])
|
|
213
|
+
max_range += 1
|
|
214
|
+
size = len(self.response.file.data) if self.response.file._data is not None else os.path.getsize(self.response.file.path)
|
|
215
|
+
status = 416 if size < max_range else 206
|
|
208
216
|
|
|
209
217
|
return status, headers, body
|
|
210
218
|
|
|
@@ -236,7 +244,7 @@ class ResponseProxy:
|
|
|
236
244
|
headers['content-type'] = f'{mime_type}; charset=utf-8'
|
|
237
245
|
headers['content-disposition'] = f'{"inline" if self.response.preview and mime_type in PREVIEWABLE_TYPES else "attachment"}; filename="{self.response.file.name}"'
|
|
238
246
|
|
|
239
|
-
if isinstance(
|
|
247
|
+
if isinstance(body, AsyncIterable):
|
|
240
248
|
if 'transfer-encoding' not in headers:
|
|
241
249
|
headers['transfer-encoding'] = 'chunked'
|
|
242
250
|
|
|
@@ -265,6 +273,8 @@ class ResponseProxy:
|
|
|
265
273
|
else:
|
|
266
274
|
encoding_quality = True
|
|
267
275
|
encoding_split[1] = float(encoding_split[1].split('=')[1])
|
|
276
|
+
if float(encoding_split[1]) <= 0: # q=0 表示该算法不可接受
|
|
277
|
+
continue
|
|
268
278
|
encodings.append(encoding_split)
|
|
269
279
|
encodings.sort(key = lambda x: float(x[1]), reverse = True)
|
|
270
280
|
encodings = [encoding[0] for encoding in encodings]
|
|
@@ -304,51 +314,55 @@ class ResponseProxy:
|
|
|
304
314
|
if cookie['http_only']:
|
|
305
315
|
_cookie += '; HttpOnly'
|
|
306
316
|
cookies.append(_cookie)
|
|
307
|
-
headers['set-cookie'] =
|
|
317
|
+
headers['set-cookie'] = cookies # set-cookie 需要写成多行,此处保留列表,由 ResponseProxy.send 展开
|
|
308
318
|
|
|
309
319
|
return status, headers, body
|
|
310
320
|
|
|
311
321
|
async def get_body(self, status: int, headers: dict[str, str], body: dict | list | str | bytes | AsyncIterable | None) -> AsyncIterable[tuple[int, dict[str, str], bytes]]:
|
|
312
322
|
if type(self.response) is FileResponse and self.request.ranges and status != 416:
|
|
313
|
-
if self.response.file._data is None:
|
|
323
|
+
if self.response.file._data is not None:
|
|
324
|
+
size = len(self.response.file.data)
|
|
325
|
+
else:
|
|
326
|
+
size = os.path.getsize(self.response.file.path)
|
|
314
327
|
handler = open(self.response.file.path, 'rb')
|
|
328
|
+
|
|
315
329
|
if len(self.request.ranges) == 1:
|
|
330
|
+
start, end = self.request.ranges[0]
|
|
331
|
+
end = size - 1 if end is None else end # 区间为闭区间,末尾含在响应内
|
|
316
332
|
if self.response.file._data is not None:
|
|
317
|
-
|
|
318
|
-
data = self.response.file.data[self.request.ranges[0][0] or 0:self.request.ranges[0][1] or len(self.response.file.data) + 1]
|
|
333
|
+
data = self.response.file.data[start:end + 1]
|
|
319
334
|
else:
|
|
320
|
-
|
|
321
|
-
handler.
|
|
322
|
-
data = handler.read((self.request.ranges[0][1] or size) - (self.request.ranges[0][0]))
|
|
335
|
+
handler.seek(start)
|
|
336
|
+
data = handler.read(end + 1 - start)
|
|
323
337
|
handler.close()
|
|
338
|
+
|
|
324
339
|
headers['content-length'] = str(len(data))
|
|
325
|
-
headers['content-range'] = f'bytes {
|
|
340
|
+
headers['content-range'] = f'bytes {start}-{end}/{size}'
|
|
326
341
|
yield status, headers, data
|
|
327
342
|
else:
|
|
328
343
|
boundary = uuid.uuid4().hex
|
|
329
344
|
content_type = headers['content-type']
|
|
330
345
|
headers['content-type'] = f'multipart/byteranges; boundary={boundary}'
|
|
331
|
-
if self.response.file._data is not None:
|
|
332
|
-
size = len(self.response.file.data)
|
|
333
|
-
else:
|
|
334
|
-
size = os.path.getsize(self.response.file.path)
|
|
335
|
-
|
|
336
|
-
content_length = 0
|
|
337
|
-
for range in self.request.ranges:
|
|
338
|
-
content_length += 2 + 32 + 2 + 14 + len(content_type) + 2 + 21 + (1 if range[0] == 0 else int(math.log10(range[0]))) + 1 + 1 + int(math.log10(range[1] or size)) + 1 + 1 + int(math.log10(size)) + 1 + 4 + (range[1] or size) - range[0] + 1
|
|
339
|
-
headers['content-length'] = str(content_length)
|
|
340
346
|
|
|
341
|
-
|
|
342
|
-
|
|
347
|
+
parts = []
|
|
348
|
+
for start, end in self.request.ranges:
|
|
349
|
+
end = size - 1 if end is None else end
|
|
343
350
|
if self.response.file._data is not None:
|
|
344
|
-
data
|
|
351
|
+
data = self.response.file.data[start:end + 1]
|
|
345
352
|
else:
|
|
346
|
-
handler.seek(
|
|
347
|
-
data
|
|
348
|
-
|
|
353
|
+
handler.seek(start)
|
|
354
|
+
data = handler.read(end + 1 - start)
|
|
355
|
+
parts.append(b''.join([b'--', boundary.encode(), b'\r\n', b'content-type: ', content_type.encode(), b'\r\n', b'content-range: bytes ', str(start).encode(), b'-', str(end).encode(), b'/', str(size).encode(), b'\r\n\r\n', data]))
|
|
356
|
+
|
|
349
357
|
if self.response.file._data is None:
|
|
350
358
|
handler.close()
|
|
351
|
-
|
|
359
|
+
|
|
360
|
+
trailer = b'--' + boundary.encode() + b'--'
|
|
361
|
+
headers['content-length'] = str(sum(len(part) for part in parts) + len(trailer))
|
|
362
|
+
|
|
363
|
+
for part in parts:
|
|
364
|
+
yield status, headers, part
|
|
365
|
+
yield status, headers, trailer
|
|
352
366
|
else:
|
|
353
367
|
if isinstance(body, AsyncIterable):
|
|
354
368
|
data = await anext(body)
|
|
@@ -386,11 +400,11 @@ class ResponseProxy:
|
|
|
386
400
|
|
|
387
401
|
async def get_encode_body(self, status: int, headers: dict[str, str], body: bytes) -> tuple[int, dict[str, str], bytes]:
|
|
388
402
|
content_length = headers.get('content-length')
|
|
389
|
-
if content_length and int(content_length) < self.app.compress_min_length:
|
|
403
|
+
if self.response.compress is None and content_length and int(content_length) < self.app.compress_min_length:
|
|
390
404
|
if 'content-encoding' in headers:
|
|
391
405
|
del headers['content-encoding']
|
|
392
406
|
|
|
393
|
-
if 'content-encoding' in headers and 'content-length' in headers and (int(headers['content-length'])
|
|
407
|
+
if 'content-encoding' in headers and 'content-length' in headers and (self.response.compress is not None or int(headers['content-length']) >= self.app.compress_min_length):
|
|
394
408
|
compress_level = self.response.compress_level if self.response.compress_level is not None else self.app.compress_level
|
|
395
409
|
if headers['content-encoding'] == 'gzip':
|
|
396
410
|
body = gzip.compress(body, compress_level)
|
|
@@ -1,5 +1,6 @@
|
|
|
1
1
|
import inspect, os
|
|
2
2
|
import datetime, uuid, threading, multiprocessing, asyncio, time, json
|
|
3
|
+
from queue import Empty
|
|
3
4
|
from typing import Callable, Literal, TYPE_CHECKING
|
|
4
5
|
|
|
5
6
|
import redis, redis.exceptions
|
|
@@ -13,22 +14,25 @@ class Task:
|
|
|
13
14
|
@classmethod
|
|
14
15
|
def from_dict(cls, data: dict[str, any], _scheduler_proxy) -> 'Task':
|
|
15
16
|
instance = cls.__new__(cls)
|
|
17
|
+
|
|
16
18
|
for key, value in data.items():
|
|
17
19
|
if key == '_queue':
|
|
18
|
-
|
|
19
|
-
|
|
20
|
-
|
|
21
|
-
else:
|
|
22
|
-
value = multiprocessing.get_context('spawn').Queue()
|
|
20
|
+
value = multiprocessing.get_context('spawn').Queue()
|
|
21
|
+
elif key == '_stop_event':
|
|
22
|
+
value = threading.Event()
|
|
23
23
|
elif key == 'first_run_timer':
|
|
24
24
|
value = datetime.datetime.fromtimestamp(value) if value else None
|
|
25
25
|
elif key == '_last_run_timer':
|
|
26
26
|
value = datetime.datetime.fromtimestamp(value) if value else None
|
|
27
27
|
setattr(instance, key, value)
|
|
28
|
+
|
|
29
|
+
if '_running_remote' not in data:
|
|
30
|
+
instance._running_remote = not data.get('_queue', True) # 旧数据只有 `_queue`(非空代表未运行)
|
|
31
|
+
|
|
28
32
|
setattr(instance, '_scheduler_proxy', _scheduler_proxy)
|
|
29
33
|
return instance
|
|
30
34
|
|
|
31
|
-
__slots__ = ('fn', 'interval_time', 'first_run_timer', 'expected_run_num', '_key', 'run_type', 'args', 'kwargs', 'auto_remove', '_last_run_timer', '_last_run_time', '_run_num', '_handler', '_queue', '_scheduler_proxy', 'timeout')
|
|
35
|
+
__slots__ = ('fn', 'interval_time', 'first_run_timer', 'expected_run_num', '_key', 'run_type', 'args', 'kwargs', 'auto_remove', '_last_run_timer', '_last_run_time', '_run_num', '_handler', '_queue', '_stop_event', '_running_remote', '_scheduler_proxy', 'timeout')
|
|
32
36
|
|
|
33
37
|
def __init__(self, fn: Callable, interval_time: float, *, first_run_timer: datetime.datetime | float | None = None, expected_run_num: int | None = None, key: str | None = None, run_type: Literal['THREAD', 'PROCESS', 'ASYNC'] = 'THREAD', args: tuple = (), kwargs: dict = {}, auto_remove: bool = False, timeout: float | None = None, _scheduler_proxy: 'SchedulerProxy'):
|
|
34
38
|
'''
|
|
@@ -67,23 +71,66 @@ class Task:
|
|
|
67
71
|
self._last_run_time: float | None = None
|
|
68
72
|
self._run_num: int = 0
|
|
69
73
|
self._handler: threading.Thread | multiprocessing.Process | asyncio.Task | None = None
|
|
74
|
+
|
|
70
75
|
self._queue = multiprocessing.get_context('spawn').Queue()
|
|
76
|
+
''' 进程任务的停止信号队列(父进程放入、子进程取出) '''
|
|
77
|
+
self._stop_event = threading.Event()
|
|
78
|
+
''' 线程 / 协程任务的停止信号(同进程内使用) '''
|
|
79
|
+
self._running_remote: bool = False
|
|
80
|
+
''' 任务是否运行在其它进程(由 `sync_server` 反序列化而来的任务) '''
|
|
81
|
+
|
|
82
|
+
'''
|
|
83
|
+
停止信号的收发
|
|
84
|
+
|
|
85
|
+
不使用 `multiprocessing.Queue.qsize()`(macOS 未实现,调用即 `NotImplementedError`):
|
|
86
|
+
线程 / 协程任务用 `threading.Event`,进程任务用队列的 `empty()` / `get_nowait()` 判断与取出信号。
|
|
87
|
+
'''
|
|
88
|
+
|
|
89
|
+
def _stop(self):
|
|
90
|
+
''' 发送停止信号 '''
|
|
91
|
+
|
|
92
|
+
if self.run_type == 'PROCESS':
|
|
93
|
+
self._queue.put(None)
|
|
94
|
+
else:
|
|
95
|
+
self._stop_event.set()
|
|
96
|
+
|
|
97
|
+
def _reset_stop(self):
|
|
98
|
+
''' 清除停止信号(启动任务前调用) '''
|
|
99
|
+
|
|
100
|
+
if self.run_type == 'PROCESS':
|
|
101
|
+
while True:
|
|
102
|
+
try:
|
|
103
|
+
self._queue.get_nowait()
|
|
104
|
+
except Empty:
|
|
105
|
+
break
|
|
106
|
+
else:
|
|
107
|
+
self._stop_event.clear()
|
|
108
|
+
|
|
109
|
+
def _stops(self) -> bool:
|
|
110
|
+
''' 是否收到了停止信号 '''
|
|
111
|
+
|
|
112
|
+
if self.run_type == 'PROCESS':
|
|
113
|
+
return not self._queue.empty()
|
|
114
|
+
return self._stop_event.is_set()
|
|
71
115
|
|
|
72
116
|
def __getstate__(self) -> tuple[None, dict[str, any]]:
|
|
73
117
|
state = {
|
|
74
118
|
key: getattr(self, key) for key in self.__slots__
|
|
75
119
|
}
|
|
76
120
|
state['_handler'] = None
|
|
121
|
+
state['_stop_event'] = None # `threading.Event` 不能被 pickle,进程任务只用队列信号
|
|
77
122
|
return None, state
|
|
78
123
|
|
|
79
124
|
def _to_dict(self) -> dict[str, any]:
|
|
80
125
|
data = {
|
|
81
126
|
key: getattr(self, key) for key in self.__slots__
|
|
82
127
|
}
|
|
83
|
-
data['_queue'] =
|
|
128
|
+
data['_queue'] = None
|
|
129
|
+
data['_stop_event'] = None
|
|
130
|
+
data['_handler'] = None
|
|
131
|
+
data['_running_remote'] = self.is_running
|
|
84
132
|
data['first_run_timer'] = self.first_run_timer.timestamp() if self.first_run_timer else None
|
|
85
133
|
data['_last_run_timer'] = self._last_run_timer.timestamp() if self._last_run_timer else None
|
|
86
|
-
data['_handler'] = None
|
|
87
134
|
data['fn'] = None
|
|
88
135
|
data['args'] = tuple()
|
|
89
136
|
data['kwargs'] = {}
|
|
@@ -95,7 +142,7 @@ class Task:
|
|
|
95
142
|
|
|
96
143
|
async def async_start(self):
|
|
97
144
|
if not self._handler:
|
|
98
|
-
self._handler = asyncio.create_task(self._scheduler_proxy.async_task_processing(self.key, self.
|
|
145
|
+
self._handler = asyncio.create_task(self._scheduler_proxy.async_task_processing(self.key, self.fn, *self.args, **self.kwargs))
|
|
99
146
|
|
|
100
147
|
@property
|
|
101
148
|
def key(self) -> str:
|
|
@@ -105,7 +152,7 @@ class Task:
|
|
|
105
152
|
def run_num_completed(self) -> bool:
|
|
106
153
|
if self.expected_run_num is None:
|
|
107
154
|
return False
|
|
108
|
-
return self.
|
|
155
|
+
return self.run_num >= self.expected_run_num
|
|
109
156
|
|
|
110
157
|
@property
|
|
111
158
|
def last_run_timer(self) -> datetime.datetime | None:
|
|
@@ -129,7 +176,11 @@ class Task:
|
|
|
129
176
|
def is_running(self) -> bool:
|
|
130
177
|
''' 任务是否在运行中 '''
|
|
131
178
|
|
|
132
|
-
|
|
179
|
+
if self._handler is None:
|
|
180
|
+
return self._running_remote
|
|
181
|
+
if isinstance(self._handler, asyncio.Task):
|
|
182
|
+
return not self._handler.done()
|
|
183
|
+
return self._handler.is_alive()
|
|
133
184
|
|
|
134
185
|
class Scheduler:
|
|
135
186
|
__slots__ = ('_proxy',)
|
|
@@ -310,7 +361,6 @@ class SchedulerProxy:
|
|
|
310
361
|
coro.close()
|
|
311
362
|
|
|
312
363
|
task = Task(fn, interval_time, first_run_timer = first_run_timer, expected_run_num = expected_run_num, key = key, run_type = run_type, args = args, kwargs = kwargs, auto_remove = auto_remove, timeout = timeout, _scheduler_proxy = self)
|
|
313
|
-
task._queue.put(None)
|
|
314
364
|
|
|
315
365
|
if task.key in self.get_tasks():
|
|
316
366
|
raise KeyError(f'Task with key "{task.key}" already exists')
|
|
@@ -339,7 +389,6 @@ class SchedulerProxy:
|
|
|
339
389
|
coro.close()
|
|
340
390
|
|
|
341
391
|
task = Task(fn, interval_time = interval_time, first_run_timer = first_run_timer, expected_run_num = expected_run_num, key = key, run_type = 'ASYNC', args = args, kwargs = kwargs, auto_remove = auto_remove, timeout = timeout, _scheduler_proxy = self)
|
|
342
|
-
task._queue.put(None)
|
|
343
392
|
|
|
344
393
|
if task.key in await self.async_get_tasks():
|
|
345
394
|
raise KeyError(f'Task with key "{task.key}" already exists')
|
|
@@ -356,16 +405,12 @@ class SchedulerProxy:
|
|
|
356
405
|
return _fn
|
|
357
406
|
return wrapper
|
|
358
407
|
|
|
359
|
-
def task_processing(self, task: Task,
|
|
408
|
+
def task_processing(self, task: Task, fn, *args, **kwargs):
|
|
360
409
|
try:
|
|
361
|
-
queue.get()
|
|
362
|
-
|
|
363
|
-
if static.scheduler_sync_servers:
|
|
364
|
-
self.get_task(task.key)._queue.get()
|
|
365
410
|
if task.first_run_timer:
|
|
366
411
|
time.sleep(max(0, task.first_run_timer.timestamp() - time.time()))
|
|
367
412
|
|
|
368
|
-
while not
|
|
413
|
+
while not task._stops():
|
|
369
414
|
now = time.time()
|
|
370
415
|
|
|
371
416
|
try:
|
|
@@ -378,7 +423,7 @@ class SchedulerProxy:
|
|
|
378
423
|
except Exception as e:
|
|
379
424
|
self.app.printer.scheduler_error(e, task)
|
|
380
425
|
|
|
381
|
-
if
|
|
426
|
+
if task._stops():
|
|
382
427
|
break
|
|
383
428
|
|
|
384
429
|
task._last_run_time = time.time() - now
|
|
@@ -390,6 +435,9 @@ class SchedulerProxy:
|
|
|
390
435
|
sync_server.hset('CheeseAPI_scheduler_tasks', task.key, json.dumps(task._to_dict()))
|
|
391
436
|
sync_server.hpexpire('CheeseAPI_scheduler_tasks', int(task.timeout * 1000), task.key)
|
|
392
437
|
|
|
438
|
+
if task.run_num_completed:
|
|
439
|
+
break
|
|
440
|
+
|
|
393
441
|
time.sleep(max(0, task.interval_time - time.time() + now))
|
|
394
442
|
except (KeyboardInterrupt, SystemExit):
|
|
395
443
|
...
|
|
@@ -399,16 +447,16 @@ class SchedulerProxy:
|
|
|
399
447
|
if _redis.hexists('CheeseAPI_scheduler_tasks', task.key):
|
|
400
448
|
_redis.hpersist('CheeseAPI_scheduler_tasks', task.key)
|
|
401
449
|
|
|
402
|
-
|
|
403
|
-
|
|
450
|
+
if task.run_type == 'PROCESS':
|
|
451
|
+
task._reset_stop() # 取出残留的停止信号,保证下次启动时队列是干净的
|
|
404
452
|
|
|
453
|
+
async def async_task_processing(self, key: str, fn, *args, **kwargs):
|
|
405
454
|
task = await self.async_get_task(key)
|
|
406
|
-
|
|
407
|
-
task._queue.get()
|
|
455
|
+
|
|
408
456
|
if task.first_run_timer:
|
|
409
457
|
await asyncio.sleep(max(0, task.first_run_timer.timestamp() - time.time()))
|
|
410
458
|
|
|
411
|
-
while not
|
|
459
|
+
while not task._stops():
|
|
412
460
|
now = time.time()
|
|
413
461
|
|
|
414
462
|
try:
|
|
@@ -421,7 +469,7 @@ class SchedulerProxy:
|
|
|
421
469
|
except Exception as e:
|
|
422
470
|
self.app.printer.scheduler_error(e, task)
|
|
423
471
|
|
|
424
|
-
if
|
|
472
|
+
if task._stops():
|
|
425
473
|
break
|
|
426
474
|
|
|
427
475
|
task._last_run_time = time.time() - now
|
|
@@ -469,10 +517,11 @@ class SchedulerProxy:
|
|
|
469
517
|
redis.Redis(connection_pool = static.scheduler_sync_servers[0]).publish('CheeseAPI_scheduler', json.dumps(['start', key]))
|
|
470
518
|
return
|
|
471
519
|
|
|
520
|
+
task._reset_stop()
|
|
472
521
|
if task.run_type == 'THREAD':
|
|
473
|
-
task._handler = threading.Thread(target = self.task_processing, args = (task, task.
|
|
522
|
+
task._handler = threading.Thread(target = self.task_processing, args = (task, task.fn, *task.args), kwargs = task.kwargs, daemon = True)
|
|
474
523
|
elif task.run_type == 'PROCESS':
|
|
475
|
-
task._handler = multiprocessing.get_context('spawn').Process(target = self.task_processing, args = (task, task.
|
|
524
|
+
task._handler = multiprocessing.get_context('spawn').Process(target = self.task_processing, args = (task, task.fn, *task.args), kwargs = task.kwargs, daemon = True)
|
|
476
525
|
task._handler.start()
|
|
477
526
|
|
|
478
527
|
if task.auto_remove:
|
|
@@ -489,7 +538,8 @@ class SchedulerProxy:
|
|
|
489
538
|
if not task and static.scheduler_sync_servers:
|
|
490
539
|
await redis.asyncio.Redis(connection_pool = static.scheduler_sync_servers[1]).publish('CheeseAPI_scheduler', json.dumps(['start', key]))
|
|
491
540
|
else:
|
|
492
|
-
task.
|
|
541
|
+
task._reset_stop()
|
|
542
|
+
task._handler = asyncio.create_task(self.async_task_processing(key, task.fn, *task.args, **task.kwargs))
|
|
493
543
|
|
|
494
544
|
def stop(self, key: str):
|
|
495
545
|
task = self.get_task(key)
|
|
@@ -503,7 +553,7 @@ class SchedulerProxy:
|
|
|
503
553
|
if static.scheduler_sync_servers:
|
|
504
554
|
redis.Redis(connection_pool = static.scheduler_sync_servers[0]).publish('CheeseAPI_scheduler', json.dumps(['stop', key]))
|
|
505
555
|
else:
|
|
506
|
-
local_task.
|
|
556
|
+
local_task._stop()
|
|
507
557
|
if static.scheduler_sync_servers:
|
|
508
558
|
redis.Redis(connection_pool = static.scheduler_sync_servers[0]).hset('CheeseAPI_scheduler_tasks', key, json.dumps(task._to_dict()))
|
|
509
559
|
|
|
@@ -519,7 +569,7 @@ class SchedulerProxy:
|
|
|
519
569
|
if static.scheduler_sync_servers:
|
|
520
570
|
redis.Redis(connection_pool = static.scheduler_sync_servers[0]).publish('CheeseAPI_scheduler', json.dumps(['remove', key]))
|
|
521
571
|
else:
|
|
522
|
-
local_task.
|
|
572
|
+
local_task._stop()
|
|
523
573
|
self._tasks.pop(key, None)
|
|
524
574
|
if static.scheduler_sync_servers:
|
|
525
575
|
redis.Redis(connection_pool = static.scheduler_sync_servers[0]).hdel('CheeseAPI_scheduler_tasks', key)
|
|
@@ -536,7 +586,7 @@ class SchedulerProxy:
|
|
|
536
586
|
if static.scheduler_sync_servers:
|
|
537
587
|
await redis.asyncio.Redis(connection_pool = static.scheduler_sync_servers[1]).publish('CheeseAPI_scheduler', json.dumps(['stop', key]))
|
|
538
588
|
else:
|
|
539
|
-
local_task.
|
|
589
|
+
local_task._stop()
|
|
540
590
|
if static.scheduler_sync_servers:
|
|
541
591
|
await redis.asyncio.Redis(connection_pool = static.scheduler_sync_servers[1]).hset('CheeseAPI_scheduler_tasks', key, json.dumps(task._to_dict()))
|
|
542
592
|
|
|
@@ -550,7 +600,7 @@ class SchedulerProxy:
|
|
|
550
600
|
if static.scheduler_sync_servers:
|
|
551
601
|
await redis.asyncio.Redis(connection_pool = static.scheduler_sync_servers[1]).publish('CheeseAPI_scheduler', json.dumps(['remove', key]))
|
|
552
602
|
else:
|
|
553
|
-
local_task.
|
|
603
|
+
local_task._stop()
|
|
554
604
|
self._tasks.pop(key, None)
|
|
555
605
|
if static.scheduler_sync_servers:
|
|
556
606
|
await redis.asyncio.Redis(connection_pool = static.scheduler_sync_servers[1]).hdel('CheeseAPI_scheduler_tasks', key)
|
|
@@ -1,4 +1,4 @@
|
|
|
1
|
-
import base64, hashlib, asyncio, ssl, struct, json
|
|
1
|
+
import base64, hashlib, asyncio, ssl, struct, json, inspect
|
|
2
2
|
from functools import partial
|
|
3
3
|
from typing import TYPE_CHECKING, AsyncIterable, Self
|
|
4
4
|
|
|
@@ -48,8 +48,8 @@ class Websocket:
|
|
|
48
48
|
self._request: 'Request' = request
|
|
49
49
|
|
|
50
50
|
self._proxy: 'WebsocketProxy' = request._proxy.app.WebsocketProxy_Class(request._proxy.app, self)
|
|
51
|
-
self._key: str = self.request.headers.get('Sec-WebSocket-Key')
|
|
52
|
-
subprotocols = self.request.headers.get('Sec-WebSocket-Protocol')
|
|
51
|
+
self._key: str = self.request.headers.get('sec-websocket-key') or self.request.headers.get('Sec-WebSocket-Key')
|
|
52
|
+
subprotocols = self.request.headers.get('sec-websocket-protocol') or self.request.headers.get('Sec-WebSocket-Protocol')
|
|
53
53
|
self._subprotocols: list[str] | None = subprotocols.strip().split(',') if subprotocols else None
|
|
54
54
|
self._subprotocol: str | None = None
|
|
55
55
|
self.response: Response | None = None
|
|
@@ -111,6 +111,37 @@ class Websocket:
|
|
|
111
111
|
class WebsocketProxy:
|
|
112
112
|
_sync_tasks: dict[str, asyncio.Task] = {}
|
|
113
113
|
|
|
114
|
+
@staticmethod
|
|
115
|
+
def _consume_exception(future: asyncio.Future):
|
|
116
|
+
''' 消费投递结果的异常,避免出现 "exception was never retrieved" '''
|
|
117
|
+
if not future.cancelled():
|
|
118
|
+
future.exception()
|
|
119
|
+
|
|
120
|
+
@staticmethod
|
|
121
|
+
def _schedule(connector: Websocket, coroutine):
|
|
122
|
+
'''
|
|
123
|
+
同步入口(`_static_send` / `_static_close`)不能 await,将协程投递到连接所属的事件循环执行
|
|
124
|
+
|
|
125
|
+
既不能直接丢弃(会报 `coroutine ... was never awaited` 且消息实际发不出去),
|
|
126
|
+
也不能在调用方的循环里执行(连接的 writer 属于创建它的那个循环)
|
|
127
|
+
'''
|
|
128
|
+
loop: asyncio.AbstractEventLoop | None = connector._proxy._loop
|
|
129
|
+
|
|
130
|
+
try:
|
|
131
|
+
current_loop = asyncio.get_running_loop()
|
|
132
|
+
except RuntimeError:
|
|
133
|
+
current_loop = None
|
|
134
|
+
|
|
135
|
+
if loop is None:
|
|
136
|
+
coroutine.close()
|
|
137
|
+
elif current_loop is loop:
|
|
138
|
+
loop.create_task(coroutine)
|
|
139
|
+
elif loop.is_running():
|
|
140
|
+
''' 调用方在其他线程 / 非协程上下文,线程安全地投递回该循环 '''
|
|
141
|
+
asyncio.run_coroutine_threadsafe(coroutine, loop).add_done_callback(WebsocketProxy._consume_exception)
|
|
142
|
+
else:
|
|
143
|
+
coroutine.close()
|
|
144
|
+
|
|
114
145
|
@staticmethod
|
|
115
146
|
def _static_send(path: str, data: bytes | list | str | dict, *, websocket_key_or_keys: str | list[str] | None = None):
|
|
116
147
|
if static.websocket_sync_servers is not None:
|
|
@@ -131,16 +162,16 @@ class WebsocketProxy:
|
|
|
131
162
|
elif path in Websocket.connectors:
|
|
132
163
|
if websocket_key_or_keys is None:
|
|
133
164
|
for connector in Websocket.connectors[path]:
|
|
134
|
-
connector.send(data)
|
|
165
|
+
WebsocketProxy._schedule(connector, connector.send(data))
|
|
135
166
|
elif isinstance(websocket_key_or_keys, str):
|
|
136
167
|
for connector in Websocket.connectors[path]:
|
|
137
168
|
if connector.key == websocket_key_or_keys:
|
|
138
|
-
connector.send(data)
|
|
169
|
+
WebsocketProxy._schedule(connector, connector.send(data))
|
|
139
170
|
break
|
|
140
171
|
elif isinstance(websocket_key_or_keys, list):
|
|
141
172
|
for connector in Websocket.connectors[path]:
|
|
142
173
|
if connector.key in websocket_key_or_keys:
|
|
143
|
-
connector.send(data)
|
|
174
|
+
WebsocketProxy._schedule(connector, connector.send(data))
|
|
144
175
|
|
|
145
176
|
@staticmethod
|
|
146
177
|
async def async_send(path: str, data: bytes | list | str | dict, *, websocket_key_or_keys: str | list[str] | None = None):
|
|
@@ -182,16 +213,16 @@ class WebsocketProxy:
|
|
|
182
213
|
elif path in Websocket.connectors:
|
|
183
214
|
if websocket_key_or_keys is None:
|
|
184
215
|
for connector in Websocket.connectors[path]:
|
|
185
|
-
connector.close(code, message)
|
|
216
|
+
WebsocketProxy._schedule(connector, connector.close(code, message))
|
|
186
217
|
elif isinstance(websocket_key_or_keys, str):
|
|
187
218
|
for connector in Websocket.connectors[path]:
|
|
188
219
|
if connector.key == websocket_key_or_keys:
|
|
189
|
-
connector.close(code, message)
|
|
220
|
+
WebsocketProxy._schedule(connector, connector.close(code, message))
|
|
190
221
|
break
|
|
191
222
|
elif isinstance(websocket_key_or_keys, list):
|
|
192
223
|
for connector in Websocket.connectors[path]:
|
|
193
224
|
if connector.key in websocket_key_or_keys:
|
|
194
|
-
connector.close(code, message)
|
|
225
|
+
WebsocketProxy._schedule(connector, connector.close(code, message))
|
|
195
226
|
|
|
196
227
|
@staticmethod
|
|
197
228
|
async def async_close(path: str, code: int = 1000, message: str = '', *, websocket_key_or_keys: str | list[str] | None = None):
|
|
@@ -218,14 +249,25 @@ class WebsocketProxy:
|
|
|
218
249
|
|
|
219
250
|
self.reader: asyncio.StreamReader | None = None
|
|
220
251
|
self.writer: asyncio.StreamWriter | None = None
|
|
252
|
+
self._loop: asyncio.AbstractEventLoop | None = None
|
|
253
|
+
self._is_disconnected: bool = False
|
|
221
254
|
|
|
222
255
|
async def running(self) -> AsyncIterable[Response]:
|
|
223
256
|
try:
|
|
224
257
|
await self.connect()
|
|
225
258
|
await self.message()
|
|
226
|
-
await self.disconnect()
|
|
227
259
|
except Exception as e:
|
|
228
|
-
|
|
260
|
+
result = self.app.printer.websocket_error(e, self.websocket)
|
|
261
|
+
if inspect.isawaitable(result):
|
|
262
|
+
await result
|
|
263
|
+
finally:
|
|
264
|
+
''' 无论连接、接收如何结束(含异常断连),都必须执行一次清理,否则心跳与推送任务不会停止 '''
|
|
265
|
+
try:
|
|
266
|
+
await self.disconnect()
|
|
267
|
+
except Exception as e:
|
|
268
|
+
result = self.app.printer.websocket_error(e, self.websocket)
|
|
269
|
+
if inspect.isawaitable(result):
|
|
270
|
+
await result
|
|
229
271
|
|
|
230
272
|
async def get_response(self) -> Response:
|
|
231
273
|
headers = {
|
|
@@ -249,6 +291,7 @@ class WebsocketProxy:
|
|
|
249
291
|
WebsocketProxy._sync_tasks[self.websocket.request.path] = asyncio.create_task(self.sync_server_running())
|
|
250
292
|
|
|
251
293
|
loop = asyncio.get_running_loop()
|
|
294
|
+
self._loop = loop
|
|
252
295
|
self.reader = asyncio.StreamReader()
|
|
253
296
|
protocol = asyncio.StreamReaderProtocol(self.reader)
|
|
254
297
|
if isinstance(self.websocket.request._proxy.client_socket, ssl.SSLSocket):
|
|
@@ -344,13 +387,52 @@ class WebsocketProxy:
|
|
|
344
387
|
await self.websocket.on_pong()
|
|
345
388
|
|
|
346
389
|
async def disconnect(self):
|
|
390
|
+
'''
|
|
391
|
+
清理连接
|
|
392
|
+
|
|
393
|
+
- 允许重复调用:接收、发送两侧可能同时发现断连,重复清理直接返回
|
|
394
|
+
- 兼容连接未完全建立的情况:reader / writer 可能为 None,也未加入 connectors
|
|
395
|
+
- 保证执行 on_disconnect():心跳任务与业务推送任务依赖它停止
|
|
396
|
+
- 用 finally 保证底层连接关闭,即使回调清理报错也要关闭
|
|
397
|
+
'''
|
|
398
|
+
if self._is_disconnected:
|
|
399
|
+
return
|
|
400
|
+
self._is_disconnected = True
|
|
401
|
+
|
|
347
402
|
self.websocket._is_running = False
|
|
348
|
-
|
|
349
|
-
|
|
350
|
-
|
|
351
|
-
|
|
352
|
-
|
|
353
|
-
|
|
403
|
+
|
|
404
|
+
connectors = Websocket.connectors.get(self.websocket.request.path)
|
|
405
|
+
if connectors is not None:
|
|
406
|
+
try:
|
|
407
|
+
connectors.remove(self.websocket)
|
|
408
|
+
except ValueError:
|
|
409
|
+
...
|
|
410
|
+
|
|
411
|
+
try:
|
|
412
|
+
self.app.printer.websocket_disconnect(self.websocket)
|
|
413
|
+
except Exception:
|
|
414
|
+
...
|
|
415
|
+
|
|
416
|
+
try:
|
|
417
|
+
await self.websocket.on_disconnect()
|
|
418
|
+
except Exception as e:
|
|
419
|
+
result = self.app.printer.websocket_error(e, self.websocket)
|
|
420
|
+
if inspect.isawaitable(result):
|
|
421
|
+
await result
|
|
422
|
+
finally:
|
|
423
|
+
writer, self.writer = self.writer, None
|
|
424
|
+
self.reader = None
|
|
425
|
+
|
|
426
|
+
if writer is not None:
|
|
427
|
+
try:
|
|
428
|
+
writer.close()
|
|
429
|
+
except Exception:
|
|
430
|
+
...
|
|
431
|
+
try:
|
|
432
|
+
await writer.wait_closed()
|
|
433
|
+
except Exception:
|
|
434
|
+
''' 对端已断开时 wait_closed 可能再报 OSError,属于预期错误,不能中断清理 '''
|
|
435
|
+
...
|
|
354
436
|
|
|
355
437
|
def encode(self, opcode: int, data: bytes) -> bytes:
|
|
356
438
|
_bytes = bytearray()
|
|
@@ -404,23 +486,51 @@ class WebsocketProxy:
|
|
|
404
486
|
except asyncio.TimeoutError:
|
|
405
487
|
return opcode, full_payload
|
|
406
488
|
|
|
489
|
+
def _check_running(self):
|
|
490
|
+
''' 发送前的连接校验,连接已断开或未建立时直接拒绝,避免继续向坏连接写入 '''
|
|
491
|
+
if self._is_disconnected or self.writer is None or not self.websocket._is_running:
|
|
492
|
+
raise ConnectionError('The websocket is disconnected, data will not be sent')
|
|
493
|
+
|
|
494
|
+
def _disconnected(self, e: Exception) -> ConnectionError:
|
|
495
|
+
''' 标记连接断开并返回统一的异常,阻止后续继续向坏连接发送 '''
|
|
496
|
+
self.websocket._is_running = False
|
|
497
|
+
self._is_disconnected = True
|
|
498
|
+
return ConnectionError('The websocket is disconnected, data will not be sent')
|
|
499
|
+
|
|
407
500
|
async def _instance_send(self, data: bytes | list | str | dict, **kwargs):
|
|
408
|
-
|
|
409
|
-
|
|
410
|
-
|
|
411
|
-
|
|
412
|
-
|
|
413
|
-
|
|
414
|
-
|
|
501
|
+
self._check_running()
|
|
502
|
+
try:
|
|
503
|
+
if isinstance(data, str):
|
|
504
|
+
self.writer.write(self.encode(0x1, data.encode()))
|
|
505
|
+
elif isinstance(data, bytes):
|
|
506
|
+
self.writer.write(self.encode(0x2, data))
|
|
507
|
+
else:
|
|
508
|
+
self.writer.write(self.encode(0x1, json.dumps(data).encode()))
|
|
509
|
+
await self.writer.drain()
|
|
510
|
+
except (ConnectionError, OSError, RuntimeError) as e:
|
|
511
|
+
raise self._disconnected(e) from e
|
|
415
512
|
|
|
416
513
|
async def _instance_close(self, code: int = 1000, message: str = ''):
|
|
417
|
-
self.
|
|
418
|
-
|
|
514
|
+
if self._is_disconnected or self.writer is None:
|
|
515
|
+
return
|
|
516
|
+
try:
|
|
517
|
+
self.writer.write(self.encode(0x8, struct.pack('!H', code) + message.encode('utf-8')))
|
|
518
|
+
await self.writer.drain()
|
|
519
|
+
except (ConnectionError, OSError, RuntimeError) as e:
|
|
520
|
+
raise self._disconnected(e) from e
|
|
419
521
|
|
|
420
522
|
async def ping(self):
|
|
421
|
-
self.
|
|
422
|
-
|
|
523
|
+
self._check_running()
|
|
524
|
+
try:
|
|
525
|
+
self.writer.write(self.encode(0x9, b''))
|
|
526
|
+
await self.writer.drain()
|
|
527
|
+
except (ConnectionError, OSError, RuntimeError) as e:
|
|
528
|
+
raise self._disconnected(e) from e
|
|
423
529
|
|
|
424
530
|
async def pong(self):
|
|
425
|
-
self.
|
|
426
|
-
|
|
531
|
+
self._check_running()
|
|
532
|
+
try:
|
|
533
|
+
self.writer.write(self.encode(0xA, b''))
|
|
534
|
+
await self.writer.drain()
|
|
535
|
+
except (ConnectionError, OSError, RuntimeError) as e:
|
|
536
|
+
raise self._disconnected(e) from e
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|