respect geofence, move from socket to async/await

This commit is contained in:
2026-02-13 18:55:07 -08:00
parent 61e440b594
commit d25287cdf9
2 changed files with 52 additions and 45 deletions

View File

@@ -22,7 +22,6 @@ from enum import IntEnum
from pathlib import Path
import asyncio
import os
import socket
import struct
import yaml
@@ -65,26 +64,22 @@ def get_environment():
return env
def recv_exactly(sock: socket.socket, n: int) -> bytes:
"""Receive exactly n bytes from socket."""
data = b''
while len(data) < n:
chunk = sock.recv(n - len(data))
if not chunk:
raise ConnectionError("Connection closed while receiving data")
data += chunk
async def recv_exactly(reader: asyncio.StreamReader, n: int) -> bytes:
"""Receive exactly n bytes from async stream."""
data = await reader.readexactly(n)
return data
def recv_message_type(sock: socket.socket) -> MessageType:
async def recv_message_type(reader: asyncio.StreamReader) -> MessageType:
"""Receive a single-byte message type."""
data = recv_exactly(sock, 1)
data = await recv_exactly(reader, 1)
return MessageType(data[0])
def send_message_type(sock: socket.socket, msg_type: MessageType):
async def send_message_type(writer: asyncio.StreamWriter, msg_type: MessageType):
"""Send a single-byte message type."""
sock.sendall(bytes([msg_type]))
writer.write(bytes([msg_type]))
await writer.drain()
class UAVRunner(BasicRunner):
@@ -116,23 +111,21 @@ class UAVRunner(BasicRunner):
print(f"[UAV] Connecting to controller at {self.server_ip}:{self.server_port}")
# Retry connection up to 10 times (~30 seconds total)
sock = None
reader, writer = None, None
for attempt in range(10):
try:
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.settimeout(5)
sock.connect((self.server_ip, self.server_port))
sock.settimeout(None)
reader, writer = await asyncio.wait_for(
asyncio.open_connection(self.server_ip, self.server_port),
timeout=5,
)
print(f"[UAV] Connected to controller")
break
except (ConnectionRefusedError, socket.timeout, OSError) as e:
if sock:
sock.close()
except (ConnectionRefusedError, asyncio.TimeoutError, OSError) as e:
print(f"[UAV] Connection attempt {attempt + 1}/10 failed: {e}")
if attempt < 9:
await asyncio.sleep(3)
if sock is None or sock.fileno() == -1:
if reader is None:
print("[UAV] Failed to connect to controller after 10 attempts")
return
@@ -145,12 +138,12 @@ class UAVRunner(BasicRunner):
# Command loop - handle TARGET, RTL, LAND, READY from controller
waypoint_num = 0
while True:
msg_type = recv_message_type(sock)
msg_type = await recv_message_type(reader)
print(f"[UAV] Received: {msg_type.name}")
if msg_type == MessageType.TARGET:
# Read 24 bytes of coordinates (3 little-endian doubles)
data = recv_exactly(sock, 24)
data = await recv_exactly(reader, 24)
enu_x, enu_y, enu_z = struct.unpack('<ddd', data)
waypoint_num += 1
print(f"[UAV] TARGET (waypoint {waypoint_num}): x={enu_x}, y={enu_y}, z={enu_z}")
@@ -159,18 +152,18 @@ class UAVRunner(BasicRunner):
target = self.origin + VectorNED(north=enu_y, east=enu_x, down=-enu_z)
print(f"[UAV] Target coord: {target.lat:.6f}, {target.lon:.6f}, {target.alt:.1f}")
send_message_type(sock, MessageType.ACK)
await send_message_type(writer, MessageType.ACK)
print(f"[UAV] Sent ACK")
print(f"[UAV] Moving to waypoint {waypoint_num}...")
await drone.goto_coordinates(target)
print(f"[UAV] Arrived at waypoint {waypoint_num}")
send_message_type(sock, MessageType.READY)
await send_message_type(writer, MessageType.READY)
print(f"[UAV] Sent READY")
elif msg_type == MessageType.RTL:
send_message_type(sock, MessageType.ACK)
await send_message_type(writer, MessageType.ACK)
print(f"[UAV] Sent ACK")
print("[UAV] Returning to home...")
home = drone.home_coords
@@ -179,11 +172,11 @@ class UAVRunner(BasicRunner):
print(f"[UAV] RTL to {home.lat:.6f}, {home.lon:.6f} at {safe_alt:.1f}m")
await drone.goto_coordinates(rtl_target)
print("[UAV] Arrived at home position")
send_message_type(sock, MessageType.READY)
await send_message_type(writer, MessageType.READY)
print(f"[UAV] Sent READY")
elif msg_type == MessageType.LAND:
send_message_type(sock, MessageType.ACK)
await send_message_type(writer, MessageType.ACK)
print(f"[UAV] Sent ACK")
print("[UAV] Landing...")
await drone.land()
@@ -191,7 +184,7 @@ class UAVRunner(BasicRunner):
from dronekit import VehicleMode
drone._vehicle.mode = VehicleMode("ALT_HOLD")
print("[UAV] Landed and disarmed (ALT_HOLD)")
send_message_type(sock, MessageType.READY)
await send_message_type(writer, MessageType.READY)
print(f"[UAV] Sent READY")
elif msg_type == MessageType.READY:
@@ -201,9 +194,10 @@ class UAVRunner(BasicRunner):
else:
print(f"[UAV] Unknown command: {msg_type}")
except (ValueError, ConnectionError) as e:
except (ValueError, asyncio.IncompleteReadError, ConnectionError) as e:
print(f"[UAV] Error: {e}")
finally:
sock.close()
print("[UAV] Socket closed")
writer.close()
await writer.wait_closed()
print("[UAV] Connection closed")

View File

@@ -3,16 +3,29 @@
# Target waypoints in ENU coordinates (meters: east, north, up)
# Grouped by client: first N entries = Client 1, next N = Client 2, etc.
# Each client must have the same number of waypoints.
#
# Flight pattern:
# UAV 1 (45m alt): straight line along x-axis, 4m intervals
# UAV 2 (30m alt): triangular path, diverges then converges
#
# WP1: x=0 (close, ~10m apart)
# WP2: x=4 (medium, ~20m apart)
# WP3: x=8 (far, ~25m apart)
# WP4: x=12 (medium, ~20m apart)
# WP5: x=16 (close, ~10m apart)
#
# Max distance from origin: ~22m (all waypoints within geofence)
#
targets:
# Client 1 waypoints
- [10.5, 20.3, 45.0]
- [30.0, 40.0, 45.0]
- [50.0, 60.0, 45.0]
- [70.0, 80.0, 45.0]
- [90.0, 100.0, 45.0]
# Client 2 waypoints
- [25.0, 30.0, 30.0]
- [35.0, 40.0, 30.0]
- [45.0, 50.0, 30.0]
- [55.0, 60.0, 30.0]
- [65.0, 70.0, 30.0]
# UAV 1: straight line along x-axis at y=5, alt=45m
- [0.0, 5.0, 45.0]
- [4.0, 5.0, 45.0]
- [8.0, 5.0, 45.0]
- [12.0, 5.0, 45.0]
- [16.0, 5.0, 45.0]
# UAV 2: triangular path diverging/converging, alt=30m
- [0.0, -5.0, 30.0]
- [4.0, -15.0, 30.0]
- [8.0, -20.0, 30.0]
- [12.0, -15.0, 30.0]
- [16.0, -5.0, 30.0]