Skip to content

Commit

Permalink
DRIVERS-2949 Test server for happy eyeballs (#544)
Browse files Browse the repository at this point in the history
  • Loading branch information
abr-egn authored Nov 12, 2024
1 parent aeb2d4a commit 5d0c7ad
Show file tree
Hide file tree
Showing 3 changed files with 258 additions and 0 deletions.
57 changes: 57 additions & 0 deletions .evergreen/happy_eyeballs/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@
# Happy Eyeballs test scripts

This folder contains a test server ([`server.py`](server.py)) for driver Happy Eyeballs TCP connection behavior. It also has has [`client.py`](client.py), a simple client for that server that can be useful for debugging but won't be needed in most cases.

**NOTE**: This server relies on network stack behavior present in Windows and MacOS but not Linux, so any tests using it must only be run on those two OSes.

## Command-line Usage

`python3 server.py [-c|--control PORT] [--wait] [--stop]`

Run with:
* no arguments, or with just `-c PORT`, to start the server in the foreground. `PORT` defaults to 10036 if not specified.
* `--wait` to wait for a server running in another process to be ready to start accepting control connections.
* `--stop` to signal a server running in another process to gracefully shutdown.

## Integration with Evergreen

The server can be incorporated into an evergreen test run via these functions:
```yaml
functions:
"start happy eyeballs server":
- command: subprocess.exec
params:
working_dir: src
background: true
binary: ${PYTHON3}
args:
- ${DRIVERS_TOOLS}/.evergreen/happy_eyeballs/server.py
- command: subprocess.exec
params:
working_dir: src
binary: ${PYTHON3}
args:
- ${DRIVERS_TOOLS}/.evergreen/happy_eyeballs/server.py
- --wait
"stop happy eyeballs server":
- command: subprocess.exec
params:
working_dir: src
binary: ${PYTHON3}
args:
- ${DRIVERS_TOOLS}/.evergreen/happy_eyeballs/server.py
- --stop
```
The `"stop happy eyeballs server"` function should be included in the `post` configuration or a `teardown_task` section to ensure that the server isn't left running after the test finishes.

## Test Usage

On opening a connection to the control port, the driver test should send a single byte: `0x04` to request a port pair with a slow IPv4 connection, or `0x06` to request one with a slow IPv6 connection. The server will respond with:
1. `0x01` (success signal), followed by
2. `uint16` (IPv4 port, big-endian), followed by
3. `uint16` (IPv6 port, big-endian)

Any other response should be treated as an error. The connection will be closed after the ports are sent; to request another pair, open a new connection to the control port.

Test connections to the two ports should be initiated immediately; the TCP handshake completion will be delayed on the slow port by two seconds from the time of the port _being bound_, not the ACK received. Only one port is expected to successfully connect. Once connected, the port will write out a single byte for verification (`0x04`
for IPv4, `0x06` for IPv6) and then immediately close both ports.
50 changes: 50 additions & 0 deletions .evergreen/happy_eyeballs/client.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,50 @@
### A test client for server.py
#
# This can be used to check that server.py is functioning properly. When run, it
# will connect to the control port on that server, request a pair of ports, open a connection to
# both ports in parallel, and assert that the byte read is the expected one for that port.

import argparse
import asyncio
import socket

parser = argparse.ArgumentParser(
prog='client',
description='client for testing the happy eyeballs test server',
)
parser.add_argument('-c', '--control', default=10036, type=int, metavar='PORT', help='control port')
parser.add_argument('-d', '--delay', default=4, choices=[4,6], type=int, help="ip protocol to request server delay")
args = parser.parse_args()

async def main():
print('connecting to control')
control_r, control_w = await asyncio.open_connection('localhost', args.control)
control_w.write(args.delay.to_bytes(1, 'big'))
await control_w.drain()
data = await control_r.read(1)
if data != b'\x01':
raise Exception(f'Expected byte 1, got {data}')
ipv4_port = int.from_bytes(await control_r.read(2), 'big')
ipv6_port = int.from_bytes(await control_r.read(2), 'big')
connect_tasks = [
asyncio.create_task(connect('IPv4', ipv4_port, socket.AF_INET, b'\x04')),
asyncio.create_task(connect('IPv6', ipv6_port, socket.AF_INET6, b'\x06')),
]
await asyncio.wait(connect_tasks)

async def connect(name: str, port: int, family: socket.AddressFamily, payload: bytes):
print(f'{name}: connecting')
try:
reader, writer = await asyncio.open_connection('localhost', port, family=family)
except Exception as e:
print(f'{name}: failed ({e})')
return
print(f'{name}: connected')
data = await reader.readexactly(1)
if data != payload:
raise Exception(f'Expected {payload}, got {data}')
writer.close()
await writer.wait_closed()
print(f'{name}: done')

asyncio.run(main())
151 changes: 151 additions & 0 deletions .evergreen/happy_eyeballs/server.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,151 @@
# A test server for driver Happy Eyeballs behavior. See README.md for more information.

import argparse
import asyncio
import socket
import sys
import platform

if platform.system() not in ["Darwin", "Windows"]:
print(f'Only macOS (Darwin) and Windows are supported, but got: {platform.system()}', file=sys.stderr)
exit(1)
parser = argparse.ArgumentParser(
prog='server',
description='Test server for happy eyeballs',
)
parser.add_argument('-c', '--control', default=10036, type=int, metavar='PORT', help='control port')
parser.add_argument('--wait', action='store_true', help='wait for a running server to be ready')
parser.add_argument('--stop', action='store_true', help='stop a running server')
args = parser.parse_args()

PREFIX='happy eyeballs server'

async def control_server():
shutdown = asyncio.Event()
srv = await asyncio.start_server(lambda reader, writer: on_control_connected(reader, writer, shutdown), 'localhost', args.control)
print(f'{PREFIX}: listening for control connections on {args.control}', file=sys.stderr)
async with srv:
await shutdown.wait()
print(f'{PREFIX}: all done', file=sys.stderr)

async def on_control_connected(reader: asyncio.StreamReader, writer: asyncio.StreamWriter, shutdown: asyncio.Event):
# Read the control request byte
data = await reader.readexactly(1)
if data == b'\x04':
print(f'{PREFIX}: ========================', file=sys.stderr)
print(f'{PREFIX}: request for delayed IPv4', file=sys.stderr)
slow = 'IPv4'
elif data == b'\x06':
print(f'{PREFIX}: ========================', file=sys.stderr)
print(f'{PREFIX}: request for delayed IPv6', file=sys.stderr)
slow = 'IPv6'
elif data == b'\xF0':
writer.write(b'\x01')
await writer.drain()
writer.close()
await writer.wait_closed()
return
elif data == b'\xFF':
print(f'{PREFIX}: shutting down', file=sys.stderr)
writer.close()
await writer.wait_closed()
shutdown.set()
return
else:
print(f'Unexpected control byte: {data}', file=sys.stderr)
exit(1)

# Bind the test ports but do not yet start accepting connections
connected = asyncio.Event()
on_ipv4_connected = lambda reader, writer: on_test_connected('IPv4', writer, b'\x04', connected, slow)
on_ipv6_connected = lambda reader, writer: on_test_connected('IPv6', writer, b'\x06', connected, slow)
# port 0: pick random unused port
srv4 = await asyncio.start_server(on_ipv4_connected, 'localhost', 0, family=socket.AF_INET, start_serving=False)
srv6 = await asyncio.start_server(on_ipv6_connected, 'localhost', 0, family=socket.AF_INET6, start_serving=False)
ipv4_port = srv4.sockets[0].getsockname()[1]
ipv6_port = srv6.sockets[0].getsockname()[1]
print(f'{PREFIX}: [slow {slow}] bound for IPv4 on {ipv4_port}', file=sys.stderr)
print(f'{PREFIX}: [slow {slow}] bound for IPv6 on {ipv6_port}', file=sys.stderr)

# Reply to control request with success byte and test server ports
writer.write(b'\x01')
writer.write(ipv4_port.to_bytes(2, 'big'))
writer.write(ipv6_port.to_bytes(2, 'big'))
await writer.drain()
writer.close()
await writer.wait_closed()

# Start test servers listening in parallel
# Hold a reference to the tasks so they aren't GC'd
test_tasks = [
asyncio.create_task(test_listen('IPv4', srv4, data == b'\x04', connected, slow)),
asyncio.create_task(test_listen('IPv6', srv6, data == b'\x06', connected, slow)),
]
await asyncio.wait(test_tasks)

# Wait for the test servers to shut down
srv4.close()
srv6.close()
close_tasks = [
asyncio.create_task(srv4.wait_closed()),
asyncio.create_task(srv6.wait_closed()),
]
await asyncio.wait(close_tasks)

print(f'{PREFIX}: [slow {slow}] connection complete, test ports closed', file=sys.stderr)
print(f'{PREFIX}: ========================', file=sys.stderr)

async def test_listen(name: str, srv, delay: bool, connected: asyncio.Event, slow: str):
# Both connections are delayed; the slow one is delayed by more than the fast one; this
# ensures that the client is comparing timing and not simply choosing an immediate success
# over a connection denied.
if delay:
print(f'{PREFIX}: [slow {slow}] delaying {name} connections', file=sys.stderr)
await asyncio.sleep(2.0)
else:
await asyncio.sleep(1.0)
async with srv:
await srv.start_serving()
print(f'{PREFIX}: [slow {slow}] accepting {name} connections', file=sys.stderr)
# Terminate this test server when either test server has handled a request
await connected.wait()

async def on_test_connected(name: str, writer: asyncio.StreamWriter, payload: bytes, connected: asyncio.Event, slow: str):
print(f'{PREFIX}: [slow {slow}] connected on {name}', file=sys.stderr)
writer.write(payload)
await writer.drain()
writer.close()
await writer.wait_closed()
connected.set()

async def stop_server():
control_r, control_w = await asyncio.open_connection('localhost', args.control)
control_w.write(b'\xFF')
await control_w.drain()
control_w.close()
await control_w.wait_closed()

async def wait_for_server():
while True:
try:
control_r, control_w = await asyncio.open_connection('localhost', args.control)
except OSError as e:
print(f'{PREFIX}: failed ({e}), will retry', file=sys.stderr)
await asyncio.sleep(1)
continue
break
control_w.write(b'\xF0')
await control_w.drain()
data = await control_r.read(1)
if data != b'\x01':
print(f'{PREFIX}: expected byte 1, got {data}', file=sys.stderr)
exit(1)
print(f'{PREFIX}: happy eyeballs server ready on port {args.control}', file=sys.stderr)


if args.stop:
asyncio.run(stop_server())
elif args.wait:
asyncio.run(wait_for_server())
else:
asyncio.run(control_server())

0 comments on commit 5d0c7ad

Please sign in to comment.