respect geofence, move from socket to async/await

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

View File

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

View File

@@ -3,16 +3,29 @@
# Target waypoints in ENU coordinates (meters: east, north, up) # Target waypoints in ENU coordinates (meters: east, north, up)
# Grouped by client: first N entries = Client 1, next N = Client 2, etc. # Grouped by client: first N entries = Client 1, next N = Client 2, etc.
# Each client must have the same number of waypoints. # 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: targets:
# Client 1 waypoints # UAV 1: straight line along x-axis at y=5, alt=45m
- [10.5, 20.3, 45.0] - [0.0, 5.0, 45.0]
- [30.0, 40.0, 45.0] - [4.0, 5.0, 45.0]
- [50.0, 60.0, 45.0] - [8.0, 5.0, 45.0]
- [70.0, 80.0, 45.0] - [12.0, 5.0, 45.0]
- [90.0, 100.0, 45.0] - [16.0, 5.0, 45.0]
# Client 2 waypoints # UAV 2: triangular path diverging/converging, alt=30m
- [25.0, 30.0, 30.0] - [0.0, -5.0, 30.0]
- [35.0, 40.0, 30.0] - [4.0, -15.0, 30.0]
- [45.0, 50.0, 30.0] - [8.0, -20.0, 30.0]
- [55.0, 60.0, 30.0] - [12.0, -15.0, 30.0]
- [65.0, 70.0, 30.0] - [16.0, -5.0, 30.0]