-
Notifications
You must be signed in to change notification settings - Fork 28
Expand file tree
/
Copy pathtcp.py
More file actions
193 lines (155 loc) · 6.64 KB
/
tcp.py
File metadata and controls
193 lines (155 loc) · 6.64 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
# Programmer friendly subprocess wrapper.
#
# Author: Peter Odding <peter@peterodding.com>
# Last Change: March 2, 2020
# URL: https://executor.readthedocs.io
"""
Miscellaneous TCP networking functionality.
The functionality in this module originated in the :class:`executor.ssh.server`
module with the purpose of facilitating a robust automated test suite for the
:class:`executor.ssh.client` module. While working on SSH tunnel support I
needed similar logic again and I decided to extract this code from the
:class:`executor.ssh.server` module.
"""
# Standard library modules.
import itertools
import logging
import random
import socket
# Modules included in our package.
from executor import ExternalCommand
# External dependencies.
from humanfriendly import Timer, format_timespan
from humanfriendly.terminal.spinners import Spinner
from humanfriendly.text import format, pluralize
from property_manager import (
PropertyManager,
lazy_property,
mutable_property,
required_property,
set_property,
)
# Public identifiers that require documentation.
__all__ = (
'EphemeralPortAllocator',
'EphemeralTCPServer',
'TimeoutError',
'WaitUntilConnected',
'logger',
)
# Initialize a logger.
logger = logging.getLogger(__name__)
class WaitUntilConnected(PropertyManager):
"""Wait for a TCP endpoint to start accepting connections."""
@mutable_property
def connect_timeout(self):
"""The timeout in seconds for individual connection attempts (a number, defaults to 2)."""
return 2
@property
def endpoint(self):
"""A human friendly representation of the TCP endpoint (a string containing a URL)."""
return format("%s://%s:%i", self.scheme, self.hostname, self.port_number)
@mutable_property
def hostname(self):
"""The host name or IP address to connect to (a string, defaults to ``localhost``)."""
return 'localhost'
@property
def is_connected(self):
""":data:`True` if a connection was accepted, :data:`False` otherwise."""
timer = Timer()
logger.debug("Checking whether %s is accepting connections ..", self.endpoint)
try:
socket.create_connection((self.hostname, self.port_number), self.connect_timeout)
logger.debug("Yes %s is accepting connections (took %s).", self.endpoint, timer)
return True
except Exception:
logger.debug("No %s isn't accepting connections (took %s).", self.endpoint, timer)
return False
@required_property
def port_number(self):
"""The port number to connect to (an integer)."""
@mutable_property
def scheme(self):
"""A URL scheme that indicates the purpose of the ephemeral port (a string, defaults to 'tcp')."""
return 'tcp'
@mutable_property
def wait_timeout(self):
"""The timeout in seconds for :func:`wait_until_connected()` (a number, defaults to 30)."""
return 30
def wait_until_connected(self):
"""
Wait until connections are being accepted.
:raises: :exc:`TimeoutError` when the SSH server isn't fast enough to
initialize.
"""
timer = Timer()
with Spinner(timer=timer) as spinner:
while not self.is_connected:
if timer.elapsed_time > self.wait_timeout:
raise TimeoutError(format(
"Failed to establish connection to %s within configured timeout of %s!",
self.endpoint, format_timespan(self.wait_timeout),
))
spinner.step(label="Waiting for %s to accept connections" % self.endpoint)
spinner.sleep()
logger.debug("Waited %s for %s to accept connections.", timer, self.endpoint)
class EphemeralPortAllocator(WaitUntilConnected):
"""
Allocate a free `ephemeral port number`_.
.. _ephemeral port number: \
http://en.wikipedia.org/wiki/List_of_TCP_and_UDP_port_numbers#Dynamic.2C_private_or_ephemeral_ports
"""
@lazy_property
def port_number(self):
"""A dynamically selected free ephemeral port number (an integer between 49152 and 65535)."""
timer = Timer()
logger.debug("Looking for free ephemeral port number ..")
for i in itertools.count(1):
value = self.ephemeral_port_number
set_property(self, 'port_number', value)
if not self.is_connected:
logger.debug("Found free ephemeral port number %s after %s (took %s).",
value, pluralize(i, "attempt"), timer)
return value
@property
def ephemeral_port_number(self):
"""A random ephemeral port number (an integer between 49152 and 65535)."""
return random.randint(49152, 65535)
class EphemeralTCPServer(ExternalCommand, EphemeralPortAllocator):
"""
Make it easy to launch ephemeral TCP servers.
The :class:`EphemeralTCPServer` class makes it easy to allocate an
`ephemeral port number`_ that is not (yet) in use.
.. _ephemeral port number: \
http://en.wikipedia.org/wiki/List_of_TCP_and_UDP_port_numbers#Dynamic.2C_private_or_ephemeral_ports
"""
@property
def asynchronous(self):
"""Ephemeral TCP servers always set :attr:`.ExternalCommand.asynchronous` to :data:`True`."""
return True
def start(self, **options):
"""
Start the TCP server and wait for it to start accepting connections.
:param options: Any keyword arguments are passed to the
:func:`~executor.ExternalCommand.start()` method of the
superclass.
:raises: Any exceptions raised by :func:`~executor.ExternalCommand.start()`
and :func:`~executor.tcp.WaitUntilConnected.wait_until_connected()`.
If the TCP server doesn't start accepting connections within the
configured timeout (see :attr:`~executor.tcp.WaitUntilConnected.wait_timeout`)
the process will be terminated and the timeout exception is propagated.
"""
if not self.was_started:
logger.debug("Preparing to start %s server ..", self.scheme.upper())
super(EphemeralTCPServer, self).start(**options)
try:
self.wait_until_connected()
except TimeoutError:
self.terminate()
raise
class TimeoutError(Exception):
"""
Raised when a TCP server doesn't start accepting connections quickly enough.
This exception is raised by :func:`~executor.tcp.WaitUntilConnected.wait_until_connected()`
when the TCP server doesn't start accepting connections within a reasonable time.
"""