import asyncio import atexit import hashlib import os import signal import sys import time from dataclasses import dataclass, field from threading import Thread from typing import Optional, Dict, Union, AsyncGenerator import tqdm.asyncio from OCP.TopoDS import TopoDS_Shape from aiohttp import web from dataclasses_json import dataclass_json, config from tqdm.contrib.logging import logging_redirect_tqdm from glbs import glb_sequence_to_glbs from mylogger import logger from pubsub import BufferedPubSub from tessellate import _hashcode, tessellate_count, tessellate FRONTEND_BASE_PATH = os.getenv('FRONTEND_BASE_PATH', '../dist') UPDATES_API_PATH = '/api/updates' OBJECTS_API_PATH = '/api/object' # /{name} @dataclass_json @dataclass class UpdatesApiData: """Data sent to the client through the updates API""" name: str """Name of the object. Should be unique unless you want to overwrite the previous object""" hash: str """Hash of the object, to detect changes without rebuilding the object""" obj: Optional[TopoDS_Shape] = field(default=None, metadata=config(exclude=lambda obj: True)) """The OCCT object, if any (not serialized)""" # noinspection PyUnusedLocal async def _index_handler(request: web.Request) -> web.Response: return web.HTTPTemporaryRedirect(location='index.html') class Server: app = web.Application() runner: web.AppRunner thread: Optional[Thread] = None do_shutdown = asyncio.Event() show_events = BufferedPubSub[UpdatesApiData]() object_events: Dict[str, BufferedPubSub[bytes]] = {} object_events_lock = asyncio.Lock() def __init__(self, *args, **kwargs): # --- Routes --- # - APIs self.app.router.add_route('GET', f'{UPDATES_API_PATH}', self._api_updates) self.app.router.add_route('GET', f'{OBJECTS_API_PATH}/{{name}}', self._api_object) # - Static files from the frontend self.app.router.add_get('/{path:(.*/|)}', _index_handler) # Any folder -> index.html self.app.router.add_static('/', path=FRONTEND_BASE_PATH, name='static_frontend') # --- Misc --- self.loop = asyncio.new_event_loop() def start(self): """Starts the web server in the background""" assert self.thread is None, "Server already started" # Start the server in a separate daemon thread self.thread = Thread(target=self._run_server, name='yacv_server', daemon=True) signal.signal(signal.SIGINT | signal.SIGTERM, self.stop) atexit.register(self.stop) self.thread.start() # noinspection PyUnusedLocal def stop(self, *args): """Stops the web server""" if self.thread is None: print('Cannot stop server because it is not running') return # FIXME: Wait for at least one client to confirm ready before stopping in case we are too fast? self.loop.call_soon_threadsafe(lambda *a: self.do_shutdown.set()) self.thread.join(timeout=12) self.thread = None if len(args) >= 1 and args[0] in (signal.SIGINT, signal.SIGTERM): sys.exit(0) # Exit with success def _run_server(self): """Runs the web server""" asyncio.set_event_loop(self.loop) self.loop.run_until_complete(self._run_server_async()) self.loop.stop() self.loop.close() async def _run_server_async(self): """Runs the web server (async)""" runner = web.AppRunner(self.app) await runner.setup() site = web.TCPSite(runner, os.getenv('YACV_HOST', 'localhost'), int(os.getenv('YACV_PORT', 32323))) await site.start() # print(f'Server started at {site.name}') # Wait for a signal to stop the server while running await self.do_shutdown.wait() # print('Shutting down server...') await runner.cleanup() async def _api_updates(self, request: web.Request) -> web.WebSocketResponse: """Handles a publish-only websocket connection that send show_object events along with their hashes and URLs""" ws = web.WebSocketResponse() await ws.prepare(request) async def _send_api_updates(): subscription = self.show_events.subscribe() try: async for data in subscription: # noinspection PyUnresolvedReferences await ws.send_str(data.to_json()) finally: await subscription.aclose() # Start sending updates to the client automatically send_task = asyncio.create_task(_send_api_updates()) try: logger.debug('Client connected: %s', request.remote) # Wait for the client to close the connection (or send a message) await ws.receive() finally: # Make sure to stop sending updates to the client and close the connection send_task.cancel() await ws.close() logger.debug('Client disconnected: %s', request.remote) return ws obj_counter = 0 def _show_common(self, name: Optional[str], hash: str, start: float, obj: Optional[TopoDS_Shape] = None): name = name or f'object_{self.obj_counter}' self.obj_counter += 1 precomputed_info = UpdatesApiData(name=name, hash=hash, obj=obj) self.show_events.publish_nowait(precomputed_info) logger.info('show_object(%s, %s) took %.3f seconds', name, hash, time.time() - start) return precomputed_info def show_gltf(self, gltf: bytes, name: Optional[str] = None, **kwargs): """Publishes any single-file GLTF object to the server (GLB format recommended).""" start = time.time() # Precompute the info and send it to the client as if it was a CAD object precomputed_info = self._show_common(hashlib.md5(gltf).hexdigest(), name, start) # Also pre-populate the GLTF data for the object API publish_to = BufferedPubSub[bytes]() publish_to.publish_nowait(gltf) self.object_events[precomputed_info.name] = publish_to def show_object(self, obj: Union[TopoDS_Shape, any], name: Optional[str] = None, **kwargs): """Publishes a CAD object to the server""" start = time.time() # Try to grab a shape if a different type of object was passed if not isinstance(obj, TopoDS_Shape): # Build123D if 'part' in dir(obj): obj = obj.part if 'sketch' in dir(obj): obj = obj.sketch if 'line' in dir(obj): obj = obj.line # Build123D & CadQuery while 'wrapped' in dir(obj) and not isinstance(obj, TopoDS_Shape): obj = obj.wrapped if not isinstance(obj, TopoDS_Shape): raise ValueError(f'Cannot show object of type {type(obj)} (submit issue?)') self._show_common(name, _hashcode(obj), start, obj) async def _api_object(self, request: web.Request) -> web.StreamResponse: """Returns the object file with the matching name, building it if necessary.""" # Start exporting the object (or fail if not found) export_data = self.export(request.match_info['name']) response = web.StreamResponse() try: # First exported element is the object itself, grab it to collect data export_obj = await anext(export_data) # Create a new stream response with custom content type and headers response.content_type = 'model/gltf-binary-sequence' response.headers['Content-Disposition'] = f'attachment; filename="{request.match_info["name"]}.glbs"' total_parts = 1 if export_obj is None else tessellate_count(export_obj) response.headers['X-Object-Parts'] = str(total_parts) await response.prepare(request) # Convert the GLB sequence to a GLBS sequence and write it to the response with logging_redirect_tqdm(tqdm_class=tqdm.asyncio.tqdm): # noinspection PyTypeChecker glb_parts = tqdm.asyncio.tqdm(export_data, total=total_parts) async for chunk in glb_sequence_to_glbs(glb_parts): await response.write(chunk) finally: # Close the export data subscription await export_data.aclose() # Close the response (if not an error) if response.prepared: await response.write_eof() return response async def export(self, name: str) -> AsyncGenerator[Union[TopoDS_Shape, bytes], None]: """Export the given previously-shown object to a sequence of GLB files, building it if necessary.""" start = time.time() # Check that the object to build exists and grab it if it does subscription = self.show_events.subscribe(include_future=False) obj: Optional[TopoDS_Shape] = None found = False async for data in subscription: if data.name == name: obj = data.obj found = True # Required because obj could be None break await subscription.aclose() if not found: raise web.HTTPNotFound(text=f'No object named {name} was previously shown') # First published element is the TopoDS_Shape, which is None for glTF objects yield obj # Use the lock to ensure that we don't build the object twice async with self.object_events_lock: # If there are no object events for this name, we need to build the object if name not in self.object_events: # Prepare the pubsub for the object publish_to = BufferedPubSub[bytes]() self.object_events[name] = publish_to def _build_object(): # Build the object part_count = 0 for tessellation_update in tessellate(obj): # We publish the object parts as soon as we have a new tessellation list_of_bytes = tessellation_update.gltf.save_to_bytes() publish_to.publish_nowait(b''.join(list_of_bytes)) part_count += 1 publish_to.publish_nowait(b'') # Signal the end of the stream logger.info('export(%s) took %.3f seconds, %d parts', name, time.time() - start, part_count) # We should build it fully even if we are cancelled, so we use a separate task # Furthermore, building is CPU-bound, so we use the default executor asyncio.get_running_loop().run_in_executor(None, _build_object) # In either case return the elements of a subscription to the async generator subscription = self.object_events[name].subscribe() try: async for chunk in subscription: if chunk == b'': break yield chunk finally: await subscription.aclose()