fix sessionStatus type, remove no connection -> no connection transition on server, extra invariant checks (#296)

## Why

we get a very very low volume of message out of order (~40 messages
across 15m connections) but these still suck

lets see if the message ordering problem is present before we send too

## What changed

add some extra invariant violation checks before message send

fix some small session status type alignment and removed an extra state
where we would emit an extra session transition for the transparent
reconnect case

## Versioning

- [ ] Breaking protocol change
- [ ] Breaking ts/js API change

<!-- Kind reminder to add tests and updated documentation if needed -->

Author: jackyzha0

package-lock.json

@@ -1,7 +1,7 @@
 {
   "name": "@replit/river",
   "description": "It's like tRPC but... with JSON Schema Support, duplex streaming and support for service multiplexing. Transport agnostic!",
-  "version": "0.205.1",
+  "version": "0.205.2",
   "type": "module",
   "exports": {
     ".": {

package.json

@@ -15,14 +15,15 @@ export class InMemoryConnection extends Connection {
   constructor(pipe: Duplex) {
     super();
     this.conn = pipe;
+    this.conn.allowHalfOpen = false;
 
     this.conn.on('data', (data: Uint8Array) => {
       for (const cb of this.dataListeners) {
         cb(data);
       }
     });
 
-    this.conn.on('end', () => {
+    this.conn.on('close', () => {
       for (const cb of this.closeListeners) {
         cb();
       }
@@ -46,6 +47,7 @@ export class InMemoryConnection extends Connection {
   close(): void {
     setImmediate(() => {
       this.conn.end();
+      this.conn.emit('close');
     });
   }
 }
@@ -153,6 +155,7 @@ export function createMockTransportNetwork(
     simulatePhantomDisconnect() {
       for (const conn of Object.values(connections.get())) {
         conn.serverToClient.pause();
+        conn.clientToServer.pause();
       }
     },
     async restartServer() {

testUtil/fixtures/mockTransport.ts

@@ -2,6 +2,7 @@ import { type Static } from '@sinclair/typebox';
 import { Connection } from './connection';
 import { OpaqueTransportMessage, HandshakeErrorResponseCodes } from './message';
 import { Session, SessionState } from './sessionStateMachine';
+import { SessionId } from './sessionStateMachine/common';
 import { TransportStatus } from './transport';
 
 export const ProtocolError = {
@@ -26,11 +27,11 @@ export interface EventMap {
         session: Pick<Session<Connection>, 'id' | 'to'>;
       };
   sessionTransition:
-    | { state: SessionState.Connected }
-    | { state: SessionState.Handshaking }
-    | { state: SessionState.Connecting }
-    | { state: SessionState.BackingOff }
-    | { state: SessionState.NoConnection };
+    | { state: SessionState.Connected; id: SessionId }
+    | { state: SessionState.Handshaking; id: SessionId }
+    | { state: SessionState.Connecting; id: SessionId }
+    | { state: SessionState.BackingOff; id: SessionId }
+    | { state: SessionState.NoConnection; id: SessionId };
   protocolError:
     | {
         type: (typeof ProtocolError)['HandshakeFailed'];

transport/events.ts

@@ -395,9 +395,8 @@ export abstract class ServerTransport<
           );
 
         oldSession = noConnectionSession;
+        this.updateSession(oldSession);
       }
-
-      this.updateSession(oldSession);
     } else if (oldSession) {
       connectCase = 'hard reconnection';
 

transport/server.ts

@@ -4,6 +4,7 @@ import {
   ControlMessageAckSchema,
   OpaqueTransportMessage,
   PartialTransportMessage,
+  TransportMessage,
   isAck,
 } from '../message';
 import {
@@ -48,10 +49,25 @@ export class SessionConnected<
     this.heartbeatMisses = 0;
   }
 
+  private assertSendOrdering(constructedMsg: TransportMessage) {
+    if (constructedMsg.seq > this.seqSent + 1) {
+      const msg = `invariant violation: would have sent out of order msg (seq: ${constructedMsg.seq}, expected: ${this.seqSent} + 1)`;
+      this.log?.error(msg, {
+        ...this.loggingMetadata,
+        transportMessage: constructedMsg,
+        tags: ['invariant-violation'],
+      });
+
+      throw new Error(msg);
+    }
+  }
+
   send(msg: PartialTransportMessage): string {
     const constructedMsg = this.constructMsg(msg);
+    this.assertSendOrdering(constructedMsg);
     this.sendBuffer.push(constructedMsg);
     this.conn.send(this.options.codec.toBuffer(constructedMsg));
+    this.seqSent = constructedMsg.seq;
 
     return constructedMsg.id;
   }
@@ -75,7 +91,9 @@ export class SessionConnected<
       );
 
       for (const msg of this.sendBuffer) {
+        this.assertSendOrdering(msg);
         this.conn.send(this.options.codec.toBuffer(msg));
+        this.seqSent = msg.seq;
       }
     }
 
@@ -165,7 +183,7 @@ export class SessionConnected<
         );
       } else {
         const reason = `received out-of-order msg, closing connection (got seq: ${parsedMsg.seq}, wanted seq: ${this.ack})`;
-        this.log?.warn(reason, {
+        this.log?.error(reason, {
           ...this.loggingMetadata,
           transportMessage: parsedMsg,
           tags: ['invariant-violation'],

transport/sessionStateMachine/SessionConnected.ts

@@ -208,6 +208,7 @@ export interface IdentifiedSessionProps extends CommonSessionProps {
   to: TransportClientId;
   seq: number;
   ack: number;
+  seqSent: number;
   sendBuffer: Array<OpaqueTransportMessage>;
   telemetry: TelemetryInfo;
   protocolVersion: ProtocolVersion;
@@ -224,15 +225,29 @@ export abstract class IdentifiedSession extends CommonSession {
    */
   seq: number;
 
+  /**
+   * Last seq we sent over the wire this session (excluding handshake) and retransmissions
+   */
+  seqSent: number;
+
   /**
    * Number of unique messages we've received this session (excluding handshake)
    */
   ack: number;
   sendBuffer: Array<OpaqueTransportMessage>;
 
   constructor(props: IdentifiedSessionProps) {
-    const { id, to, seq, ack, sendBuffer, telemetry, log, protocolVersion } =
-      props;
+    const {
+      id,
+      to,
+      seq,
+      ack,
+      sendBuffer,
+      telemetry,
+      log,
+      protocolVersion,
+      seqSent: messagesSent,
+    } = props;
     super(props);
     this.id = id;
     this.to = to;
@@ -242,6 +257,7 @@ export abstract class IdentifiedSession extends CommonSession {
     this.telemetry = telemetry;
     this.log = log;
     this.protocolVersion = protocolVersion;
+    this.seqSent = messagesSent;
   }
 
   get loggingMetadata(): MessageMetadata {

transport/sessionStateMachine/common.ts

@@ -1942,16 +1942,19 @@ describe('session state machine', () => {
       expect(conn.send).toHaveBeenCalledTimes(0);
 
       // send a heartbeat
-      conn.emitData(
-        session.options.codec.toBuffer(
-          session.constructMsg({
-            streamId: 'heartbeat',
-            controlFlags: ControlFlags.AckBit,
-            payload: {
-              type: 'ACK',
-            } satisfies Static<typeof ControlMessageAckSchema>,
-          }),
-        ),
+      conn.onData(
+        session.options.codec.toBuffer({
+          id: 'msgid',
+          to: 'SERVER',
+          from: 'client',
+          seq: 0,
+          ack: 0,
+          streamId: 'heartbeat',
+          controlFlags: ControlFlags.AckBit,
+          payload: {
+            type: 'ACK',
+          } satisfies Static<typeof ControlMessageAckSchema>,
+        }),
       );
 
       // make sure the session acks the heartbeat
@@ -1964,16 +1967,19 @@ describe('session state machine', () => {
       const conn = session.conn;
 
       // send a heartbeat
-      conn.emitData(
-        session.options.codec.toBuffer(
-          session.constructMsg({
-            streamId: 'heartbeat',
-            controlFlags: ControlFlags.AckBit,
-            payload: {
-              type: 'ACK',
-            } satisfies Static<typeof ControlMessageAckSchema>,
-          }),
-        ),
+      conn.onData(
+        session.options.codec.toBuffer({
+          id: 'msgid',
+          to: 'SERVER',
+          from: 'client',
+          seq: 0,
+          ack: 0,
+          streamId: 'heartbeat',
+          controlFlags: ControlFlags.AckBit,
+          payload: {
+            type: 'ACK',
+          } satisfies Static<typeof ControlMessageAckSchema>,
+        }),
       );
 
       expect(sessionHandle.onMessage).not.toHaveBeenCalled();

transport/sessionStateMachine/stateMachine.test.ts

@@ -50,6 +50,7 @@ function inheritSharedSession(
     to: session.to,
     seq: session.seq,
     ack: session.ack,
+    seqSent: session.seqSent,
     sendBuffer: session.sendBuffer,
     telemetry: session.telemetry,
     options: session.options,
@@ -90,6 +91,7 @@ export const SessionStateGraph = {
         to,
         seq: 0,
         ack: 0,
+        seqSent: 0,
         graceExpiryTime: Date.now() + options.sessionDisconnectGraceMs,
         sendBuffer,
         telemetry,
@@ -251,12 +253,13 @@ export const SessionStateGraph = {
         ? // old session exists, inherit state
           inheritSharedSession(oldSession)
         : // old session does not exist, create new state
-          {
+          ({
             id: sessionId,
             from,
             to,
             seq: 0,
             ack: 0,
+            seqSent: 0,
             sendBuffer: [],
             telemetry: createSessionTelemetryInfo(
               pendingSession.tracer,
@@ -269,7 +272,7 @@ export const SessionStateGraph = {
             tracer: pendingSession.tracer,
             log: pendingSession.log,
             protocolVersion,
-          };
+          } satisfies IdentifiedSessionProps);
 
       pendingSession._handleStateExit();
       oldSession?._handleStateExit();