Skip to content

Room

Source code in pycrdt_websocket/yroom.py 
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
class YRoom:
    clients: set[Websocket]
    ydoc: Doc
    ystore: BaseYStore | None
    ready_event: Event
    _on_message: Callable[[bytes], Awaitable[bool] | bool] | None
    _update_send_stream: MemoryObjectSendStream
    _update_receive_stream: MemoryObjectReceiveStream
    _task_group: TaskGroup | None = None
    _started: Event | None = None
    _stopped: Event
    __start_lock: Lock | None = None
    _subscription: Subscription | None = None

    def __init__(
        self,
        ready: bool = True,
        ystore: BaseYStore | None = None,
        exception_handler: Callable[[Exception, Logger], bool] | None = None,
        log: Logger | None = None,
        ydoc: Doc | None = None,
    ):
        """Initialize the object.

        The YRoom instance should preferably be used as an async context manager:
        ```py
        async with room:
            ...
        ```
        However, a lower-level API can also be used:
        ```py
        task = asyncio.create_task(room.start())
        await room.started.wait()
        ...
        await room.stop()
        ```

        Arguments:
            ready: Whether the internal YDoc is ready to be synchronized right away.
            ystore: An optional store in which to persist document updates.
            exception_handler: An optional callback to call when an exception is raised, that
                returns True if the exception was handled.
            log: An optional logger.
            ydoc: An optional document for the room (a new one is created otherwise).
        """
        self.ydoc = Doc() if ydoc is None else ydoc
        self.ready_event = Event()
        self.ready = ready
        self.ystore = ystore
        self.log = log or getLogger(__name__)
        self.awareness = Awareness(self.ydoc)
        self.awareness.observe(self.send_server_awareness)
        self.clients = set()
        self._on_message = None
        self.exception_handler = exception_handler
        self._stopped = Event()

    @property
    def _start_lock(self) -> Lock:
        if self.__start_lock is None:
            self.__start_lock = Lock()
        return self.__start_lock

    @property
    def started(self):
        """An async event that is set when the YRoom provider has started."""
        if self._started is None:
            self._started = Event()
        return self._started

    @property
    def ready(self) -> bool:
        """
        Returns:
            True is the internal YDoc is ready to be synchronized.
        """
        return self.ready_event.is_set()

    @ready.setter
    def ready(self, value: bool) -> None:
        """
        Arguments:
            value: True if the internal YDoc is ready to be synchronized, False otherwise."""
        if value and not self.ready_event.is_set():
            self.ready_event.set()

    async def _watch_ready(self):
        await self.ready_event.wait()
        self._subscription = self.ydoc.observe(partial(put_updates, self._update_send_stream))

    @property
    def on_message(self) -> Callable[[bytes], Awaitable[bool] | bool] | None:
        """
        Returns:
            The optional callback to call when a message is received.
        """
        return self._on_message

    @on_message.setter
    def on_message(self, value: Callable[[bytes], Awaitable[bool] | bool] | None):
        """
        Arguments:
            value: An optional callback to call when a message is received.
            If the callback returns True, the message is skipped.
        """
        self._on_message = value

    async def _broadcast_updates(self):
        if self.ystore is not None:
            async with self.ystore.start_lock:
                if not self.ystore.started.is_set():
                    await self._task_group.start(self.ystore.start)

        async with self._update_receive_stream:
            async for update in self._update_receive_stream:
                if self._task_group.cancel_scope.cancel_called:
                    return
                # broadcast internal ydoc's update to all clients, that includes changes from the
                # clients and changes from the backend (out-of-band changes)
                if self.clients:
                    message = create_update_message(update)
                    for client in self.clients:
                        try:
                            self.log.debug(
                                "Sending Y update to client with endpoint: %s", client.path
                            )
                            self._task_group.start_soon(client.send, message)
                        except Exception as exception:
                            self._handle_exception(exception)
                if self.ystore:
                    try:
                        self._task_group.start_soon(self.ystore.write, update)
                        self.log.debug("Writing Y update to YStore")
                    except Exception as exception:
                        self._handle_exception(exception)

    async def __aenter__(self) -> YRoom:
        async with self._start_lock:
            if self._task_group is not None:
                raise RuntimeError("YRoom already running")

            async with AsyncExitStack() as exit_stack:
                self._task_group = await exit_stack.enter_async_context(create_task_group())
                self._exit_stack = exit_stack.pop_all()
                await self._task_group.start(partial(self.start, from_context_manager=True))

        return self

    async def __aexit__(self, exc_type, exc_value, exc_tb):
        await self.stop()
        return await self._exit_stack.__aexit__(exc_type, exc_value, exc_tb)

    def _handle_exception(self, exception: Exception) -> None:
        exception_handled = False
        if self.exception_handler is not None:
            exception_handled = self.exception_handler(exception, self.log)
        if not exception_handled:
            raise exception

    async def start(
        self,
        *,
        task_status: TaskStatus[None] = TASK_STATUS_IGNORED,
        from_context_manager: bool = False,
    ):
        """Start the room.

        Arguments:
            task_status: The status to set when the task has started.
        """
        if from_context_manager:
            task_status.started()
            self.started.set()
            self._update_send_stream, self._update_receive_stream = create_memory_object_stream(
                max_buffer_size=65536
            )
            assert self._task_group is not None
            self._task_group.start_soon(self._stopped.wait)
            self._task_group.start_soon(self._watch_ready)
            self._task_group.start_soon(self._broadcast_updates)
            self._task_group.start_soon(self.awareness.start)
            return

        async with self._start_lock:
            if self._task_group is not None:
                raise RuntimeError("YRoom already running")

            while True:
                try:
                    async with create_task_group() as self._task_group:
                        if not self.started.is_set():
                            task_status.started()
                            self.started.set()
                        self._update_send_stream, self._update_receive_stream = (
                            create_memory_object_stream(max_buffer_size=65536)
                        )
                        self._task_group.start_soon(self._stopped.wait)
                        self._task_group.start_soon(self._watch_ready)
                        self._task_group.start_soon(self._broadcast_updates)
                        self._task_group.start_soon(self.awareness.start)
                    return
                except Exception as exception:
                    await self.awareness.stop()
                    self._handle_exception(exception)

    async def stop(self) -> None:
        """Stop the room."""
        if self._task_group is None:
            raise RuntimeError("YRoom not running")
        self._stopped.set()
        await self.awareness.stop()
        self._task_group.cancel_scope.cancel()
        self._task_group = None
        if self._subscription is not None:
            self.ydoc.unobserve(self._subscription)

    async def serve(self, websocket: Websocket):
        """Serve a client.

        Arguments:
            websocket: The WebSocket through which to serve the client.
        """
        try:
            async with create_task_group() as tg:
                self.clients.add(websocket)
                sync_message = create_sync_message(self.ydoc)
                self.log.debug(
                    "Sending %s message to endpoint: %s",
                    YSyncMessageType.SYNC_STEP1.name,
                    websocket.path,
                )
                await websocket.send(sync_message)
                async for message in websocket:
                    # filter messages (e.g. awareness)
                    skip = False
                    if self.on_message:
                        _skip = self.on_message(message)
                        skip = await _skip if isawaitable(_skip) else _skip
                    if skip:
                        continue
                    message_type = message[0]
                    if message_type == YMessageType.SYNC:
                        # update our internal state in the background
                        # changes to the internal state are then forwarded to all clients
                        # and stored in the YStore (if any)
                        self.log.debug(
                            "Received %s message from endpoint: %s",
                            YSyncMessageType(message[1]).name,
                            websocket.path,
                        )
                        reply = handle_sync_message(message[1:], self.ydoc)
                        if reply is not None:
                            self.log.debug(
                                "Sending %s message to endpoint: %s",
                                YSyncMessageType.SYNC_STEP2.name,
                                websocket.path,
                            )
                            tg.start_soon(websocket.send, reply)
                    elif message_type == YMessageType.AWARENESS:
                        # forward awareness messages from this client to all clients,
                        # including itself, because it's used to keep the connection alive
                        self.log.debug(
                            "Received %s message from endpoint: %s",
                            YMessageType.AWARENESS.name,
                            websocket.path,
                        )
                        for client in self.clients:
                            self.log.debug(
                                "Sending Y awareness from client with endpoint "
                                "%s to client with endpoint: %s",
                                websocket.path,
                                client.path,
                            )
                            tg.start_soon(client.send, message)
                        # apply awareness update to the server's awareness
                        self.awareness.apply_awareness_update(read_message(message[1:]), self)
        except Exception as exception:
            self._handle_exception(exception)
        finally:
            # remove this client
            self.clients.remove(websocket)

    def send_server_awareness(self, type: str, changes: tuple[dict[str, Any], Any]) -> None:
        """
        Callback to broadcast the server awareness to clients.

        Arguments:
            type: The change type.
            changes: The awareness changes.
        """
        if type != "update" or changes[1] != "local":
            return

        if self._task_group is not None:
            updated_clients = [v for value in changes[0].values() for v in value]
            state = self.awareness.encode_awareness_update(updated_clients)
            message = create_awareness_message(state)
            self._task_group.start_soon(self._send_server_awareness, message)
        else:
            self.log.error("Cannot broadcast server awareness: YRoom not started")

    async def _send_server_awareness(self, state: bytes) -> None:
        try:
            async with create_task_group() as tg:
                for client in self.clients:
                    self.log.debug(
                        "Sending awareness from server to client with endpoint: %s",
                        client.path,
                    )
                    tg.start_soon(client.send, state)
        except Exception as e:
            self.log.error("Error while broadcasting awareness changes: %s", e)

on_message: Callable[[bytes], Awaitable[bool] | bool] | None property writable

Returns:

Type Description
Callable[[bytes], Awaitable[bool] | bool] | None

The optional callback to call when a message is received.

ready: bool property writable

Returns:

Type Description
bool

True is the internal YDoc is ready to be synchronized.

started property

An async event that is set when the YRoom provider has started.

__init__(ready=True, ystore=None, exception_handler=None, log=None, ydoc=None)

Initialize the object.

The YRoom instance should preferably be used as an async context manager: 

async with room:
    ...
However, a lower-level API can also be used: 
task = asyncio.create_task(room.start())
await room.started.wait()
...
await room.stop()

Parameters:

Name Type Description Default
ready bool

Whether the internal YDoc is ready to be synchronized right away.

 True
ystore BaseYStore | None

An optional store in which to persist document updates.

 None
exception_handler Callable[[Exception, Logger], bool] | None

An optional callback to call when an exception is raised, that returns True if the exception was handled.

 None
log Logger | None

An optional logger.

 None
ydoc Doc | None

An optional document for the room (a new one is created otherwise).

 None
Source code in pycrdt_websocket/yroom.py 
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
def __init__(
    self,
    ready: bool = True,
    ystore: BaseYStore | None = None,
    exception_handler: Callable[[Exception, Logger], bool] | None = None,
    log: Logger | None = None,
    ydoc: Doc | None = None,
):
    """Initialize the object.

    The YRoom instance should preferably be used as an async context manager:
    ```py
    async with room:
        ...
    ```
    However, a lower-level API can also be used:
    ```py
    task = asyncio.create_task(room.start())
    await room.started.wait()
    ...
    await room.stop()
    ```

    Arguments:
        ready: Whether the internal YDoc is ready to be synchronized right away.
        ystore: An optional store in which to persist document updates.
        exception_handler: An optional callback to call when an exception is raised, that
            returns True if the exception was handled.
        log: An optional logger.
        ydoc: An optional document for the room (a new one is created otherwise).
    """
    self.ydoc = Doc() if ydoc is None else ydoc
    self.ready_event = Event()
    self.ready = ready
    self.ystore = ystore
    self.log = log or getLogger(__name__)
    self.awareness = Awareness(self.ydoc)
    self.awareness.observe(self.send_server_awareness)
    self.clients = set()
    self._on_message = None
    self.exception_handler = exception_handler
    self._stopped = Event()

send_server_awareness(type, changes)

Callback to broadcast the server awareness to clients.

Parameters:

Name Type Description Default
type str

The change type.

 required
changes tuple[dict[str, Any], Any]

The awareness changes.

 required
Source code in pycrdt_websocket/yroom.py 
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
def send_server_awareness(self, type: str, changes: tuple[dict[str, Any], Any]) -> None:
    """
    Callback to broadcast the server awareness to clients.

    Arguments:
        type: The change type.
        changes: The awareness changes.
    """
    if type != "update" or changes[1] != "local":
        return

    if self._task_group is not None:
        updated_clients = [v for value in changes[0].values() for v in value]
        state = self.awareness.encode_awareness_update(updated_clients)
        message = create_awareness_message(state)
        self._task_group.start_soon(self._send_server_awareness, message)
    else:
        self.log.error("Cannot broadcast server awareness: YRoom not started")

serve(websocket) async

Serve a client.

Parameters:

Name Type Description Default
websocket Websocket

The WebSocket through which to serve the client.

 required
Source code in pycrdt_websocket/yroom.py 
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
async def serve(self, websocket: Websocket):
    """Serve a client.

    Arguments:
        websocket: The WebSocket through which to serve the client.
    """
    try:
        async with create_task_group() as tg:
            self.clients.add(websocket)
            sync_message = create_sync_message(self.ydoc)
            self.log.debug(
                "Sending %s message to endpoint: %s",
                YSyncMessageType.SYNC_STEP1.name,
                websocket.path,
            )
            await websocket.send(sync_message)
            async for message in websocket:
                # filter messages (e.g. awareness)
                skip = False
                if self.on_message:
                    _skip = self.on_message(message)
                    skip = await _skip if isawaitable(_skip) else _skip
                if skip:
                    continue
                message_type = message[0]
                if message_type == YMessageType.SYNC:
                    # update our internal state in the background
                    # changes to the internal state are then forwarded to all clients
                    # and stored in the YStore (if any)
                    self.log.debug(
                        "Received %s message from endpoint: %s",
                        YSyncMessageType(message[1]).name,
                        websocket.path,
                    )
                    reply = handle_sync_message(message[1:], self.ydoc)
                    if reply is not None:
                        self.log.debug(
                            "Sending %s message to endpoint: %s",
                            YSyncMessageType.SYNC_STEP2.name,
                            websocket.path,
                        )
                        tg.start_soon(websocket.send, reply)
                elif message_type == YMessageType.AWARENESS:
                    # forward awareness messages from this client to all clients,
                    # including itself, because it's used to keep the connection alive
                    self.log.debug(
                        "Received %s message from endpoint: %s",
                        YMessageType.AWARENESS.name,
                        websocket.path,
                    )
                    for client in self.clients:
                        self.log.debug(
                            "Sending Y awareness from client with endpoint "
                            "%s to client with endpoint: %s",
                            websocket.path,
                            client.path,
                        )
                        tg.start_soon(client.send, message)
                    # apply awareness update to the server's awareness
                    self.awareness.apply_awareness_update(read_message(message[1:]), self)
    except Exception as exception:
        self._handle_exception(exception)
    finally:
        # remove this client
        self.clients.remove(websocket)

start(*, task_status=TASK_STATUS_IGNORED, from_context_manager=False) async

Start the room.

Parameters:

Name Type Description Default
task_status TaskStatus[None]

The status to set when the task has started.

 TASK_STATUS_IGNORED
Source code in pycrdt_websocket/yroom.py 
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
async def start(
    self,
    *,
    task_status: TaskStatus[None] = TASK_STATUS_IGNORED,
    from_context_manager: bool = False,
):
    """Start the room.

    Arguments:
        task_status: The status to set when the task has started.
    """
    if from_context_manager:
        task_status.started()
        self.started.set()
        self._update_send_stream, self._update_receive_stream = create_memory_object_stream(
            max_buffer_size=65536
        )
        assert self._task_group is not None
        self._task_group.start_soon(self._stopped.wait)
        self._task_group.start_soon(self._watch_ready)
        self._task_group.start_soon(self._broadcast_updates)
        self._task_group.start_soon(self.awareness.start)
        return

    async with self._start_lock:
        if self._task_group is not None:
            raise RuntimeError("YRoom already running")

        while True:
            try:
                async with create_task_group() as self._task_group:
                    if not self.started.is_set():
                        task_status.started()
                        self.started.set()
                    self._update_send_stream, self._update_receive_stream = (
                        create_memory_object_stream(max_buffer_size=65536)
                    )
                    self._task_group.start_soon(self._stopped.wait)
                    self._task_group.start_soon(self._watch_ready)
                    self._task_group.start_soon(self._broadcast_updates)
                    self._task_group.start_soon(self.awareness.start)
                return
            except Exception as exception:
                await self.awareness.stop()
                self._handle_exception(exception)

stop() async

Stop the room.

Source code in pycrdt_websocket/yroom.py 
241
242
243
244
245
246
247
248
249
250
async def stop(self) -> None:
    """Stop the room."""
    if self._task_group is None:
        raise RuntimeError("YRoom not running")
    self._stopped.set()
    await self.awareness.stop()
    self._task_group.cancel_scope.cancel()
    self._task_group = None
    if self._subscription is not None:
        self.ydoc.unobserve(self._subscription)