Stripping WIP logging

Author: blast-hardcheese

src/replit_river/common_session.py

@@ -61,7 +61,6 @@ async def setup_heartbeat(
     send_message: SendMessage,
     increment_and_get_heartbeat_misses: Callable[[], int],
 ) -> None:
-    logger.debug("Start heartbeat")
     while True:
         await asyncio.sleep(heartbeat_ms / 1000)
         state = get_state()
@@ -141,26 +140,24 @@ async def buffered_message_sender(
     get_next_pending: Callable[[], TransportMessage | None],
     commit: Callable[[TransportMessage], None],
 ) -> None:
-    logger.debug("Entering buffered_message_sender")
     while True:
         await message_enqueued.acquire()
         while (ws := get_ws()) is None:
             # Block until we have a handle
-            logger.debug("buffered_message_sender: connection_condition.acquire() %r %r %r", ws, get_ws(), connection_condition)
+            logger.debug(
+                "buffered_message_sender: Waiting until ws is connected (condition=%r)",
+                connection_condition,
+            )
             async with connection_condition:
                 await connection_condition.wait()
-                logger.debug("buffered_message_sender: connection_condition UNLOCKED")
-        logger.debug("buffered_message_sender: acquired")
         if msg := get_next_pending():
             logger.debug(
                 "buffered_message_sender: Dequeued %r to send over %r",
                 msg,
                 ws,
             )
             try:
-                logger.debug("buffered_message_sender: Sending...")
                 await send_transport_message(msg, ws, websocket_closed_callback)
-                logger.debug("buffered_message_sender: committing!")
                 commit(msg)
             except WebsocketClosedException as e:
                 logger.debug(
@@ -179,7 +176,6 @@ async def buffered_message_sender(
             except Exception:
                 logger.exception("Error attempting to send buffered messages")
                 break
-    print("buffered_message_sender exit")
 
 
 async def add_msg_to_stream(

src/replit_river/rate_limiter.py

@@ -91,7 +91,6 @@ def has_budget_or_throw(
         Returns:
             bool: True if budget is available, False otherwise.
         """
-        logger.debug("self.get_budget_consumed(user)=%r < self.options.attempt_budget_capacity=%r", self.get_budget_consumed(user), self.options.attempt_budget_capacity)
         if self.get_budget_consumed(user) > self.options.attempt_budget_capacity:
             logger.debug("No retry budget for %s.", user)
             raise BudgetExhaustedException(

src/replit_river/v2/client_transport.py

@@ -50,7 +50,6 @@ def __init__(
     async def close(self) -> None:
         self._rate_limiter.close()
         if self._session:
-            logger.info(f"start closing session {self._transport_id}")
             await self._session.close()
             logger.info(f"Transport closed {self._transport_id}")
 
@@ -60,7 +59,6 @@ async def get_or_create_session(self) -> Session:
         call ensure_connected on whatever session is active.
         """
         existing_session = self._session
-        logger.debug(f"if not existing_session={existing_session} or existing_session.is_closed()={existing_session and existing_session.is_closed()}:")
         if not existing_session or existing_session.is_closed():
             logger.info("Creating new session")
             new_session = Session(

src/replit_river/v2/session.py

@@ -11,7 +11,6 @@
     Coroutine,
     Literal,
     TypeAlias,
-    cast,
 )
 
 import nanoid  # type: ignore
@@ -217,7 +216,6 @@ def get_next_pending() -> TransportMessage | None:
 
         # TODO: Just return _ws_unwrapped once we are no longer using the legacy client
         def get_ws() -> WebSocketCommonProtocol | ClientConnection | None:
-            logger.debug("get_ws: %r %r", self.is_connected(), self._ws_unwrapped)
             if self.is_connected():
                 return self._ws_unwrapped
             return None
@@ -250,7 +248,7 @@ async def ensure_connected[HandshakeMetadata](
         logic that actually establishes the connection.
         """
 
-        logger.debug("ensure_connected: %r", self.is_connected())
+        logger.debug("ensure_connected: is_connected=%r", self.is_connected())
         if self.is_connected():
             return
 
@@ -264,9 +262,7 @@ async def ensure_connected[HandshakeMetadata](
                 )
             )
 
-        logger.debug("BEFORE await _do_ensure_connected")
         await self._connecting_task
-        logger.debug("AFTER await _do_ensure_connected")
 
     async def _do_ensure_connected[HandshakeMetadata](
         self,
@@ -391,9 +387,7 @@ async def websocket_closed_callback() -> None:
                     rate_limiter.start_restoring_budget(client_id)
                     self._state = SessionState.ACTIVE
                     self._ws_unwrapped = ws
-                    logger.debug("Before notify_all: %r %r %r", self._state, self._ws_unwrapped, self._connection_condition)
                     self._connection_condition.notify_all()
-                    self._connection_condition.release()
                     break
                 except RiverException as e:
                     await ws.close()
@@ -422,13 +416,16 @@ async def websocket_closed_callback() -> None:
         ):
             self._connecting_task = None
 
+        # Release the lock we took earlier so we can use it again in the next
+        # connection attempt
+        self._connection_condition.release()
+
         if last_error is not None:
             raise RiverException(
                 ERROR_HANDSHAKE,
                 f"Failed to create ws after retrying {max_retry} number of times",
             ) from last_error
 
-        logger.debug("EXITING _do_ensure_connected")
         return True
 
     def is_closed(self) -> bool:
@@ -495,7 +492,6 @@ async def send_message(
             serviceName=service_name,
             procedureName=procedure_name,
         )
-        logger.debug("SENDING MESSAGE: %r", msg)
 
         if span:
             with use_span(span):
@@ -514,9 +510,7 @@ async def send_message(
             self._queue_full_lock.locked()
             or len(self._send_buffer) >= self._transport_options.buffer_size
         ):
-            logger.warning("LOCK ACQUIRED %r", repr(payload))
             await self._queue_full_lock.acquire()
-            logger.warning("LOCK RELEASED %r", repr(payload))
         self._send_buffer.append(msg)
         # Wake up buffered_message_sender
         self._message_enqueued.release()