v2 avoid message ordering issues

Author: blast-hardcheese

src/replit_river/v2/session.py

@@ -331,6 +331,19 @@ async def _send_message(
             service_name,
             procedure_name,
         )
+        # Ensure the buffer isn't full before we enqueue
+        await self._space_available.wait()
+
+        # Before we append, do an important check
+        if self._state in TerminalStates:
+            # session is closing / closed, raise
+            raise SessionClosedRiverServiceException(
+                "river session is closed, dropping message",
+                service_name,
+                procedure_name,
+            )
+
+        # Begin critical section: Avoid any await between here and _send_buffer.append
         msg = TransportMessage(
             streamId=stream_id,
             id=nanoid.generate(),
@@ -348,18 +361,6 @@ async def _send_message(
             with use_span(span):
                 trace_propagator.inject(msg, None, trace_setter)
 
-        # Ensure the buffer isn't full before we enqueue
-        await self._space_available.wait()
-
-        # Before we append, do an important check
-        if self._state in TerminalStates:
-            # session is closing / closed, raise
-            raise SessionClosedRiverServiceException(
-                "river session is closed, dropping message",
-                service_name,
-                procedure_name,
-            )
-
         self._send_buffer.append(msg)
 
         # If the buffer is now full, reset the block