1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
|
import os
import signal
import asyncio
import time
from typing import List
from typing import Set
from typing import Callable
from typing import Optional
import aiohttp
from .atx import Atx
from .streamer import Streamer
from .ps2 import Ps2Keyboard
from .logging import get_logger
# =====
def _system_task(method: Callable) -> Callable:
async def wrap(self: "Server") -> None:
try:
await method(self)
except asyncio.CancelledError:
pass
except Exception:
get_logger().exception("Unhandled exception, killing myself ...")
os.kill(os.getpid(), signal.SIGTERM)
return wrap
class Server: # pylint: disable=too-many-instance-attributes
def __init__(
self,
atx: Atx,
streamer: Streamer,
keyboard: Ps2Keyboard,
heartbeat: float,
atx_leds_poll: float,
video_shutdown_delay: float,
loop: asyncio.AbstractEventLoop,
) -> None:
self.__atx = atx
self.__streamer = streamer
self.__heartbeat = heartbeat
self.__keyboard = keyboard
self.__video_shutdown_delay = video_shutdown_delay
self.__atx_leds_poll = atx_leds_poll
self.__loop = loop
self.__sockets: Set[aiohttp.web.WebSocketResponse] = set()
self.__sockets_lock = asyncio.Lock()
self.__system_tasks: List[asyncio.Task] = []
def run(self, host: str, port: int) -> None:
self.__keyboard.start()
app = aiohttp.web.Application(loop=self.__loop)
app.router.add_get("/", self.__root_handler)
app.router.add_get("/ws", self.__ws_handler)
app.on_shutdown.append(self.__on_shutdown)
app.on_cleanup.append(self.__on_cleanup)
self.__system_tasks.extend([
self.__loop.create_task(self.__monitor_keyboard()),
self.__loop.create_task(self.__stream_controller()),
self.__loop.create_task(self.__poll_dead_sockets()),
self.__loop.create_task(self.__poll_atx_leds()),
])
aiohttp.web.run_app(
app=app,
host=host,
port=port,
print=(lambda text: [get_logger().info(line.strip()) for line in text.strip().splitlines()]), # type: ignore
)
async def __root_handler(self, _: aiohttp.web.Request) -> aiohttp.web.Response:
return aiohttp.web.Response(text="OK")
async def __ws_handler(self, request: aiohttp.web.Request) -> aiohttp.web.WebSocketResponse:
ws = aiohttp.web.WebSocketResponse(heartbeat=self.__heartbeat)
await ws.prepare(request)
await self.__register_socket(ws)
async for msg in ws:
if msg.type == aiohttp.web.WSMsgType.TEXT:
retval = await self.__execute_command(msg.data)
if retval:
await ws.send_str(retval)
else:
break
return ws
async def __on_shutdown(self, _: aiohttp.web.Application) -> None:
logger = get_logger(0)
logger.info("Cancelling system tasks ...")
for task in self.__system_tasks:
task.cancel()
await asyncio.gather(*self.__system_tasks)
logger.info("Disconnecting clients ...")
for ws in list(self.__sockets):
await self.__remove_socket(ws)
async def __on_cleanup(self, _: aiohttp.web.Application) -> None:
if self.__keyboard.is_alive():
self.__keyboard.stop()
if self.__streamer.is_running():
await self.__streamer.stop()
@_system_task
async def __monitor_keyboard(self) -> None:
while self.__keyboard.is_alive():
await asyncio.sleep(0.1)
raise RuntimeError("Keyboard dead")
@_system_task
async def __stream_controller(self) -> None:
prev = 0
shutdown_at = 0.0
while True:
cur = len(self.__sockets)
if prev == 0 and cur > 0:
if not self.__streamer.is_running():
await self.__streamer.start()
elif prev > 0 and cur == 0:
shutdown_at = time.time() + self.__video_shutdown_delay
elif prev == 0 and cur == 0 and time.time() > shutdown_at:
if self.__streamer.is_running():
await self.__streamer.stop()
prev = cur
await asyncio.sleep(0.1)
@_system_task
async def __poll_dead_sockets(self) -> None:
while True:
for ws in list(self.__sockets):
if ws.closed or not ws._req.transport: # pylint: disable=protected-access
await self.__remove_socket(ws)
await asyncio.sleep(0.1)
@_system_task
async def __poll_atx_leds(self) -> None:
while True:
if self.__sockets:
await self.__broadcast("EVENT atx_leds %d %d" % (self.__atx.get_leds()))
await asyncio.sleep(self.__atx_leds_poll)
async def __broadcast(self, msg: str) -> None:
await asyncio.gather(*[
ws.send_str(msg)
for ws in list(self.__sockets)
if not ws.closed and ws._req.transport # pylint: disable=protected-access
], return_exceptions=True)
async def __execute_command(self, command: str) -> Optional[str]:
(command, args) = (command.strip().split(" ", maxsplit=1) + [""])[:2]
if command == "CLICK":
method = {
"power": self.__atx.click_power,
"power_long": self.__atx.click_power_long,
"reset": self.__atx.click_reset,
}.get(args)
if method:
await method()
return None
get_logger().warning("Received an incorrect command: %r", command)
return "ERROR incorrect command"
async def __register_socket(self, ws: aiohttp.web.WebSocketResponse) -> None:
async with self.__sockets_lock:
self.__sockets.add(ws)
get_logger().info("Registered new client socket: remote=%s; id=%d; active=%d",
ws._req.remote, id(ws), len(self.__sockets)) # pylint: disable=protected-access
async def __remove_socket(self, ws: aiohttp.web.WebSocketResponse) -> None:
async with self.__sockets_lock:
try:
self.__sockets.remove(ws)
get_logger().info("Removed client socket: remote=%s; id=%d; active=%d",
ws._req.remote, id(ws), len(self.__sockets)) # pylint: disable=protected-access
await ws.close()
except Exception:
pass
|