python/qemu: Change ConsoleSocket to optionally drain socket. (80ded8e9) · Commits · SUMMER2020 / students / proj-2021291

python/qemu/console_socket.py

+55 −37

Original line number	Diff line number	Diff line
		@@ -13,68 +13,75 @@ which can drain a socket and optionally dump the bytes to file.
		# the COPYING file in the top-level directory.
		#

		import asyncore
		import socket
		import threading
		from collections import deque
		import time


		class ConsoleSocket(asyncore.dispatcher):
		class ConsoleSocket(socket.socket):
		"""
		ConsoleSocket represents a socket attached to a char device.

		Drains the socket and places the bytes into an in memory buffer
		for later processing.
		Optionally (if drain==True), drains the socket and places the bytes
		into an in memory buffer for later processing.

		Optionally a file path can be passed in and we will also
		dump the characters to this file for debugging purposes.
		"""
		def __init__(self, address, file=None):
		def __init__(self, address, file=None, drain=False):
		self._recv_timeout_sec = 300
		self._sleep_time = 0.5
		self._buffer = deque()
		self._asyncore_thread = None
		self._sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
		self._sock.connect(address)
		socket.socket.__init__(self, socket.AF_UNIX, socket.SOCK_STREAM)
		self.connect(address)
		self._logfile = None
		if file:
		self._logfile = open(file, "w")
		asyncore.dispatcher.__init__(self, sock=self._sock)
		self._open = True
		self._thread_start()
		if drain:
		self._drain_thread = self._thread_start()
		else:
		self._drain_thread = None

		def _thread_start(self):
		"""Kick off a thread to wait on the asyncore.loop"""
		if self._asyncore_thread is not None:
		return
		self._asyncore_thread = threading.Thread(target=asyncore.loop,
		kwargs={'timeout':1})
		self._asyncore_thread.daemon = True
		self._asyncore_thread.start()
		def _drain_fn(self):
		"""Drains the socket and runs while the socket is open."""
		while self._open:
		try:
		self._drain_socket()
		except socket.timeout:
		# The socket is expected to timeout since we set a
		# short timeout to allow the thread to exit when
		# self._open is set to False.
		time.sleep(self._sleep_time)

		def handle_close(self):
		"""redirect close to base class"""
		# Call the base class close, but not self.close() since
		# handle_close() occurs in the context of the thread which
		# self.close() attempts to join.
		asyncore.dispatcher.close(self)
		def _thread_start(self):
		"""Kick off a thread to drain the socket."""
		# Configure socket to not block and timeout.
		# This allows our drain thread to not block
		# on recieve and exit smoothly.
		socket.socket.setblocking(self, False)
		socket.socket.settimeout(self, 1)
		drain_thread = threading.Thread(target=self._drain_fn)
		drain_thread.daemon = True
		drain_thread.start()
		return drain_thread

		def close(self):
		"""Close the base object and wait for the thread to terminate"""
		if self._open:
		self._open = False
		asyncore.dispatcher.close(self)
		if self._asyncore_thread is not None:
		thread, self._asyncore_thread = self._asyncore_thread, None
		if self._drain_thread is not None:
		thread, self._drain_thread = self._drain_thread, None
		thread.join()
		socket.socket.close(self)
		if self._logfile:
		self._logfile.close()
		self._logfile = None

		def handle_read(self):
		def _drain_socket(self):
		"""process arriving characters into in memory _buffer"""
		data = asyncore.dispatcher.recv(self, 1)
		data = socket.socket.recv(self, 1)
		# latin1 is needed since there are some chars
		# we are receiving that cannot be encoded to utf-8
		# such as 0xe2, 0x80, 0xA6.
		@@ -85,27 +92,38 @@ class ConsoleSocket(asyncore.dispatcher):
		for c in string:
		self._buffer.extend(c)

		def recv(self, buffer_size=1):
		def recv(self, bufsize=1):
		"""Return chars from in memory buffer.
		Maintains the same API as socket.socket.recv.
		"""
		if self._drain_thread is None:
		# Not buffering the socket, pass thru to socket.
		return socket.socket.recv(self, bufsize)
		start_time = time.time()
		while len(self._buffer) < buffer_size:
		while len(self._buffer) < bufsize:
		time.sleep(self._sleep_time)
		elapsed_sec = time.time() - start_time
		if elapsed_sec > self._recv_timeout_sec:
		raise socket.timeout
		chars = ''.join([self._buffer.popleft() for i in range(buffer_size)])
		chars = ''.join([self._buffer.popleft() for i in range(bufsize)])
		# We choose to use latin1 to remain consistent with
		# handle_read() and give back the same data as the user would
		# receive if they were reading directly from the
		# socket w/o our intervention.
		return chars.encode("latin1")

		def set_blocking(self):
		"""Maintain compatibility with socket API"""
		pass
		def setblocking(self, value):
		"""When not draining we pass thru to the socket,
		since when draining we control socket blocking.
		"""
		if self._drain_thread is None:
		socket.socket.setblocking(self, value)

		def settimeout(self, seconds):
		"""Set current timeout on recv"""
		"""When not draining we pass thru to the socket,
		since when draining we control the timeout.
		"""
		if seconds is not None:
		self._recv_timeout_sec = seconds
		if self._drain_thread is None:
		socket.socket.settimeout(self, seconds)

python/qemu/machine.py

+4 −9

Original line number	Diff line number	Diff line
		@@ -23,7 +23,6 @@ import os
		import subprocess
		import shutil
		import signal
		import socket
		import tempfile
		from typing import Optional, Type
		from types import TracebackType
		@@ -673,12 +672,8 @@ class QEMUMachine:
		Returns a socket connected to the console
		"""
		if self._console_socket is None:
		if self._drain_console:
		self._console_socket = console_socket.ConsoleSocket(
		self._console_address,
		file=self._console_log_path)
		else:
		self._console_socket = socket.socket(socket.AF_UNIX,
		socket.SOCK_STREAM)
		self._console_socket.connect(self._console_address)
		file=self._console_log_path,
		drain=self._drain_console)
		return self._console_socket