| #!/usr/bin/python -u |
| # -*- coding: utf-8 -*- |
| # |
| # Copyright (c) 2014 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 jsonrpclib |
| import logging |
| import socket |
| import threading |
| import time |
| |
| import factory_common # pylint: disable=W0611 |
| from cros.factory.goofy_split.discoverer import DUTDiscoverer |
| from cros.factory.goofy_split.discoverer import PresenterDiscoverer |
| from cros.factory.test import utils |
| from cros.factory.test.network import GetAllIPs |
| from cros.factory.utils.jsonrpc_utils import JSONRPCServer |
| from cros.factory.utils.jsonrpc_utils import TimeoutJSONRPCTransport |
| from cros.factory.utils.net_utils import GetEthernetInterfaces |
| from cros.factory.utils.net_utils import GetEthernetIp |
| |
| |
| # Standard RPC ports. These may be replaced by unit tests. |
| PRESENTER_LINK_RPC_PORT = 4020 |
| DUT_LINK_RPC_PORT = 4021 |
| |
| |
| class LinkDownError(Exception): |
| """The exception raised on RPC calls when the link is down.""" |
| pass |
| |
| class PresenterLinkManager(object): |
| """The manager that runs on the DUT to maintain a link to the presenter.""" |
| def __init__(self, |
| check_interval=5, |
| methods=None, |
| handshake_timeout=0.3, |
| rpc_timeout=1, |
| connect_hook=None, |
| disconnect_hook=None): |
| self._check_interval = check_interval |
| self._handshake_timeout = handshake_timeout |
| self._rpc_timeout = rpc_timeout |
| self._connect_hook = connect_hook |
| self._disconnect_hook = disconnect_hook |
| self._suspend_deadline = None |
| self._methods = methods or {} |
| self._methods.update({'Announce': self._PresenterAnnounce}) |
| self._reported_failure = set() |
| self._presenter_connected = False |
| self._presenter_ip = None |
| self._presenter_proxy = None |
| self._presenter_announcement = None |
| self._discoverer = PresenterDiscoverer(PRESENTER_LINK_RPC_PORT) |
| self._kick_event = threading.Event() |
| self._abort_event = threading.Event() |
| self._server = JSONRPCServer(port=DUT_LINK_RPC_PORT, methods=self._methods) |
| self._server.Start() |
| self._thread = threading.Thread(target=self.MonitorLink) |
| self._thread.start() |
| |
| def __getattr__(self, name): |
| """A wrapper that proxies the RPC calls to the real server proxy.""" |
| if not self._presenter_connected: |
| raise LinkDownError() |
| try: |
| return self._presenter_proxy.__getattr__(name) |
| except AttributeError: |
| # _presenter_proxy is None. Link is probably down. |
| raise LinkDownError() |
| |
| def Stop(self): |
| """Stops and destroys the link manager.""" |
| self._server.Destroy() |
| self._abort_event.set() |
| self._kick_event.set() # Kick the thread |
| self._thread.join() |
| |
| def SuspendMonitoring(self, interval_sec): |
| """Suspend monitoring of connection for a given period. |
| |
| Args: |
| interval_sec: Number of seconds to suspend. |
| """ |
| self._suspend_deadline = time.time() + interval_sec |
| self._presenter_proxy.SuspendMonitoring(interval_sec) |
| |
| def ResumeMonitoring(self): |
| """Immediately resume suspended monitoring of connection.""" |
| self._suspend_deadline = None |
| self.Kick() |
| self._presenter_proxy.ResumeMonitoring() |
| |
| def PresenterIsAlive(self): |
| """Pings the presenter.""" |
| if not self._presenter_connected: |
| return False |
| try: |
| return self._presenter_proxy.IsAlive() |
| except (socket.error, socket.timeout, AttributeError): |
| return False |
| |
| def _PresenterAnnounce(self, my_ip, presenter_ips): |
| self._presenter_announcement = (my_ip, presenter_ips) |
| self._kick_event.set() |
| |
| def _HandlePresenterAnnouncement(self): |
| my_ip, presenter_ips = self._presenter_announcement |
| self._presenter_announcement = None |
| for presenter_ip in presenter_ips: |
| self._MakePresenterConnection(my_ip, presenter_ip) |
| if self._presenter_connected: |
| return |
| |
| def _MakeTimeoutServerProxy(self, presenter_ip, timeout): |
| return jsonrpclib.Server('http://%s:%d/' % |
| (presenter_ip,PRESENTER_LINK_RPC_PORT), |
| transport=TimeoutJSONRPCTransport(timeout)) |
| |
| def _MakePresenterConnection(self, my_ip, presenter_ip): |
| """Attempts to connect the the presenter. |
| |
| Args: |
| my_ip: The IP address of this DUT received from the presenter; None to guess. |
| presenter_ip: The IP address of the presenter. |
| """ |
| if self._presenter_connected and self._presenter_ip == presenter_ip: |
| return |
| |
| log = (logging.info if presenter_ip not in self._reported_failure else |
| lambda *args: None) |
| |
| try: |
| log('Attempting to connect to presenter %s', presenter_ip) |
| self._presenter_proxy = self._MakeTimeoutServerProxy(presenter_ip, |
| self._handshake_timeout) |
| self._presenter_ip = presenter_ip |
| self._presenter_proxy.IsAlive() |
| |
| # Presenter is alive. Let's register! |
| log('Registering to presenter %s', presenter_ip) |
| if not my_ip: |
| if utils.in_chroot(): |
| my_ip = '127.0.0.1' |
| else: |
| my_ip = map(GetEthernetIp, GetEthernetInterfaces()) |
| my_ip = [x for x in my_ip if x != '127.0.0.1'] |
| log('Trying available IP addresses %s', my_ip) |
| elif type(my_ip) != list: |
| my_ip = [my_ip] |
| |
| for ip in my_ip: |
| log('Trying IP address %s', ip) |
| self._presenter_proxy.Register(ip) |
| |
| # Make sure the presenter sees us |
| log('Registered. Checking connection.') |
| if not self._presenter_proxy.ConnectionGood(): |
| log('Registration failed.') |
| continue |
| self._presenter_connected = True |
| logging.info('Connected to presenter %s', presenter_ip) |
| # Now that we are connected, use a longer timeout for the proxy |
| self._presenter_proxy = self._MakeTimeoutServerProxy(presenter_ip, |
| self._rpc_timeout) |
| if presenter_ip in self._reported_failure: |
| self._reported_failure.remove(presenter_ip) |
| if self._connect_hook: |
| self._connect_hook(presenter_ip) |
| return |
| except (socket.error, socket.timeout): |
| pass |
| |
| # If we are here, we failed to make connection. Clean up. |
| self._presenter_ip = None |
| self._presenter_proxy = None |
| self._reported_failure.add(presenter_ip) |
| log('Connection failed.') |
| |
| def CheckPresenterConnection(self): |
| """Check the connection to the presenter. |
| |
| If the connection is down, put ourselves into disconnected state and attempt |
| to establish the connection again. |
| """ |
| if self._presenter_connected: |
| if self.PresenterIsAlive(): |
| return # everything's fine |
| else: |
| logging.info('Lost connection to presenter %s', self._presenter_ip) |
| self._presenter_connected = False |
| self._presenter_ip = None |
| self._presenter_proxy = None |
| if self._disconnect_hook: |
| self._disconnect_hook() |
| |
| ips = self._discoverer.Discover() |
| if not ips: |
| return |
| if type(ips) != list: |
| ips = [ips] |
| for ip in ips: |
| self._MakePresenterConnection(None, ip) |
| if self._presenter_connected: |
| return |
| |
| def MonitorLink(self): |
| while True: |
| self._kick_event.wait(self._check_interval) |
| self._kick_event.clear() |
| if self._abort_event.isSet(): |
| return |
| if self._suspend_deadline: |
| if time.time() > self._suspend_deadline: |
| self._suspend_deadline = None |
| else: |
| if self._presenter_announcement: |
| self._HandlePresenterAnnouncement() |
| else: |
| self.CheckPresenterConnection() |
| |
| |
| class DUTLinkManager(object): |
| """The manager that runs on the presenter to maintain the link to the DUT.""" |
| def __init__(self, |
| check_interval=5, |
| methods=None, |
| rpc_timeout=1, |
| connect_hook=None, |
| disconnect_hook=None): |
| self._check_interval = check_interval |
| self._rpc_timeout = rpc_timeout |
| self._connect_hook = connect_hook |
| self._disconnect_hook = disconnect_hook |
| self._suspend_deadline = None |
| self._methods = methods or {} |
| self._methods.update({'Register': self._DUTRegister, |
| 'ConnectionGood': self.DUTIsAlive, |
| 'SuspendMonitoring': self.SuspendMonitoring, |
| 'ResumeMonitoring': self.ResumeMonitoring}) |
| self._reported_announcement = set() |
| self._dut_proxy = None |
| self._dut_ip = None |
| self._dut_connected = False |
| self._lock = threading.Lock() |
| self._kick_event = threading.Event() |
| self._abort_event = threading.Event() |
| self._discoverer = DUTDiscoverer(DUT_LINK_RPC_PORT) |
| self._server = JSONRPCServer(port=PRESENTER_LINK_RPC_PORT, |
| methods=self._methods) |
| self._server.Start() |
| self._thread = threading.Thread(target=self.MonitorLink) |
| self._thread.start() |
| |
| def __getattr__(self, name): |
| """A wrapper that proxies the RPC calls to the real server proxy.""" |
| if not self._dut_connected: |
| raise LinkDownError() |
| try: |
| return self._dut_proxy.__getattr__(name) |
| except AttributeError: |
| # _dut_proxy is None. Link is probably down. |
| raise LinkDownError() |
| |
| def Stop(self): |
| """Stops and destroys the link manager.""" |
| self._server.Destroy() |
| self._abort_event.set() |
| self._kick_event.set() |
| self._thread.join() |
| |
| def SuspendMonitoring(self, interval_sec): |
| """Suspend monitoring of connection for a given period. |
| |
| Args: |
| interval_sec: Number of seconds to suspend. |
| """ |
| self._suspend_deadline = time.time() + interval_sec |
| |
| def ResumeMonitoring(self): |
| """Immediately resume suspended monitoring of connection.""" |
| self._suspend_deadline = None |
| self.Kick() |
| |
| def _MakeTimeoutServerProxy(self, dut_ip, timeout): |
| return jsonrpclib.Server('http://%s:%d/' % (dut_ip, DUT_LINK_RPC_PORT), |
| transport=TimeoutJSONRPCTransport(timeout)) |
| |
| def _DUTRegister(self, dut_ip): |
| with self._lock: |
| try: |
| self._dut_ip = dut_ip |
| self._dut_proxy = self._MakeTimeoutServerProxy(dut_ip, |
| self._rpc_timeout) |
| self._dut_proxy.IsAlive() |
| self._dut_connected = True |
| logging.info('DUT %s registered', dut_ip) |
| self._reported_announcement.clear() |
| if self._connect_hook: |
| self._connect_hook(dut_ip) |
| except (socket.error, socket.timeout): |
| self._dut_ip = None |
| self._dut_proxy = None |
| |
| def DUTIsAlive(self): |
| """Pings the DUT.""" |
| if not self._dut_connected: |
| return False |
| try: |
| self._dut_proxy.IsAlive() |
| return True |
| except (socket.error, socket.timeout, AttributeError): |
| return False |
| |
| def CheckDUTConnection(self): |
| """Check the connection to the DUT. |
| |
| If the connection is down, put ourselves into disconnected state and start |
| announcing ourselves to potential DUTs again. |
| """ |
| if self._lock.acquire(False): |
| try: |
| if self._dut_connected: |
| if self.DUTIsAlive(): |
| return # All good! |
| else: |
| logging.info('Disconnected from DUT %s', self._dut_ip) |
| self._dut_connected = False |
| self._dut_ip = None |
| self._dut_proxy = None |
| if self._disconnect_hook: |
| self._disconnect_hook() |
| |
| ips = self._discoverer.Discover() |
| if not ips: |
| return |
| if type(ips) != list: |
| ips = [ips] |
| for ip in ips: |
| try: |
| # We don't get response from the DUT for announcement, so let's |
| # keep the timeout short. |
| proxy = self._MakeTimeoutServerProxy(ip, timeout=0.05) |
| my_ips = GetAllIPs() |
| if ip not in self._reported_announcement: |
| logging.info('Announcing to DUT %s: presenter ip is %s', |
| ip, my_ips) |
| self._reported_announcement.add(ip) |
| proxy.Announce(ip, my_ips) |
| except (socket.error, socket.timeout): |
| pass |
| finally: |
| self._lock.release() |
| |
| def Kick(self): |
| """Kick the link manager to check the connection or announce ourselves.""" |
| self._kick_event.set() |
| |
| def MonitorLink(self): |
| while True: |
| if self._suspend_deadline: |
| if time.time() > self._suspend_deadline: |
| self._suspend_deadline = None |
| else: |
| self.CheckDUTConnection() |
| self._kick_event.wait(self._check_interval) |
| self._kick_event.clear() |
| if self._abort_event.isSet(): |
| return |
