blob: c9ef11303562332b2075324246ae578d83f7b698 [file] [log] [blame]
#!/usr/bin/python -u
#
# Copyright 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.
from __future__ import print_function
import collections
import logging
import subprocess
import threading
import time
from ws4py.websocket import WebSocket
import factory_common # pylint: disable=W0611
from cros.factory.test import factory
from cros.factory.test.event import Event
from cros.factory.test.event import EventClient
from cros.factory.test.web_socket_utils import WebSocketHandshake
from cros.factory.utils.process_utils import Spawn
from cros.factory.utils.string_utils import DecodeUTF8
# Number of lines to buffer for new clients.
TAIL_BUFFER_SIZE = 10
class WebSocketManager(object):
'''Object to manage web sockets for Goofy.
Brokers between events in the event client infrastructure
and on web sockets. Also tails the console log and sends
events on web sockets when new bytes become available.
Each Goofy instance is associated with a UUID. When a new web
socket is created, we send a hello event on the socket with the
current UUID. If we receive a keepalive event with the wrong
UUID, we disconnect the client. This insures that we are always
talking to a client that has a complete picture of our state
(i.e., if the server restarts, the client must restart as well).
Properties:
tail_buffer: A rotating buffer of the last TAIL_BUFFER_SIZE lines,
to give to new web clients.
'''
def __init__(self, uuid):
self.uuid = uuid
self.lock = threading.Lock()
self.web_sockets = set()
self.event_client = None
self.tail_process = None
self.has_confirmed_socket = threading.Event()
self.event_client = EventClient(callback=self._handle_event,
name='WebSocketManager')
self.tail_process = Spawn(
["tail", "-F", factory.CONSOLE_LOG_PATH],
ignore_stdin=True,
stdout=subprocess.PIPE)
self.tail_thread = threading.Thread(target=self._tail_console)
self.tail_thread.start()
self.closed = False
self.tail_buffer = collections.deque()
def close(self):
with self.lock:
if self.closed:
return
self.closed = True
if self.event_client:
self.event_client.close()
self.event_client = None
with self.lock:
web_sockets = list(self.web_sockets)
for web_socket in web_sockets:
web_socket.close_connection()
if self.tail_process:
self.tail_process.kill()
self.tail_process.wait()
if self.tail_thread:
self.tail_thread.join()
def has_sockets(self):
'''Returns true if any web sockets are currently connected.'''
with self.lock:
return len(self.web_sockets) > 0
def handle_web_socket(self, request):
'''Runs a web socket in the current thread.
request: A RequestHandler object containing the request.
'''
if not WebSocketHandshake(request):
return
class MyWebSocket(WebSocket):
def __init__(self, **kwargs):
# Add a per-socket lock to use for sending, since ws4py is not
# thread-safe.
self.send_lock = threading.Lock()
super(MyWebSocket, self).__init__(**kwargs)
def received_message(socket_self, message): # pylint: disable=E0213
event = Event.from_json(str(message))
if event.type == Event.Type.KEEPALIVE:
if event.uuid == self.uuid:
if not self.has_confirmed_socket.is_set():
logging.info('Chrome UI has come up')
self.has_confirmed_socket.set()
else:
logging.warning('Disconnecting web socket with '
'incorrect UUID')
socket_self.close_connection()
else:
self.event_client.post_event(event)
web_socket = MyWebSocket(sock=request.connection)
with self.lock:
lines = list(self.tail_buffer)
with web_socket.send_lock:
web_socket.send(Event(Event.Type.HELLO,
uuid=self.uuid).to_json())
for line in lines:
# Send the last n lines.
web_socket.send(
Event(Event.Type.LOG,
message=DecodeUTF8(line)).to_json())
try:
with self.lock:
self.web_sockets.add(web_socket)
logging.info('Running web socket')
web_socket.run()
logging.info('Web socket closed gracefully')
except: # pylint: disable=W0702
logging.exception('Web socket closed with exception')
finally:
with self.lock:
self.web_sockets.discard(web_socket)
def wait(self):
'''Waits for one socket to connect successfully.'''
while not self.has_confirmed_socket.is_set():
# Wait at most 100 ms at a time; without a timeout, this seems
# to eat SIGINT signals.
self.has_confirmed_socket.wait(0.1)
def _tail_console(self):
'''Tails the console log, generating an event whenever a new
line is available.
We send this event only to web sockets (not to event clients
in general) since only the UI is interested in these log
lines.
'''
# tail seems to have a bug where, when outputting to a pipe, it
# doesn't output the first batch of data until it receives some
# new output. Let tail start up, then output a single line to
# wake it up. This is a terrible hack, but it's better than
# missing a bunch of lines. A better fix might involve emulating
# tail directly in Python.
def target():
time.sleep(0.5)
factory.console.info('Opened console.')
thread = threading.Thread(target=target)
thread.daemon = True
thread.start()
while True:
line = self.tail_process.stdout.readline()
if line == '':
break
with self.lock:
self.tail_buffer.append(line)
while len(self.tail_buffer) > TAIL_BUFFER_SIZE:
self.tail_buffer.popleft()
self._handle_event(
Event(Event.Type.LOG,
message=DecodeUTF8(line).rstrip("\n")))
def _handle_event(self, event):
'''Sends an event to each open WebSocket client.'''
with self.lock:
web_sockets = list(self.web_sockets)
if not web_sockets:
return
event_json = event.to_json()
for web_socket in web_sockets:
try:
with web_socket.send_lock:
web_socket.send(event_json)
except: # pylint: disable=W0702
logging.exception('Unable to send event on web socket')