@replit/river 0.23.7 → 0.23.9

package/dist/transport/impls/ws/server.cjs

@@ -196,7 +196,7 @@ var import_nanoid2 = require("nanoid");
 var import_api = require("@opentelemetry/api");
 
 // package.json
-var version = "0.23.7";
+var version = "0.23.9";
 
 // tracing/index.ts
 function createSessionTelemetryInfo(session, propagationCtx) {
@@ -267,6 +267,11 @@ var Session = class {
    * The active connection associated with this session
    */
   connection;
+  /**
+   * A connection that is currently undergoing handshaking. Used to distinguish between the active
+   * connection, but still be able to close it if needed.
+   */
+  handshakingConnection;
   from;
   to;
   /**
@@ -446,23 +451,60 @@ var Session = class {
     this.cancelGrace();
     this.sendBufferedMessages(newConn);
     this.connection = newConn;
+    this.handshakingConnection = void 0;
+  }
+  replaceWithNewHandshakingConnection(newConn) {
+    this.handshakingConnection = newConn;
   }
   beginGrace(cb) {
     this.log?.info(
       `starting ${this.options.sessionDisconnectGraceMs}ms grace period until session to ${this.to} is closed`,
       this.loggingMetadata
     );
+    this.cancelGrace({ keepHeartbeatMisses: true });
     this.disconnectionGrace = setTimeout(() => {
-      this.close();
+      if (this.connection !== void 0) {
+        this.log?.warn(
+          `grace period for ${this.to} elapsed while connected. not calling callback`,
+          {
+            ...this.loggingMetadata,
+            connId: this.connection.id,
+            tags: ["invariant-violation"]
+          }
+        );
+        return;
+      }
+      this.log?.info(
+        `grace period for ${this.to} elapsed`,
+        this.loggingMetadata
+      );
       cb();
     }, this.options.sessionDisconnectGraceMs);
   }
   // called on reconnect of the underlying session
-  cancelGrace() {
-    this.heartbeatMisses = 0;
+  cancelGrace({ keepHeartbeatMisses } = {
+    keepHeartbeatMisses: false
+  }) {
+    if (!keepHeartbeatMisses) {
+      this.heartbeatMisses = 0;
+    }
+    if (this.disconnectionGrace === void 0)
+      return;
     clearTimeout(this.disconnectionGrace);
     this.disconnectionGrace = void 0;
   }
+  /**
+   * Used to close the handshaking connection, if set.
+   */
+  closeHandshakingConnection(expectedHandshakingConn) {
+    if (this.handshakingConnection === void 0)
+      return;
+    if (expectedHandshakingConn !== void 0 && this.handshakingConnection === expectedHandshakingConn) {
+      return;
+    }
+    this.handshakingConnection.close();
+    this.handshakingConnection = void 0;
+  }
   // closed when we want to discard the whole session
   // (i.e. shutdown or session disconnect)
   close() {
@@ -672,7 +714,13 @@ var Transport = class {
     });
     return session;
   }
-  getOrCreateSession(to, conn, sessionId, propagationCtx) {
+  getOrCreateSession({
+    to,
+    conn,
+    handshakingConn,
+    sessionId,
+    propagationCtx
+  }) {
     let session = this.sessions.get(to);
     let isReconnect = session !== void 0;
     if (session?.advertisedSessionId !== void 0 && sessionId !== void 0 && session.advertisedSessionId !== sessionId) {
@@ -680,7 +728,11 @@ var Transport = class {
         `session for ${to} already exists but has a different session id (expected: ${session.advertisedSessionId}, got: ${sessionId}), creating a new one`,
         session.loggingMetadata
       );
-      this.deleteSession(session);
+      this.deleteSession({
+        session,
+        closeHandshakingConnection: handshakingConn !== void 0,
+        handshakingConn
+      });
       isReconnect = false;
       session = void 0;
     }
@@ -694,9 +746,19 @@ var Transport = class {
     if (sessionId !== void 0) {
       session.advertisedSessionId = sessionId;
     }
+    if (handshakingConn !== void 0) {
+      session.replaceWithNewHandshakingConnection(handshakingConn);
+    }
     return { session, isReconnect };
   }
-  deleteSession(session) {
+  deleteSession({
+    session,
+    closeHandshakingConnection,
+    handshakingConn
+  }) {
+    if (closeHandshakingConnection) {
+      session.closeHandshakingConnection(handshakingConn);
+    }
     session.close();
     session.telemetry.span.end();
     this.sessions.delete(session.to);
@@ -723,7 +785,11 @@ var Transport = class {
     session.connection = void 0;
     session.beginGrace(() => {
       session.telemetry.span.addEvent("session grace period expired");
-      this.deleteSession(session);
+      this.deleteSession({
+        session,
+        closeHandshakingConnection: true,
+        handshakingConn: conn
+      });
     });
   }
   /**
@@ -803,7 +869,7 @@ var Transport = class {
           code: import_api3.SpanStatusCode.ERROR,
           message: "message order violated"
         });
-        session.close();
+        this.deleteSession({ session, closeHandshakingConnection: true });
       }
       return;
     }
@@ -857,7 +923,7 @@ var Transport = class {
       });
       return void 0;
     }
-    return this.getOrCreateSession(to).session.send(msg);
+    return this.getOrCreateSession({ to }).session.send(msg);
   }
   // control helpers
   sendCloseStream(to, streamId) {
@@ -880,7 +946,7 @@ var Transport = class {
   close() {
     this.state = "closed";
     for (const session of this.sessions.values()) {
-      this.deleteSession(session);
+      this.deleteSession({ session, closeHandshakingConnection: true });
     }
     this.log?.info(`manually closed transport`, { clientId: this.clientId });
   }
@@ -892,7 +958,7 @@ var Transport = class {
   destroy() {
     this.state = "destroyed";
     for (const session of this.sessions.values()) {
-      this.deleteSession(session);
+      this.deleteSession({ session, closeHandshakingConnection: true });
     }
     this.log?.info(`manually destroyed transport`, { clientId: this.clientId });
   }
@@ -1126,12 +1192,12 @@ var ServerTransport = class extends Transport {
     if (parsedMetadata === false) {
       return false;
     }
-    const { session, isReconnect } = this.getOrCreateSession(
-      parsed.from,
+    const { session, isReconnect } = this.getOrCreateSession({
+      to: parsed.from,
       conn,
-      parsed.payload.sessionId,
-      parsed.tracing
-    );
+      sessionId: parsed.payload.sessionId,
+      propagationCtx: parsed.tracing
+    });
     this.sessionHandshakeMetadata.set(session, parsedMetadata);
     this.log?.debug(
       `handshake from ${parsed.from} ok, responding with handshake success`,
@@ -1167,7 +1233,6 @@ var WebSocketConnection = class extends Connection {
             `websocket closed with code and reason: ${code} - ${reason}`
           )
         );
-        return;
       }
       if (this.closeCb) {
         this.closeCb();