<commit_before>import subprocess
import unittest
import platform

class NoFreePortException(Exception):
	pass

class FreePortFinder:
	def __init__(self):
		self.last_allocated = None
		self.preallocated = set()

	def add_allocated_port(self, port):
		if isinstance(port, list):
			self.preallocated.update(port)
		else:
			self.preallocated.add(port)

	def remove_allocated_port(self, port):
		if port in self.preallocated:
			self.preallocated.remove(port)

	def free_in_range(self, lower, upper, complete_scan=False):
		start = lower
		using_last_allocated = False
		if self.last_allocated and not complete_scan:
			using_last_allocated = True
			if self.last_allocated < lower or self.last_allocated > upper:
				# Range has changed!
				start = lower
			else:
				start = self.last_allocated + 1
		# Fetch the current list of ports.
		output = self.fetch_netstat()
		# Now begin the search.
		port = start
		free = (output.find(":%d " % port) == -1) and (port not in self.preallocated)
		while not free and port < upper:
			port += 1
			free = (output.find(":%d " % port) == -1) and (port not in self.preallocated)
		if not free:
			# If we were using the last allocated port,
			# we might not have scanned the bottom of the range.
			if using_last_allocated:
				# Scan again from the bottom.
				self.free_in_range(lower, upper, complete_scan=True)
			else:
				raise NoFreePortException("Can't find a free port between %d and %d." % (lower, upper))
		# Return the port.
		self.last_allocated = port
		return port

	def in_use(self, port):
		"""
		Quickly check if a port is in use. This is not a way to allocate
		ports.
		"""
		output = self.fetch_netstat()
		return (output.find(":%d " % port) != -1)

	def fetch_netstat(self):
		# TODO: This is HIGHLY platform specific, but it's fast.
		# TODO: This will block the whole process... but only for ~20ms.
		# NOTE: The docs say not to use subprocess.PIPE with the command
		# below because it can cause the call to block if that pipe
		# fills up. However, I'm using this to prevent the stderr from
		# being printed out - and relying on the fact that it won't
		# print very much output to stderr.
		if platform.system() == 'Linux':
			command = ['netstat', '-ltnp']
		if platform.system() == 'Windows':
			command = ['netstat', '-na', '-p', 'tcp']
		output = subprocess.check_output(command, stderr=subprocess.PIPE)
		# Post process the output a little bit.
		raw_lines = output.split("\n")
		lines = []
		for line in raw_lines:
			if line.find("LISTEN") != -1:
				lines.append(line)
		output = "\n".join(lines)
		return output

class FreePortFinderTest(unittest.TestCase):
	def test_free(self):
		finder = FreePortFinder()

		port = finder.free_in_range(10000, 10100)
		self.assertTrue(port >= 10000 and port <= 10100)

	def test_allocated(self):
		finder = FreePortFinder()

		port = finder.free_in_range(10000, 10100)
		self.assertEquals(port, 10000)

		# Mask out the following ports.
		finder.add_allocated_port(10000)
		finder.add_allocated_port(10002)

		# And now we should be returned ports around that.
		port = finder.free_in_range(10000, 10100)
		self.assertNotIn(port, [10000, 10002])
		port = finder.free_in_range(10000, 10100)
		self.assertNotIn(port, [10000, 10002])

		# Add a few more, via a list.
		finder.add_allocated_port([10004, 10006])<commit_msg>Updated the port allocator to have a helper function to wait for a port to be in a certain state.<commit_after>import subprocess
import unittest
import platform
import time

class NoFreePortException(Exception):
	pass

class FreePortFinder(object):
	def __init__(self):
		self.last_allocated = None
		self.preallocated = set()

	def add_allocated_port(self, port):
		if isinstance(port, list):
			self.preallocated.update(port)
		else:
			self.preallocated.add(port)

	def remove_allocated_port(self, port):
		if port in self.preallocated:
			self.preallocated.remove(port)

	def free_in_range(self, lower, upper, complete_scan=False):
		start = lower
		using_last_allocated = False
		if self.last_allocated and not complete_scan:
			using_last_allocated = True
			if self.last_allocated < lower or self.last_allocated > upper:
				# Range has changed!
				start = lower
			else:
				start = self.last_allocated + 1
		# Fetch the current list of ports.
		output = self.fetch_netstat()
		# Now begin the search.
		port = start
		free = (output.find(":%d " % port) == -1) and (port not in self.preallocated)
		while not free and port < upper:
			port += 1
			free = (output.find(":%d " % port) == -1) and (port not in self.preallocated)
		if not free:
			# If we were using the last allocated port,
			# we might not have scanned the bottom of the range.
			if using_last_allocated:
				# Scan again from the bottom.
				self.free_in_range(lower, upper, complete_scan=True)
			else:
				raise NoFreePortException("Can't find a free port between %d and %d." % (lower, upper))
		# Return the port.
		self.last_allocated = port
		return port

	def in_use(self, port):
		"""
		Quickly check if a port is in use. This is not a way to allocate
		ports.
		"""
		output = self.fetch_netstat()
		return (output.find(":%d " % port) != -1)

	def fetch_netstat(self):
		# TODO: This is HIGHLY platform specific, but it's fast.
		# TODO: This will block the whole process... but only for ~20ms.
		# NOTE: The docs say not to use subprocess.PIPE with the command
		# below because it can cause the call to block if that pipe
		# fills up. However, I'm using this to prevent the stderr from
		# being printed out - and relying on the fact that it won't
		# print very much output to stderr.
		if platform.system() == 'Linux':
			command = ['netstat', '-ltnp']
		if platform.system() == 'Windows':
			command = ['netstat', '-na', '-p', 'tcp']
		output = subprocess.check_output(command, stderr=subprocess.PIPE)
		# Post process the output a little bit.
		raw_lines = output.split("\n")
		lines = []
		for line in raw_lines:
			if line.find("LISTEN") != -1:
				lines.append(line)
		output = "\n".join(lines)
		return output

	def wait_until_port_used(self, io_loop, port, timeout, callback, timeout_callback):
		self.wait_until_port_state(io_loop, port, True, timeout, callback, timeout_callback)
	def wait_until_port_free(self, io_loop, port, timeout, callback, timeout_callback):
		self.wait_until_port_state(io_loop, port, False, timeout, callback, timeout_callback)

	def wait_until_port_state(self, io_loop, port, state, timeout, callback, timeout_callback):
		# Wait until the port is no longer free.
		end_timeout = time.time() + timeout

		def wait_for_state():
			if self.in_use(port) == state:
				# And say that we're done.
				callback("In appropriate state.")
			else:
				if time.time() > end_timeout:
					timeout_callback("Failed to end up in appropriate state in time.")
				else:
					# Wait a little bit longer.
					io_loop.add_timeout(time.time() + 0.1, wait_for_state)

		io_loop.add_timeout(time.time() + 0.1, wait_for_state)

class FreePortFinderTest(unittest.TestCase):
	def test_free(self):
		finder = FreePortFinder()

		port = finder.free_in_range(10000, 10100)
		self.assertTrue(port >= 10000 and port <= 10100)

	def test_allocated(self):
		finder = FreePortFinder()

		port = finder.free_in_range(10000, 10100)
		self.assertEquals(port, 10000)

		# Mask out the following ports.
		finder.add_allocated_port(10000)
		finder.add_allocated_port(10002)

		# And now we should be returned ports around that.
		port = finder.free_in_range(10000, 10100)
		self.assertNotIn(port, [10000, 10002])
		port = finder.free_in_range(10000, 10100)
		self.assertNotIn(port, [10000, 10002])

		# Add a few more, via a list.
		finder.add_allocated_port([10004, 10006])
{'fraction_non_alphanumeric': 0.10708252111760884, 'fraction_numerical': 0.02729044834307992, 'mean_word_length': 3.855520504731861, 'pattern_counts': {'":': 6, '<': 9, '<?xml version=': 0, '>': 8, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2812435', 'n_tokens_mistral': 2968, 'n_tokens_neox': 2522, 'n_words': 809}