Skip to content

Commit

Permalink
Fixed error on trying to send ws message to disconnected client
Browse files Browse the repository at this point in the history
This could result in a major filament usage track failure since the Use http request would fail and Moonraker would keep retrying to save the used filament. However it was actually saved since the db.commit() happens before the websocket message fails.
  • Loading branch information
Donkie committed Aug 16, 2024
1 parent 5c5bd46 commit e5f7949
Show file tree
Hide file tree
Showing 4 changed files with 32 additions and 12 deletions.
6 changes: 3 additions & 3 deletions spoolman/database/filament.py
Original file line number Diff line number Diff line change
Expand Up @@ -76,8 +76,8 @@ async def create(
extra=[models.FilamentField(key=k, value=v) for k, v in (extra or {}).items()],
)
db.add(filament)
await db.commit()
await filament_changed(filament, EventType.ADDED)
await db.commit()
return filament


Expand Down Expand Up @@ -172,8 +172,8 @@ async def update(
filament.multi_color_direction = v.value if v is not None else None
else:
setattr(filament, k, v)
await db.commit()
await filament_changed(filament, EventType.UPDATED)
await db.commit()
return filament


Expand All @@ -182,8 +182,8 @@ async def delete(db: AsyncSession, filament_id: int) -> None:
filament = await get_by_id(db, filament_id)
await db.delete(filament)
try:
await db.commit() # Flush immediately so any errors are propagated in this request.
await filament_changed(filament, EventType.DELETED)
await db.commit() # Flush immediately so any errors are propagated in this request.
except IntegrityError as exc:
await db.rollback()
raise ItemDeleteError("Failed to delete filament.") from exc
Expand Down
16 changes: 10 additions & 6 deletions spoolman/database/spool.py
Original file line number Diff line number Diff line change
Expand Up @@ -92,8 +92,8 @@ async def create(
extra=[models.SpoolField(key=k, value=v) for k, v in (extra or {}).items()],
)
db.add(spool)
await db.commit()
await spool_changed(spool, EventType.ADDED)
await db.commit()
return spool


Expand Down Expand Up @@ -228,16 +228,16 @@ async def update(
spool.extra = [models.SpoolField(key=k, value=v) for k, v in v.items()]
else:
setattr(spool, k, v)
await db.commit()
await spool_changed(spool, EventType.UPDATED)
await db.commit()
return spool


async def delete(db: AsyncSession, spool_id: int) -> None:
"""Delete a spool object."""
spool = await get_by_id(db, spool_id)
await db.delete(spool)
await spool_changed(spool, EventType.DELETED)
await db.delete(spool)


async def clear_extra_field(db: AsyncSession, key: str) -> None:
Expand Down Expand Up @@ -291,8 +291,8 @@ async def use_weight(db: AsyncSession, spool_id: int, weight: float) -> models.S
spool.first_used = datetime.utcnow().replace(microsecond=0)
spool.last_used = datetime.utcnow().replace(microsecond=0)

await db.commit()
await spool_changed(spool, EventType.UPDATED)
await db.commit()
return spool


Expand Down Expand Up @@ -337,8 +337,12 @@ async def use_length(db: AsyncSession, spool_id: int, length: float) -> models.S
spool.first_used = datetime.utcnow().replace(microsecond=0)
spool.last_used = datetime.utcnow().replace(microsecond=0)

await db.commit()
await spool_changed(spool, EventType.UPDATED)

# Commit should be the last action, everything after that must never fail
# Otherwise you can end up in a non-atomic thing where the http use request fails
# but the data still has been committed.
await db.commit()
return spool


Expand Down Expand Up @@ -449,6 +453,6 @@ async def reset_initial_weight(db: AsyncSession, spool_id: int, weight: float) -

spool.initial_weight = weight
spool.used_weight = 0
await db.commit()
await spool_changed(spool, EventType.UPDATED)
await db.commit()
return spool
4 changes: 2 additions & 2 deletions spoolman/database/vendor.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,8 +33,8 @@ async def create(
extra=[models.VendorField(key=k, value=v) for k, v in (extra or {}).items()],
)
db.add(vendor)
await db.commit()
await vendor_changed(vendor, EventType.ADDED)
await db.commit()
return vendor


Expand Down Expand Up @@ -101,8 +101,8 @@ async def update(
vendor.extra = [models.VendorField(key=k, value=v) for k, v in v.items()]
else:
setattr(vendor, k, v)
await db.commit()
await vendor_changed(vendor, EventType.UPDATED)
await db.commit()
return vendor


Expand Down
18 changes: 17 additions & 1 deletion spoolman/ws.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import logging

from fastapi import WebSocket
from starlette.websockets import WebSocketState

from spoolman.api.v1.models import Event

Expand Down Expand Up @@ -46,7 +47,22 @@ async def send(self, path: tuple[str, ...], evt: Event) -> None:
"""Send a message to all websockets in this branch of the tree."""
# Broadcast to all subscribers on this level
for websocket in self.subscribers:
await websocket.send_text(evt.json())
if (
websocket.client_state == WebSocketState.DISCONNECTED # noqa: PLR1714
or websocket.application_state == WebSocketState.DISCONNECTED
):
# A bad disconnection may have occurred
self.remove(path, websocket)
logger.info(
"Forcing disconnection of client %s on pool %s",
websocket.client.host if websocket.client else "?",
",".join(path),
)
elif (
websocket.client_state == WebSocketState.CONNECTED
and websocket.application_state == WebSocketState.CONNECTED
):
await websocket.send_text(evt.json())

# Send the message further down the tree
if len(path) > 0 and path[0] in self.children:
Expand Down

0 comments on commit e5f7949

Please sign in to comment.