rewrite: Implement shard IPC server.
This commit is contained in:
0
bot/meta/ipc/__init__.py
Normal file
0
bot/meta/ipc/__init__.py
Normal file
200
bot/meta/ipc/client.py
Normal file
200
bot/meta/ipc/client.py
Normal file
@@ -0,0 +1,200 @@
|
||||
from typing import Optional
|
||||
import asyncio
|
||||
import logging
|
||||
import pickle
|
||||
|
||||
|
||||
logging.basicConfig(level=logging.DEBUG)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AppClient:
|
||||
routes = {} # route_name -> Callable[Any, Awaitable[Any]]
|
||||
|
||||
def __init__(self, appid, client_address, server_address):
|
||||
self.appid = appid
|
||||
self.address = client_address
|
||||
self.server_address = server_address
|
||||
|
||||
self.peers = {appid: client_address} # appid -> address
|
||||
|
||||
self._listener: Optional[asyncio.Server] = None # Local client server
|
||||
self._server = None # Connection to the registry server
|
||||
|
||||
self.register_route('new_peer')(self.new_peer)
|
||||
self.register_route('drop_peer')(self.drop_peer)
|
||||
self.register_route('peer_list')(self.peer_list)
|
||||
|
||||
def register_route(self, name=None):
|
||||
def wrapper(coro):
|
||||
route = AppRoute(coro, name)
|
||||
self.routes[route.name] = route
|
||||
return route
|
||||
return wrapper
|
||||
|
||||
async def server_connection(self):
|
||||
"""Establish a connection to the registry server"""
|
||||
try:
|
||||
reader, writer = await asyncio.open_connection(**self.server_address)
|
||||
|
||||
payload = ('connect', (), {'appid': self.appid, 'address': self.address})
|
||||
writer.write(pickle.dumps(payload))
|
||||
writer.write(b'\n')
|
||||
await writer.drain()
|
||||
|
||||
data = await reader.readline()
|
||||
peers = pickle.loads(data)
|
||||
self.peers = peers
|
||||
self._server = (reader, writer)
|
||||
except Exception:
|
||||
logger.exception("Could not connect to registry server. Trying again in 30 seconds.")
|
||||
await asyncio.sleep(30)
|
||||
asyncio.create_task(self.server_connection())
|
||||
else:
|
||||
logger.info("Connected to the registry server, launching keepalive.")
|
||||
asyncio.create_task(self._server_keepalive())
|
||||
|
||||
async def _server_keepalive(self):
|
||||
if self._server is None:
|
||||
raise ValueError("Cannot keepalive non-existent server!")
|
||||
reader, write = self._server
|
||||
try:
|
||||
await reader.read()
|
||||
except Exception:
|
||||
logger.exception("Lost connection to address server. Reconnecting...")
|
||||
else:
|
||||
# Connection ended or broke
|
||||
logger.info("Lost connection to address server. Reconnecting...")
|
||||
await asyncio.sleep(30)
|
||||
asyncio.create_task(self.server_connection())
|
||||
|
||||
async def new_peer(self, appid, address):
|
||||
self.peers[appid] = address
|
||||
|
||||
async def peer_list(self, peers):
|
||||
self.peers = peers
|
||||
|
||||
async def drop_peer(self, appid):
|
||||
self.peers.pop(appid, None)
|
||||
|
||||
async def close(self):
|
||||
# Close connection to the server
|
||||
# TODO
|
||||
...
|
||||
|
||||
async def request(self, appid, payload: 'AppPayload'):
|
||||
try:
|
||||
if appid not in self.peers:
|
||||
raise ValueError(f"Peer '{appid}' not found.")
|
||||
logger.debug(f"Sending request to app '{appid}' with payload {payload}")
|
||||
|
||||
address = self.peers[appid]
|
||||
reader, writer = await asyncio.open_connection(**address)
|
||||
|
||||
writer.write(payload.encoded())
|
||||
await writer.drain()
|
||||
writer.write_eof()
|
||||
result = await reader.read()
|
||||
writer.close()
|
||||
decoded = payload.route.decode(result)
|
||||
return decoded
|
||||
except Exception:
|
||||
logging.exception(f"Failed to send request to {appid}'")
|
||||
return None
|
||||
|
||||
async def requestall(self, payload):
|
||||
results = await asyncio.gather(*(self.request(appid, payload) for appid in self.peers))
|
||||
return dict(zip(self.peers.keys(), results))
|
||||
|
||||
async def handle_request(self, reader, writer):
|
||||
data = await reader.read()
|
||||
loaded = pickle.loads(data)
|
||||
route, args, kwargs = loaded
|
||||
|
||||
logger.debug(f"AppClient {self.appid} handling request on route '{route}' with args {args} and kwargs {kwargs}")
|
||||
|
||||
if route in self.routes:
|
||||
try:
|
||||
await self.routes[route].run((reader, writer), args, kwargs)
|
||||
except Exception:
|
||||
logger.exception(f"Fatal exception during route '{route}'. This should never happen!")
|
||||
else:
|
||||
logger.warning(f"Appclient '{self.appid}' recieved unknown route {route}. Ignoring.")
|
||||
writer.write_eof()
|
||||
|
||||
async def connect(self):
|
||||
"""
|
||||
Start the local peer server.
|
||||
Connect to the address server.
|
||||
"""
|
||||
# Start the client server
|
||||
self._listener = await asyncio.start_server(self.handle_request, **self.address, start_serving=True)
|
||||
|
||||
logger.info(f"Serving on {self.address}")
|
||||
await self.server_connection()
|
||||
|
||||
|
||||
class AppPayload:
|
||||
__slots__ = ('route', 'args', 'kwargs')
|
||||
|
||||
def __init__(self, route, *args, **kwargs):
|
||||
self.route = route
|
||||
self.args = args
|
||||
self.kwargs = kwargs
|
||||
|
||||
def __await__(self):
|
||||
return self.route.execute(*self.args, **self.kwargs).__await__()
|
||||
|
||||
def encoded(self):
|
||||
return pickle.dumps((self.route.name, self.args, self.kwargs))
|
||||
|
||||
|
||||
class AppRoute:
|
||||
__slots__ = ('func', 'name')
|
||||
|
||||
def __init__(self, func, name=None):
|
||||
self.func = func
|
||||
self.name = name or func.__name__
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
return AppPayload(self, *args, **kwargs)
|
||||
|
||||
def encode(self, output):
|
||||
return pickle.dumps(output)
|
||||
|
||||
def decode(self, encoded):
|
||||
# TODO: Handle exceptions here somehow
|
||||
if len(encoded) > 0:
|
||||
return pickle.loads(encoded)
|
||||
else:
|
||||
return ''
|
||||
|
||||
def encoder(self, func):
|
||||
self.encode = func
|
||||
|
||||
def decoder(self, func):
|
||||
self.decode = func
|
||||
|
||||
async def execute(self, *args, **kwargs):
|
||||
"""
|
||||
Execute the underlying function, with the given arguments.
|
||||
"""
|
||||
return await self.func(*args, **kwargs)
|
||||
|
||||
async def run(self, connection, args, kwargs):
|
||||
"""
|
||||
Run the route, with the given arguments, using the given connection.
|
||||
"""
|
||||
# TODO: ContextVar here for logging? Or in handle_request?
|
||||
# Get encoded result
|
||||
# TODO: handle exceptions in the execution process
|
||||
try:
|
||||
result = await self.execute(*args, **kwargs)
|
||||
payload = self.encode(result)
|
||||
except Exception:
|
||||
logger.exception(f"Exception occured running route '{self.name}' with args: {args} and kwargs: {kwargs}")
|
||||
payload = b''
|
||||
_, writer = connection
|
||||
writer.write(payload)
|
||||
await writer.drain()
|
||||
writer.close()
|
||||
173
bot/meta/ipc/server.py
Normal file
173
bot/meta/ipc/server.py
Normal file
@@ -0,0 +1,173 @@
|
||||
import asyncio
|
||||
import pickle
|
||||
import logging
|
||||
import string
|
||||
import random
|
||||
|
||||
from ..logger import log_action, log_context, log_app
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
uuid_alphabet = string.ascii_lowercase + string.digits
|
||||
|
||||
|
||||
def short_uuid():
|
||||
return ''.join(random.choices(uuid_alphabet, k=10))
|
||||
|
||||
|
||||
class AppServer:
|
||||
routes = {} # route name -> bound method
|
||||
|
||||
def __init__(self):
|
||||
self.clients = {} # AppID -> (info, connection)
|
||||
|
||||
self.route('ping')(self.route_ping)
|
||||
self.route('whereis')(self.route_whereis)
|
||||
self.route('peers')(self.route_peers)
|
||||
self.route('connect')(self.client_connection)
|
||||
|
||||
@classmethod
|
||||
def route(cls, route_name):
|
||||
"""
|
||||
Decorator to add a route to the server.
|
||||
"""
|
||||
def wrapper(coro):
|
||||
cls.routes[route_name] = coro
|
||||
return coro
|
||||
return wrapper
|
||||
|
||||
async def route_ping(self, connection):
|
||||
"""
|
||||
Pong.
|
||||
"""
|
||||
reader, writer = connection
|
||||
writer.write(b"Pong")
|
||||
writer.write_eof()
|
||||
|
||||
async def route_whereis(self, connection, appid):
|
||||
"""
|
||||
Return an address for the given client appid.
|
||||
Returns None if the client does not have a connection.
|
||||
"""
|
||||
reader, writer = connection
|
||||
if appid in self.clients:
|
||||
writer.write(pickle.dumps(self.clients[appid][0]))
|
||||
else:
|
||||
writer.write(b'')
|
||||
writer.write_eof()
|
||||
|
||||
async def route_peers(self, connection):
|
||||
"""
|
||||
Send back a map of current peers.
|
||||
"""
|
||||
reader, writer = connection
|
||||
peers = self.peer_list()
|
||||
payload = pickle.dumps(('peer_list', (peers,)))
|
||||
writer.write(payload)
|
||||
writer.write_eof()
|
||||
|
||||
async def client_connection(self, connection, appid, address):
|
||||
"""
|
||||
Register and hold a new client connection.
|
||||
"""
|
||||
log_action.set("CONN " + appid)
|
||||
reader, writer = connection
|
||||
# Add the new client
|
||||
self.clients[appid] = (address, connection)
|
||||
|
||||
# Send the new client a client list
|
||||
peers = self.peer_list()
|
||||
writer.write(pickle.dumps(peers))
|
||||
writer.write(b'\n')
|
||||
await writer.drain()
|
||||
|
||||
# Announce the new client to everyone
|
||||
await self.broadcast('new_peer', (), {'appid': appid, 'address': address})
|
||||
|
||||
# Keep the connection open until socket closed or EOF (indicating client death)
|
||||
try:
|
||||
await reader.read()
|
||||
finally:
|
||||
# Connection ended or it broke
|
||||
logger.info(f"Lost client '{appid}'")
|
||||
await self.deregister_client(appid)
|
||||
|
||||
async def handle_connection(self, reader, writer):
|
||||
data = await reader.readline()
|
||||
route, args, kwargs = pickle.loads(data)
|
||||
|
||||
rqid = short_uuid()
|
||||
log_context.set("RQID:" + rqid)
|
||||
log_action.set("SERV ROUTE " + route)
|
||||
|
||||
logger.info(f"AppServer handling request on route '{route}' with args {args} and kwargs {kwargs}")
|
||||
|
||||
if route in self.routes:
|
||||
# Execute route
|
||||
try:
|
||||
await self.routes[route]((reader, writer), *args, **kwargs)
|
||||
except Exception:
|
||||
logger.exception(f"AppServer recieved exception during route '{route}'")
|
||||
else:
|
||||
logger.warning(f"AppServer recieved unknown route '{route}'. Ignoring.")
|
||||
|
||||
def peer_list(self):
|
||||
return {appid: address for appid, (address, _) in self.clients.items()}
|
||||
|
||||
async def deregister_client(self, appid):
|
||||
self.clients.pop(appid, None)
|
||||
await self.broadcast('drop_peer', (), {'appid': appid})
|
||||
|
||||
async def broadcast(self, route, args, kwargs):
|
||||
logger.debug(f"Sending broadcast on route '{route}' with args {args} and kwargs {kwargs}.")
|
||||
payload = pickle.dumps((route, args, kwargs))
|
||||
if self.clients:
|
||||
await asyncio.gather(
|
||||
*(self._send(appid, payload) for appid in self.clients),
|
||||
return_exceptions=True
|
||||
)
|
||||
|
||||
async def message_client(self, appid, route, args, kwargs):
|
||||
"""
|
||||
Send a message to client `appid` along `route` with given arguments.
|
||||
"""
|
||||
logger.debug(f"Sending '{route}' to '{appid}' with args {args} and kwargs {kwargs}.")
|
||||
if appid not in self.clients:
|
||||
raise ValueError(f"Client '{appid}' is not connected.")
|
||||
|
||||
payload = pickle.dumps((route, args, kwargs))
|
||||
return await self._send(appid, payload)
|
||||
|
||||
async def _send(self, appid, payload):
|
||||
"""
|
||||
Send the encoded `payload` to the client `appid`.
|
||||
"""
|
||||
address, _ = self.clients[appid]
|
||||
try:
|
||||
reader, writer = await asyncio.open_connection(**address)
|
||||
writer.write(payload)
|
||||
writer.write_eof()
|
||||
await writer.drain()
|
||||
writer.close()
|
||||
except Exception as ex:
|
||||
# TODO: Close client if we can't connect?
|
||||
logger.exception(f"Failed to send message to '{appid}'")
|
||||
raise ex
|
||||
|
||||
async def start(self, address):
|
||||
log_app.set("APPSERVER")
|
||||
server = await asyncio.start_server(self.handle_connection, **address)
|
||||
logger.info(f"Serving on {address}")
|
||||
async with server:
|
||||
await server.serve_forever()
|
||||
|
||||
|
||||
async def start_server():
|
||||
address = {'host': '127.0.0.1', 'port': '5000'}
|
||||
server = AppServer()
|
||||
await server.start(address)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
asyncio.run(start_server())
|
||||
Reference in New Issue
Block a user