/fbnet/command_runner/command_server.py

https://github.com/facebookincubator/FCR · Python · 85 lines · 48 code · 24 blank · 13 comment · 1 complexity · 872f4b626bcd4b29d0469883ca5e51b5 MD5 · raw file

  1. #!/usr/bin/env python3
  2. # -*- coding: utf-8 -*-
  3. # Copyright (c) Facebook, Inc. and its affiliates.
  4. #
  5. # This source code is licensed under the MIT license found in the
  6. # LICENSE file in the root directory of this source tree.
  7. from fbnet.command_runner_asyncio.CommandRunner.Command import Processor
  8. from thrift.server.TAsyncioServer import ThriftServerProtocolFactory
  9. from thrift.server.TServer import TServerEventHandler
  10. from thrift.Thrift import TProcessorEventHandler
  11. from .base_service import ServiceTask
  12. from .command_handler import CommandHandler
  13. from .options import Option
  14. class CommandServer(ServiceTask):
  15. """
  16. Command server for thrift commands.
  17. """
  18. PORT = Option(
  19. "-p", "--port", help="TCP port for FCR service", type=int, default=5000
  20. )
  21. def __init__(self, service, loop=None):
  22. super().__init__(service, "CommandServer")
  23. self._handler = None
  24. self._server = None
  25. self._backlog = 100
  26. def _get_processor_class(self):
  27. return Processor
  28. async def run(self):
  29. # Wait for FBNet to finish its run
  30. await self.service.device_db.wait_for_data()
  31. event_handler = self._create_thrift_event_handler()
  32. thrift_handler = self._create_thrift_handler(event_handler)
  33. self._handler = thrift_handler
  34. processor = self._get_processor_class()(self._handler, loop=self.loop)
  35. processor.setEventHandler(event_handler)
  36. pfactory = ThriftServerProtocolFactory(
  37. processor, self._create_server_event_handler(), loop=self.loop
  38. )
  39. self._server = await self.loop.create_server(
  40. pfactory, port=self.PORT, backlog=self._backlog
  41. )
  42. self.logger.info("server started: %d ", self.PORT)
  43. # Wait for the server to be closed. Typically this will never close.
  44. await self._server.wait_closed()
  45. self.logger.info("server done: %d", self.PORT)
  46. self._server = None
  47. def _create_thrift_handler(self, event_handler):
  48. return CommandHandler(self.service)
  49. def _create_server_event_handler(self):
  50. return TServerEventHandler()
  51. def _create_thrift_event_handler(self):
  52. return TProcessorEventHandler()
  53. async def cleanup(self):
  54. # Let the handler cleanup any relevent state
  55. self._handler.cleanup()
  56. self.close()
  57. def close(self):
  58. # close the server port
  59. if self._server:
  60. self.logger.info("closing the server")
  61. self._server.close()