/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
- #!/usr/bin/env python3
- # -*- coding: utf-8 -*-
- # Copyright (c) Facebook, Inc. and its affiliates.
- #
- # This source code is licensed under the MIT license found in the
- # LICENSE file in the root directory of this source tree.
- from fbnet.command_runner_asyncio.CommandRunner.Command import Processor
- from thrift.server.TAsyncioServer import ThriftServerProtocolFactory
- from thrift.server.TServer import TServerEventHandler
- from thrift.Thrift import TProcessorEventHandler
- from .base_service import ServiceTask
- from .command_handler import CommandHandler
- from .options import Option
- class CommandServer(ServiceTask):
- """
- Command server for thrift commands.
- """
- PORT = Option(
- "-p", "--port", help="TCP port for FCR service", type=int, default=5000
- )
- def __init__(self, service, loop=None):
- super().__init__(service, "CommandServer")
- self._handler = None
- self._server = None
- self._backlog = 100
- def _get_processor_class(self):
- return Processor
- async def run(self):
- # Wait for FBNet to finish its run
- await self.service.device_db.wait_for_data()
- event_handler = self._create_thrift_event_handler()
- thrift_handler = self._create_thrift_handler(event_handler)
- self._handler = thrift_handler
- processor = self._get_processor_class()(self._handler, loop=self.loop)
- processor.setEventHandler(event_handler)
- pfactory = ThriftServerProtocolFactory(
- processor, self._create_server_event_handler(), loop=self.loop
- )
- self._server = await self.loop.create_server(
- pfactory, port=self.PORT, backlog=self._backlog
- )
- self.logger.info("server started: %d ", self.PORT)
- # Wait for the server to be closed. Typically this will never close.
- await self._server.wait_closed()
- self.logger.info("server done: %d", self.PORT)
- self._server = None
- def _create_thrift_handler(self, event_handler):
- return CommandHandler(self.service)
- def _create_server_event_handler(self):
- return TServerEventHandler()
- def _create_thrift_event_handler(self):
- return TProcessorEventHandler()
- async def cleanup(self):
- # Let the handler cleanup any relevent state
- self._handler.cleanup()
- self.close()
- def close(self):
- # close the server port
- if self._server:
- self.logger.info("closing the server")
- self._server.close()