Skip to content

Commit

Permalink
implement spark USB discovery
Browse files Browse the repository at this point in the history
  • Loading branch information
steersbob committed Jun 13, 2024
1 parent d9df9d8 commit e35be81
Show file tree
Hide file tree
Showing 21 changed files with 209 additions and 372 deletions.
4 changes: 2 additions & 2 deletions .github/workflows/build.yml
Original file line number Diff line number Diff line change
Expand Up @@ -8,11 +8,11 @@ on:
workflow_dispatch: {}

env:
DOCKER_IMAGE: ghcr.io/you/your-package
DOCKER_IMAGE: ghcr.io/BrewBlox/brewblox-usb-proxy

jobs:
build:
if: false
if: github.repository_owner == 'BrewBlox'
runs-on: ubuntu-22.04

steps:
Expand Down
6 changes: 6 additions & 0 deletions Dockerfile
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,12 @@ COPY ./entrypoint.sh ./entrypoint.sh
RUN <<EOF
set -ex

apt-get update
apt-get install -y --no-install-recommends \
socat \
usbutils
rm -rf /var/cache/apt/archives /var/lib/apt/lists

python3 -m venv $VENV
pip3 install --no-index your_package
pip3 freeze
Expand Down
File renamed without changes.
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,7 @@

from fastapi import FastAPI

from . import (http_example, mqtt, mqtt_publish_example,
mqtt_subscribe_example, utils)
from . import connected_api, discovery, utils

LOGGER = logging.getLogger(__name__)

Expand Down Expand Up @@ -43,8 +42,7 @@ async def lifespan(app: FastAPI):
# With an AsyncExitStack, we can combine multiple context managers
# without having to increase indentation
async with AsyncExitStack() as stack:
await stack.enter_async_context(mqtt.lifespan())
await stack.enter_async_context(mqtt_publish_example.lifespan())
await stack.enter_async_context(discovery.lifespan())
yield


Expand All @@ -59,8 +57,7 @@ def create_app() -> FastAPI:
setup_logging(config.debug)

# Call setup functions for modules
mqtt.setup()
mqtt_subscribe_example.setup()
discovery.setup()

# Create app
# OpenApi endpoints are set to /api/doc for backwards compatibility
Expand All @@ -70,7 +67,6 @@ def create_app() -> FastAPI:
redoc_url=f'{prefix}/api/redoc',
openapi_url=f'{prefix}/openapi.json')

# Include all endpoints declared by modules
app.include_router(http_example.router, prefix=prefix)
app.include_router(connected_api.router, prefix=prefix)

return app
17 changes: 17 additions & 0 deletions brewblox_usb_proxy/connected_api.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
import logging

from fastapi import APIRouter

from . import discovery

LOGGER = logging.getLogger(__name__)

router = APIRouter(prefix='/connected', tags=['Connected'])


@router.get('/spark')
async def discovery_connected() -> dict[str, int]:
"""
Get device ID and port for all connected Spark USB devices
"""
return discovery.CV.get().connection_index
117 changes: 117 additions & 0 deletions brewblox_usb_proxy/discovery.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,117 @@
import asyncio
import logging
import os
import signal
from asyncio.subprocess import Process
from contextlib import asynccontextmanager, suppress
from contextvars import ContextVar
from dataclasses import dataclass

from serial.tools import list_ports

from . import utils

SPARK_HWIDS = [
r'USB VID\:PID=2B04\:C006.*', # Photon
r'USB VID\:PID=2B04\:C008.*', # P1
]

# Construct a regex OR'ing all allowed hardware ID matches
# Example result: (?:HWID_REGEX_ONE|HWID_REGEX_TWO)
SPARK_DEVICE_REGEX = f'(?:{"|".join([dev for dev in SPARK_HWIDS])})'
USB_BAUD_RATE = 115200


LOGGER = logging.getLogger(__name__)
CV: ContextVar['SparkDiscovery'] = ContextVar('discovery.SparkDiscovery')


@dataclass(frozen=True)
class SparkConnection:
device_id: str
port: int
proc: Process

def close(self):
with suppress(Exception):
os.killpg(os.getpgid(self.proc.pid), signal.SIGINT)


class SparkDiscovery:
def __init__(self) -> None:
self.config = utils.get_config()
self._connections: set[SparkConnection] = set()

@property
def connection_index(self) -> dict[str, int]:
return dict((c.device_id, c.port)
for c in self._connections
if c.proc.returncode is None)

def _next_port(self):
ports = [c.port for c in self._connections]
return next((p for p
in range(self.config.port_start, self.config.port_end)
if p not in ports))

async def connect_usb(self, device_id: str, device_serial: str) -> SparkConnection:
port = self._next_port()
proc = await asyncio.create_subprocess_exec('/usr/bin/socat',
f'tcp-listen:{port},reuseaddr,fork',
f'file:{device_serial},raw,echo=0,b{USB_BAUD_RATE}',
preexec_fn=os.setsid,
shell=False)
return SparkConnection(device_id=device_id,
port=port,
proc=proc)

async def discover_usb(self) -> None:
connected_ids = set(c.device_id for c in self._connections)
discovered_ids = set()

for usb_port in list_ports.grep(SPARK_DEVICE_REGEX):
device_id = usb_port.serial_number.lower()
discovered_ids.add(device_id)
if device_id not in connected_ids:
conn = await self.connect_usb(device_id, usb_port.device)
LOGGER.info(f'Added {conn}')
self._connections.add(conn)

def test_connected(conn: SparkConnection) -> bool:
if conn.proc.returncode is not None:
LOGGER.info(f'Removed {conn} (returncode={conn.proc.returncode})')
return False

if conn.device_id not in discovered_ids:
conn.close()

return True

LOGGER.debug(f'{discovered_ids=}')
self._connections = set(filter(test_connected, self._connections))

async def repeat(self):
interval = self.config.discovery_interval
while True:
try:
await self.discover_usb()
await asyncio.sleep(interval.total_seconds())
except Exception as ex:
LOGGER.error(utils.strex(ex))
await asyncio.sleep(interval.total_seconds())

def close(self):
for conn in self._connections:
conn.close()
self._connections.clear()


@asynccontextmanager
async def lifespan():
async with utils.task_context(CV.get().repeat()):
yield
CV.get().close()


def setup():
CV.set(SparkDiscovery())
30 changes: 30 additions & 0 deletions brewblox_usb_proxy/models.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@
"""
Pydantic models are declared here, and then imported wherever needed
"""

from datetime import timedelta

from pydantic_settings import BaseSettings, SettingsConfigDict


class ServiceConfig(BaseSettings):
"""
Global service configuration.
Pydantic Settings (https://docs.pydantic.dev/latest/concepts/pydantic_settings/)
provides the `BaseSettings` model that loads values from environment variables.
To access the loaded model, we use `utils.get_config()`.
"""
model_config = SettingsConfigDict(
env_prefix='brewblox_usb_proxy_',
case_sensitive=False,
json_schema_extra='ignore',
)

name: str = 'usb-proxy'
debug: bool = False

discovery_interval: timedelta = timedelta(seconds=5)
port_start: int = 9000
port_end: int = 9500
7 changes: 7 additions & 0 deletions your_package/utils.py → brewblox_usb_proxy/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,13 @@ def get_config() -> ServiceConfig: # pragma: no cover
return ServiceConfig()


def strex(ex: Exception):
"""
Formats exception as `Exception(message)`
"""
return f'{type(ex).__name__}({str(ex)})'


@asynccontextmanager
async def task_context(coro: Coroutine,
cancel_timeout=timedelta(seconds=5)
Expand Down
2 changes: 1 addition & 1 deletion entrypoint.sh
Original file line number Diff line number Diff line change
Expand Up @@ -5,4 +5,4 @@ exec uvicorn \
--host 0.0.0.0 \
--port 5000 \
--factory \
your_package.app_factory:create_app
brewblox_usb_proxy.app_factory:create_app
50 changes: 16 additions & 34 deletions poetry.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

8 changes: 4 additions & 4 deletions pyproject.toml
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
[tool.poetry]
name = "your_package"
name = "brewblox_usb_proxy"
version = "0.1.0"
description = "ADD DESCRIPTION HERE"
authors = ["Code Fairies Inc <noreply@santa.com>"]
description = "USB communication proxy for Brewblox devices"
authors = ["BrewPi B.V. <development@brewpi.com>"]
license = "GPL-3.0"
readme = "README.md"

Expand All @@ -11,7 +11,7 @@ python = ">=3.11,<4"
fastapi = "^0.109.2"
uvicorn = { extras = ["standard"], version = "^0.27.1" }
pydantic-settings = "^2.1.0"
fastapi-mqtt = "^2.1.0"
pyserial = "^3.5"

[tool.poetry.group.dev.dependencies]
pytest-mock = "*"
Expand Down
2 changes: 1 addition & 1 deletion tasks.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
from invoke import Context, task

ROOT = Path(__file__).parent.resolve()
IMAGE = 'ghcr.io/you/your-package'
IMAGE = 'ghcr.io/BrewBlox/brewblox-usb-proxy'


@task
Expand Down
4 changes: 2 additions & 2 deletions test/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,8 @@
from httpx import AsyncClient
from pydantic_settings import BaseSettings, PydanticBaseSettingsSource

from your_package import app_factory, utils
from your_package.models import ServiceConfig
from brewblox_usb_proxy import app_factory, utils
from brewblox_usb_proxy.models import ServiceConfig

LOGGER = logging.getLogger(__name__)

Expand Down
2 changes: 2 additions & 0 deletions test/test_discovery.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
def test_dummy():
pass
Loading

0 comments on commit e35be81

Please sign in to comment.