| #!/usr/bin/python |
| # Copyright (c) 2012 The Chromium OS Authors. All rights reserved. |
| # Use of this source code is governed by a BSD-style license that can be |
| # found in the LICENSE file. |
| |
| import errno |
| import heapq |
| import logging |
| import os |
| import sys |
| import socket |
| import threading |
| import xmlrpclib |
| |
| import rpm_logging_config |
| from config import rpm_config |
| from MultiThreadedXMLRPCServer import MultiThreadedXMLRPCServer |
| from rpm_infrastructure_exception import RPMInfrastructureException |
| |
| import common |
| from autotest_lib.server import frontend |
| from autotest_lib.site_utils.rpm_control_system import utils |
| |
| DEFAULT_RPM_COUNT = 0 |
| TERMINATED = -1 |
| |
| # Indexes for accessing heap entries. |
| RPM_COUNT = 0 |
| DISPATCHER_URI = 1 |
| |
| LOG_FILENAME_FORMAT = rpm_config.get('GENERAL','frontend_logname_format') |
| DEFAULT_RPM_ID = rpm_config.get('RPM_INFRASTRUCTURE', 'default_rpm_id') |
| |
| # Valid state values. |
| VALID_STATE_VALUES = ['ON', 'OFF', 'CYCLE'] |
| |
| # Servo-interface mapping file |
| MAPPING_FILE = os.path.join( |
| os.path.dirname(__file__), |
| rpm_config.get('CiscoPOE', 'servo_interface_mapping_file')) |
| |
| # Size of the LRU that holds power management unit information related |
| # to a device, e.g. rpm_hostname, outlet, hydra_hostname, etc. |
| LRU_SIZE = rpm_config.getint('RPM_INFRASTRUCTURE', 'lru_size') |
| |
| |
| class DispatcherDownException(Exception): |
| """Raised when a particular RPMDispatcher is down.""" |
| |
| |
| class RPMFrontendServer(object): |
| """ |
| This class is the frontend server of the RPM Infrastructure. All clients |
| will send their power state requests to this central server who will |
| forward the requests to an avaliable or already assigned RPM dispatcher |
| server. |
| |
| Once the dispatcher processes the request it will return the result |
| to this frontend server who will send the result back to the client. |
| |
| All calls to this server are blocking. |
| |
| @var _dispatcher_minheap: Min heap that returns a list of format- |
| [ num_rpm's, dispatcher_uri ] |
| Used to choose the least loaded dispatcher. |
| @var _entry_dict: Maps dispatcher URI to an entry (list) inside the min |
| heap. If a dispatcher server shuts down this allows us to |
| invalidate the entry in the minheap. |
| @var _lock: Used to protect data from multiple running threads all |
| manipulating the same data. |
| @var _rpm_dict: Maps rpm hostname's to an already assigned dispatcher |
| server. |
| @var _mapping_last_modified: Last-modified time of the servo-interface |
| mapping file. |
| @var _servo_interface: Maps servo hostname to (switch_hostname, interface). |
| @var _rpm_info: An LRU cache to hold recently visited rpm information |
| so that we don't hit AFE too often. The elements in |
| the cache are instances of PowerUnitInfo indexed by |
| dut hostnames. POE info is not stored in the cache. |
| @var _afe: AFE instance to talk to autotest. Used to retrieve rpm hostname. |
| @var _email_handler: Email handler to use to control email notifications. |
| """ |
| |
| |
| def __init__(self, email_handler=None): |
| """ |
| RPMFrontendServer constructor. |
| |
| Initializes instance variables. |
| """ |
| self._dispatcher_minheap = [] |
| self._entry_dict = {} |
| self._lock = threading.Lock() |
| self._mapping_last_modified = os.path.getmtime(MAPPING_FILE) |
| self._servo_interface = utils.load_servo_interface_mapping() |
| self._rpm_dict = {} |
| self._afe = frontend.AFE() |
| self._rpm_info = utils.LRUCache(size=LRU_SIZE) |
| self._email_handler = email_handler |
| |
| |
| def set_power_via_poe(self, device_hostname, new_state): |
| """Sets power state of the device to the requested state via POE. |
| |
| @param device_hostname: Hostname of the servo to control. |
| @param new_state: [ON, OFF, CYCLE] State to which we want to set the |
| device's outlet to. |
| |
| @return: True if the attempt to change power state was successful, |
| False otherwise. |
| |
| @raise RPMInfrastructureException: No dispatchers are available or can |
| be reached. |
| """ |
| # Remove any DNS Zone information and simplify down to just the hostname. |
| device_hostname = device_hostname.split('.')[0] |
| new_state = new_state.upper() |
| if new_state not in VALID_STATE_VALUES: |
| logging.error('Received request to set servo %s to invalid ' |
| 'state %s', device_hostname, new_state) |
| return False |
| logging.info('Received request to set servo: %s to state: %s', |
| device_hostname, new_state) |
| powerunit_info = self._get_poe_powerunit_info(device_hostname) |
| try: |
| return self._queue_once(powerunit_info, new_state) |
| except DispatcherDownException: |
| # Retry forwarding the request. |
| return self.set_power_via_poe(device_hostname, new_state) |
| |
| |
| def set_power_via_rpm(self, device_hostname, rpm_hostname, |
| rpm_outlet, hydra_hostname, new_state): |
| """Sets power state of a device to the requested state via RPM. |
| |
| Unlike the special case of POE, powerunit information is not available |
| on the RPM server, so must be provided as arguments. |
| |
| @param device_hostname: Hostname of the servo to control. |
| @param rpm_hostname: Hostname of the RPM to use. |
| @param rpm_outlet: The RPM outlet to control. |
| @param hydra_hostname: If required, the hydra device to SSH through to |
| get to the RPM. |
| @param new_state: [ON, OFF, CYCLE] State to which we want to set the |
| device's outlet to. |
| |
| @return: True if the attempt to change power state was successful, |
| False otherwise. |
| |
| @raise RPMInfrastructureException: No dispatchers are available or can |
| be reached. |
| """ |
| powerunit_info = utils.PowerUnitInfo( |
| device_hostname=device_hostname, |
| powerunit_type=utils.PowerUnitInfo.POWERUNIT_TYPES.RPM, |
| powerunit_hostname=rpm_hostname, |
| outlet=rpm_outlet, |
| hydra_hostname=hydra_hostname, |
| ) |
| try: |
| return self._queue_once(powerunit_info, new_state) |
| except DispatcherDownException: |
| # Retry forwarding the request. |
| return self.set_power_via_rpm(device_hostname, rpm_hostname, |
| rpm_outlet, hydra_hostname, new_state) |
| |
| |
| def queue_request(self, device_hostname, new_state): |
| """ |
| Forwards a request to change a device's (a dut or a servo) power state |
| to the appropriate dispatcher server. |
| |
| This call will block until the forwarded request returns. |
| |
| @param device_hostname: Hostname of the device whose power state we want |
| to change. |
| @param new_state: [ON, OFF, CYCLE] State to which we want to set the |
| device's outlet to. |
| |
| @return: True if the attempt to change power state was successful, |
| False otherwise. |
| |
| @raise RPMInfrastructureException: No dispatchers are available or can |
| be reached. |
| """ |
| # Remove any DNS Zone information and simplify down to just the hostname. |
| device_hostname = device_hostname.split('.')[0] |
| new_state = new_state.upper() |
| # Put new_state in all uppercase letters |
| if new_state not in VALID_STATE_VALUES: |
| logging.error('Received request to set device %s to invalid ' |
| 'state %s', device_hostname, new_state) |
| return False |
| logging.info('Received request to set device: %s to state: %s', |
| device_hostname, new_state) |
| powerunit_info = self._get_powerunit_info(device_hostname) |
| try: |
| return self._queue_once(powerunit_info, new_state) |
| except DispatcherDownException: |
| # Retry forwarding the request. |
| return self.queue_request(device_hostname, new_state) |
| |
| |
| def _queue_once(self, powerunit_info, new_state): |
| """Queue one request to the dispatcher.""" |
| dispatcher_uri = self._get_dispatcher(powerunit_info) |
| if not dispatcher_uri: |
| # No dispatchers available. |
| raise RPMInfrastructureException('No dispatchers available.') |
| client = xmlrpclib.ServerProxy(dispatcher_uri, allow_none=True) |
| try: |
| # Block on the request and return the result once it arrives. |
| return client.queue_request(powerunit_info, new_state) |
| except socket.error as er: |
| # Dispatcher Server is not reachable. Unregister it and retry. |
| logging.error("Can't reach Dispatch Server: %s. Error: %s", |
| dispatcher_uri, errno.errorcode[er.errno]) |
| if self.is_network_infrastructure_down(): |
| # No dispatchers can handle this request so raise an Exception |
| # to the caller. |
| raise RPMInfrastructureException('No dispatchers can be' |
| 'reached.') |
| logging.info('Will attempt forwarding request to other dispatch ' |
| 'servers.') |
| logging.error('Unregistering %s due to error. Recommend resetting ' |
| 'that dispatch server.', dispatcher_uri) |
| self.unregister_dispatcher(dispatcher_uri) |
| raise DispatcherDownException(dispatcher_uri) |
| |
| |
| def is_network_infrastructure_down(self): |
| """ |
| Check to see if we can communicate with any dispatcher servers. |
| |
| Only called in the situation that queuing a request to a dispatcher |
| server failed. |
| |
| @return: False if any dispatcher server is up and the rpm infrastructure |
| can still function. True otherwise. |
| """ |
| for dispatcher_entry in self._dispatcher_minheap: |
| dispatcher = xmlrpclib.ServerProxy( |
| dispatcher_entry[DISPATCHER_URI], allow_none=True) |
| try: |
| if dispatcher.is_up(): |
| # Atleast one dispatcher is alive so our network is fine. |
| return False |
| except socket.error: |
| # Can't talk to this dispatcher so keep looping. |
| pass |
| logging.error("Can't reach any dispatchers. Check frontend network " |
| 'status or all dispatchers are down.') |
| return True |
| |
| |
| def _get_powerunit_info(self, device_hostname): |
| """Get the power management unit information for a device. |
| |
| A device could be a chromeos dut or a servo. |
| 1) ChromeOS dut |
| Chromeos dut is managed by RPM. The related information |
| we need to know include rpm hostname, rpm outlet, hydra hostname. |
| Such information can be retrieved from afe_host_attributes table |
| from afe. A local LRU cache is used avoid hitting afe too often. |
| |
| 2) Servo |
| Servo is managed by POE. The related information we need to know |
| include poe hostname, poe interface. Such information is |
| stored in a local file and read into memory. |
| |
| @param device_hostname: A string representing the device's hostname. |
| |
| @returns: A PowerUnitInfo object. |
| @raises RPMInfrastructureException if failed to get the power |
| unit info. |
| |
| """ |
| if device_hostname.endswith('servo'): |
| return self._get_poe_powerunit_info(device_hostname) |
| else: |
| return self._get_rpm_powerunit_info(device_hostname) |
| |
| |
| def _get_poe_powerunit_info(self, device_hostname): |
| """Get the power management unit information for a POE controller. |
| |
| Servo is managed by POE. The related information we need to know |
| include poe hostname, poe interface. Such information is |
| stored in a local file and read into memory. |
| |
| @param device_hostname: A string representing the device's hostname. |
| |
| @returns: A PowerUnitInfo object. |
| @raises RPMInfrastructureException if failed to get the power |
| unit info. |
| |
| """ |
| with self._lock: |
| reload_info = utils.reload_servo_interface_mapping_if_necessary( |
| self._mapping_last_modified) |
| if reload_info: |
| self._mapping_last_modified, self._servo_interface = reload_info |
| switch_if_tuple = self._servo_interface.get(device_hostname) |
| if not switch_if_tuple: |
| raise RPMInfrastructureException( |
| 'Could not determine POE hostname for %s. ' |
| 'Please check the servo-interface mapping file.', |
| device_hostname) |
| else: |
| return utils.PowerUnitInfo( |
| device_hostname=device_hostname, |
| powerunit_type=utils.PowerUnitInfo.POWERUNIT_TYPES.POE, |
| powerunit_hostname=switch_if_tuple[0], |
| outlet=switch_if_tuple[1], |
| hydra_hostname=None) |
| |
| |
| |
| def _get_rpm_powerunit_info(self, device_hostname): |
| """Get the power management unit information for an RPM controller. |
| |
| Chromeos dut is managed by RPM. The related information |
| we need to know include rpm hostname, rpm outlet, hydra hostname. |
| Such information can be retrieved from afe_host_attributes table |
| from afe. A local LRU cache is used avoid hitting afe too often. |
| |
| @param device_hostname: A string representing the device's hostname. |
| |
| @returns: A PowerUnitInfo object. |
| @raises RPMInfrastructureException if failed to get the power |
| unit info. |
| |
| """ |
| with self._lock: |
| # Regular DUTs are managed by RPMs. |
| if device_hostname in self._rpm_info: |
| return self._rpm_info[device_hostname] |
| else: |
| hosts = self._afe.get_hosts(hostname=device_hostname) |
| if not hosts: |
| raise RPMInfrastructureException( |
| 'Can not retrieve rpm information ' |
| 'from AFE for %s, no host found.' % device_hostname) |
| else: |
| info = utils.PowerUnitInfo.get_powerunit_info(hosts[0]) |
| self._rpm_info[device_hostname] = info |
| return info |
| |
| |
| |
| def _get_dispatcher(self, powerunit_info): |
| """ |
| Private method that looks up or assigns a dispatcher server |
| responsible for communicating with the given RPM/POE. |
| |
| Will also call _check_dispatcher to make sure it is up before returning |
| it. |
| |
| @param powerunit_info: A PowerUnitInfo instance. |
| |
| @return: URI of dispatcher server responsible for the rpm/poe. |
| None if no dispatcher servers are available. |
| """ |
| powerunit_type = powerunit_info.powerunit_type |
| powerunit_hostname = powerunit_info.powerunit_hostname |
| with self._lock: |
| if self._rpm_dict.get(powerunit_hostname): |
| return self._rpm_dict[powerunit_hostname] |
| logging.info('No Dispatcher assigned for %s %s.', |
| powerunit_type, powerunit_hostname) |
| # Choose the least loaded dispatcher to communicate with the RPM. |
| try: |
| heap_entry = heapq.heappop(self._dispatcher_minheap) |
| except IndexError: |
| logging.error('Infrastructure Error: Frontend has no' |
| 'registered dispatchers to field out this ' |
| 'request!') |
| return None |
| dispatcher_uri = heap_entry[DISPATCHER_URI] |
| # Put this entry back in the heap with an RPM Count + 1. |
| heap_entry[RPM_COUNT] = heap_entry[RPM_COUNT] + 1 |
| heapq.heappush(self._dispatcher_minheap, heap_entry) |
| logging.info('Assigning %s for %s %s', dispatcher_uri, |
| powerunit_type, powerunit_hostname) |
| self._rpm_dict[powerunit_hostname] = dispatcher_uri |
| return dispatcher_uri |
| |
| |
| def register_dispatcher(self, dispatcher_uri): |
| """ |
| Called by a dispatcher server so that the frontend server knows it is |
| available to field out RPM requests. |
| |
| Adds an entry to the min heap and entry map for this dispatcher. |
| |
| @param dispatcher_uri: Address of dispatcher server we are registering. |
| """ |
| logging.info('Registering uri: %s as a rpm dispatcher.', dispatcher_uri) |
| with self._lock: |
| heap_entry = [DEFAULT_RPM_COUNT, dispatcher_uri] |
| heapq.heappush(self._dispatcher_minheap, heap_entry) |
| self._entry_dict[dispatcher_uri] = heap_entry |
| |
| |
| def unregister_dispatcher(self, uri_to_unregister): |
| """ |
| Called by a dispatcher server as it exits so that the frontend server |
| knows that it is no longer available to field out requests. |
| |
| Assigns an rpm count of -1 to this dispatcher so that it will be pushed |
| out of the min heap. |
| |
| Removes from _rpm_dict all entries with the value of this dispatcher so |
| that those RPM's can be reassigned to a new dispatcher. |
| |
| @param uri_to_unregister: Address of dispatcher server we are |
| unregistering. |
| """ |
| logging.info('Unregistering uri: %s as a rpm dispatcher.', |
| uri_to_unregister) |
| with self._lock: |
| heap_entry = self._entry_dict.get(uri_to_unregister) |
| if not heap_entry: |
| logging.warning('%s was not registered.', uri_to_unregister) |
| return |
| # Set this entry's RPM_COUNT to TERMINATED (-1). |
| heap_entry[RPM_COUNT] = TERMINATED |
| # Remove all RPM mappings. |
| for rpm, dispatcher in self._rpm_dict.items(): |
| if dispatcher == uri_to_unregister: |
| self._rpm_dict[rpm] = None |
| self._entry_dict[uri_to_unregister] = None |
| # Re-sort the heap and remove any terminated dispatchers. |
| heapq.heapify(self._dispatcher_minheap) |
| self._remove_terminated_dispatchers() |
| |
| |
| def _remove_terminated_dispatchers(self): |
| """ |
| Peek at the head of the heap and keep popping off values until there is |
| a non-terminated dispatcher at the top. |
| """ |
| # Heapq guarantees the head of the heap is in the '0' index. |
| try: |
| # Peek at the next element in the heap. |
| top_of_heap = self._dispatcher_minheap[0] |
| while top_of_heap[RPM_COUNT] is TERMINATED: |
| # Pop off the top element. |
| heapq.heappop(self._dispatcher_minheap) |
| # Peek at the next element in the heap. |
| top_of_heap = self._dispatcher_minheap[0] |
| except IndexError: |
| # No more values in the heap. Can be thrown by both minheap[0] |
| # statements. |
| pass |
| |
| |
| def suspend_emails(self, hours): |
| """Suspend email notifications. |
| |
| @param hours: How many hours to suspend email notifications. |
| """ |
| if self._email_handler: |
| self._email_handler.suspend_emails(hours) |
| |
| |
| def resume_emails(self): |
| """Resume email notifications.""" |
| if self._email_handler: |
| self._email_handler.resume_emails() |
| |
| |
| if __name__ == '__main__': |
| """ |
| Main function used to launch the frontend server. Creates an instance of |
| RPMFrontendServer and registers it to a MultiThreadedXMLRPCServer instance. |
| """ |
| if len(sys.argv) != 2: |
| print 'Usage: ./%s <log_file_dir>.' % sys.argv[0] |
| sys.exit(1) |
| |
| email_handler = rpm_logging_config.set_up_logging_to_file( |
| sys.argv[1], LOG_FILENAME_FORMAT) |
| frontend_server = RPMFrontendServer(email_handler=email_handler) |
| # We assume that external clients will always connect to us via the |
| # hostname, so listening on the hostname ensures we pick the right network |
| # interface. |
| address = socket.gethostname() |
| port = rpm_config.getint('RPM_INFRASTRUCTURE', 'frontend_port') |
| server = MultiThreadedXMLRPCServer((address, port), allow_none=True) |
| server.register_instance(frontend_server) |
| logging.info('Listening on %s port %d', address, port) |
| server.serve_forever() |
