xtcp.py 1.9 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859
  1. import asyncio
  2. import struct
  3. class CustomProtocol(asyncio.Protocol):
  4. clients = {}
  5. def connection_made(self, client):
  6. self.client = client
  7. self.client_info = client.get_extra_info('peername')
  8. print(f"Connection from {self.client_info}", flush=True)
  9. CustomProtocol.clients[self.client_info] = client
  10. self.on_connect()
  11. def data_received(self, data):
  12. asyncio.create_task(self.handle_data(data))
  13. def connection_lost(self, exc):
  14. print(f"Connection closed by {self.client_info}", flush=True)
  15. if self.client_info in CustomProtocol.clients:
  16. del CustomProtocol.clients[self.client_info]
  17. self.on_disconnect()
  18. def on_connect(self):
  19. print(f"Connected to {self.client_info}", flush=True)
  20. async def handle_data(self, data):
  21. # Simulating a blocking operation with asyncio.run_in_executor
  22. result = await asyncio.get_running_loop().run_in_executor(None, self.process_data, data)
  23. self.client.write(result)
  24. def process_data(self, data):
  25. # Simulate a CPU-bound task
  26. values = struct.unpack('!hh', data)
  27. processed_values = (values[0] + 1, values[1] + 1)
  28. return struct.pack('!hh', *processed_values)
  29. def on_disconnect(self):
  30. print(f"Disconnected from {self.client_info}", flush=True)
  31. @classmethod
  32. async def broadcast(cls, message):
  33. data = struct.pack('!hh', *message)
  34. for client in cls.clients.values():
  35. client.write(data)
  36. await asyncio.gather(*(client.drain() for client in cls.clients.values()))
  37. async def main():
  38. loop = asyncio.get_running_loop()
  39. server = await loop.create_server(
  40. lambda: CustomProtocol(),
  41. '0.0.0.0', 20917
  42. )
  43. async with server:
  44. print("Server listening on 0.0.0.0:20917", flush=True)
  45. await server.serve_forever()
  46. if __name__ == "__main__":
  47. asyncio.run(main())