@replit/river 0.21.1 → 0.23.0

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

@@ -26,7 +26,6 @@ module.exports = __toCommonJS(server_exports);
 
 // transport/transport.ts
 var import_value = require("@sinclair/typebox/value");
-var import_api2 = require("@opentelemetry/api");
 
 // transport/message.ts
 var import_typebox = require("@sinclair/typebox");
@@ -142,17 +141,62 @@ var EventDispatcher = class {
 
 // transport/session.ts
 var import_nanoid2 = require("nanoid");
+
+// tracing/index.ts
+var import_api = require("@opentelemetry/api");
+
+// package.json
+var version = "0.23.0";
+
+// tracing/index.ts
+function createSessionTelemetryInfo(session, propagationCtx) {
+  const ctx = propagationCtx ? import_api.propagation.extract(import_api.context.active(), propagationCtx) : import_api.context.active();
+  const span = tracer.startSpan(
+    `session ${session.id}`,
+    {
+      attributes: {
+        component: "river",
+        "river.session.id": session.id,
+        "river.session.to": session.to,
+        "river.session.from": session.from
+      }
+    },
+    ctx
+  );
+  return { span, ctx };
+}
+function createConnectionTelemetryInfo(connection, sessionSpan) {
+  const ctx = import_api.trace.setSpan(import_api.context.active(), sessionSpan);
+  const span = tracer.startSpan(
+    `connection ${connection.id}`,
+    {
+      attributes: {
+        component: "river",
+        "river.connection.id": connection.id
+      },
+      links: [{ context: sessionSpan.spanContext() }]
+    },
+    ctx
+  );
+  return { span, ctx };
+}
+var tracer = import_api.trace.getTracer("river", version);
+
+// transport/session.ts
+var import_api2 = require("@opentelemetry/api");
 var nanoid2 = (0, import_nanoid2.customAlphabet)("1234567890abcdefghijklmnopqrstuvxyz", 6);
 var unsafeId = () => nanoid2();
 var Connection = class {
-  debugId;
+  id;
+  telemetry;
   constructor() {
-    this.debugId = `conn-${unsafeId()}`;
+    this.id = `conn-${nanoid2(12)}`;
   }
 };
 var Session = class {
   codec;
   options;
+  telemetry;
   /**
    * The buffer of messages that have been sent but not yet acknowledged.
    */
@@ -172,12 +216,6 @@ var Session = class {
    * for this session.
    */
   advertisedSessionId;
-  /**
-   * The metadata for this session, as parsed from the handshake.
-   *
-   * Will only ever be populated on the server side.
-   */
-  metadata;
   /**
    * Number of messages we've sent along this session (excluding handshake and acks)
    */
@@ -199,7 +237,7 @@ var Session = class {
    * The interval for sending heartbeats.
    */
   heartbeat;
-  constructor(conn, from, to, options) {
+  constructor(conn, from, to, options, propagationCtx) {
     this.id = `session-${nanoid2(12)}`;
     this.options = options;
     this.from = from;
@@ -211,13 +249,14 @@ var Session = class {
       () => this.sendHeartbeat(),
       options.heartbeatIntervalMs
     );
+    this.telemetry = createSessionTelemetryInfo(this, propagationCtx);
   }
   get loggingMetadata() {
     return {
       clientId: this.from,
       connectedTo: this.to,
       sessionId: this.id,
-      connId: this.connection?.debugId
+      connId: this.connection?.id
     };
   }
   /**
@@ -262,6 +301,7 @@ var Session = class {
           `closing connection to ${this.to} due to inactivity (missed ${misses} heartbeats which is ${missDuration}ms)`,
           this.loggingMetadata
         );
+        this.telemetry.span.addEvent("closing connection due to inactivity");
         this.closeStaleConnection();
       }
       return;
@@ -283,21 +323,25 @@ var Session = class {
   sendBufferedMessages(conn) {
     log?.info(`resending ${this.sendBuffer.length} buffered messages`, {
       ...this.loggingMetadata,
-      connId: conn.debugId
+      connId: conn.id
     });
     for (const msg of this.sendBuffer) {
       log?.debug(`resending msg`, {
         ...this.loggingMetadata,
         fullTransportMessage: msg,
-        connId: conn.debugId
+        connId: conn.id
       });
       const ok = conn.send(this.codec.toBuffer(msg));
       if (!ok) {
         const errMsg = `failed to send buffered message to ${this.to} (sus, this is a fresh connection)`;
+        conn.telemetry?.span.setStatus({
+          code: import_api2.SpanStatusCode.ERROR,
+          message: errMsg
+        });
         log?.error(errMsg, {
           ...this.loggingMetadata,
           fullTransportMessage: msg,
-          connId: conn.debugId,
+          connId: conn.id,
           tags: ["invariant-violation"]
         });
         conn.close();
@@ -380,11 +424,6 @@ var Session = class {
   }
 };
 
-// tracing/index.ts
-var import_api = require("@opentelemetry/api");
-var tracer = import_api.trace.getTracer("river");
-var tracing_default = tracer;
-
 // util/stringify.ts
 function coerceErrorString(err) {
   if (err instanceof Error) {
@@ -446,6 +485,7 @@ var NaiveJsonCodec = {
 };
 
 // transport/transport.ts
+var import_api3 = require("@opentelemetry/api");
 var defaultTransportOptions = {
   heartbeatIntervalMs: 1e3,
   heartbeatsUntilDead: 2,
@@ -524,17 +564,22 @@ var Transport = class {
       status: "connect",
       conn
     });
+    conn.telemetry = createConnectionTelemetryInfo(
+      conn,
+      session.telemetry.span
+    );
     if (isReconnect) {
       session.replaceWithNewConnection(conn);
       log?.info(`reconnected to ${connectedTo}`, session.loggingMetadata);
     }
   }
-  createSession(to, conn) {
+  createSession(to, conn, propagationCtx) {
     const session = new Session(
       conn,
       this.clientId,
       to,
-      this.options
+      this.options,
+      propagationCtx
     );
     this.sessions.set(session.to, session);
     this.eventDispatcher.dispatchEvent("sessionStatus", {
@@ -543,11 +588,11 @@ var Transport = class {
     });
     return session;
   }
-  getOrCreateSession(to, conn, sessionId) {
+  getOrCreateSession(to, conn, sessionId, propagationCtx) {
     let session = this.sessions.get(to);
     let isReconnect = session !== void 0;
     if (session?.advertisedSessionId !== void 0 && sessionId !== void 0 && session.advertisedSessionId !== sessionId) {
-      log?.warn(
+      log?.info(
         `session for ${to} already exists but has a different session id (expected: ${session.advertisedSessionId}, got: ${sessionId}), creating a new one`,
         session.loggingMetadata
       );
@@ -556,7 +601,7 @@ var Transport = class {
       session = void 0;
     }
     if (!session) {
-      session = this.createSession(to, conn);
+      session = this.createSession(to, conn, propagationCtx);
       log?.info(
         `no session for ${to}, created a new one`,
         session.loggingMetadata
@@ -569,6 +614,7 @@ var Transport = class {
   }
   deleteSession(session) {
     session.close();
+    session.telemetry.span.end();
     this.sessions.delete(session.to);
     log?.info(
       `session ${session.id} disconnect from ${session.to}`,
@@ -585,12 +631,16 @@ var Transport = class {
    * @param connectedTo The peer we are connected to.
    */
   onDisconnect(conn, session) {
+    conn.telemetry?.span.end();
     this.eventDispatcher.dispatchEvent("connectionStatus", {
       status: "disconnect",
       conn
     });
     session.connection = void 0;
-    session.beginGrace(() => this.deleteSession(session));
+    session.beginGrace(() => {
+      session.telemetry.span.addEvent("session grace period expired");
+      this.deleteSession(session);
+    });
   }
   /**
    * Parses a message from a Uint8Array into a {@link OpaqueTransportMessage}.
@@ -650,6 +700,10 @@ var Transport = class {
           tags: ["invariant-violation"]
         });
         this.protocolError(ProtocolError.MessageOrderingViolated, errMsg);
+        session.telemetry.span.setStatus({
+          code: import_api3.SpanStatusCode.ERROR,
+          message: "message order violated"
+        });
         session.close();
       }
       return;
@@ -748,280 +802,240 @@ var ServerTransport = class extends Transport {
    * The options for this transport.
    */
   options;
+  /**
+   * Optional handshake options for the server.
+   */
+  handshakeExtensions;
+  /**
+   * A map of session handshake data for each session.
+   */
+  sessionHandshakeMetadata;
   constructor(clientId, providedOptions) {
     super(clientId, providedOptions);
     this.options = {
       ...defaultServerTransportOptions,
       ...providedOptions
     };
+    this.sessionHandshakeMetadata = /* @__PURE__ */ new WeakMap();
     log?.info(`initiated server transport`, {
       clientId: this.clientId,
       protocolVersion: PROTOCOL_VERSION
     });
   }
+  extendHandshake(options) {
+    this.handshakeExtensions = options;
+  }
   handleConnection(conn) {
-    tracing_default.startActiveSpan(
-      "handleConnection",
-      {
-        attributes: {
-          component: "river",
-          "span.kind": "server"
-        },
-        kind: import_api2.SpanKind.SERVER
-      },
-      (span) => {
-        if (this.state !== "open")
-          return;
-        log?.info(`new incoming connection`, {
-          clientId: this.clientId,
-          connId: conn.debugId
+    if (this.state !== "open")
+      return;
+    log?.info(`new incoming connection`, {
+      clientId: this.clientId,
+      connId: conn.id
+    });
+    let session = void 0;
+    const client = () => session?.to ?? "unknown";
+    const handshakeTimeout = setTimeout(() => {
+      if (!session) {
+        log?.warn(
+          `connection to ${client()} timed out waiting for handshake, closing`,
+          {
+            clientId: this.clientId,
+            connectedTo: client(),
+            connId: conn.id
+          }
+        );
+        conn.telemetry?.span.setStatus({
+          code: import_api3.SpanStatusCode.ERROR,
+          message: "handshake timeout"
         });
-        let session = void 0;
-        const client = () => session?.to ?? "unknown";
-        const handshakeTimeout = setTimeout(() => {
-          if (!session) {
-            log?.warn(
-              `connection to ${client()} timed out waiting for handshake, closing`,
-              {
-                clientId: this.clientId,
-                connectedTo: client(),
-                connId: conn.debugId
-              }
-            );
-            span.setStatus({ code: import_api2.SpanStatusCode.ERROR });
-            span.end();
+        conn.close();
+      }
+    }, this.options.sessionDisconnectGraceMs);
+    const buffer = [];
+    let receivedHandshakeMessage = false;
+    const handshakeHandler = (data) => {
+      if (receivedHandshakeMessage) {
+        buffer.push(data);
+        return;
+      }
+      receivedHandshakeMessage = true;
+      clearTimeout(handshakeTimeout);
+      void this.receiveHandshakeRequestMessage(data, conn).then(
+        (maybeSession) => {
+          if (!maybeSession) {
             conn.close();
-          }
-        }, this.options.sessionDisconnectGraceMs);
-        const buffer = [];
-        let receivedHandshakeMessage = false;
-        const handshakeHandler = (data) => {
-          if (receivedHandshakeMessage) {
-            buffer.push(data);
             return;
           }
-          receivedHandshakeMessage = true;
-          clearTimeout(handshakeTimeout);
-          void this.receiveHandshakeRequestMessage(data, conn).then(
-            (maybeSession) => {
-              if (!maybeSession) {
-                span.setStatus({ code: import_api2.SpanStatusCode.ERROR });
-                span.end();
-                conn.close();
-                return;
-              }
-              session = maybeSession;
-              const dataHandler = (data2) => {
-                const parsed = this.parseMsg(data2);
-                if (!parsed) {
-                  conn.close();
-                  return;
-                }
-                this.handleMsg(parsed);
-              };
-              conn.removeDataListener(handshakeHandler);
-              conn.addDataListener(dataHandler);
-              for (const data2 of buffer) {
-                dataHandler(data2);
-              }
-              buffer.length = 0;
+          session = maybeSession;
+          const dataHandler = (data2) => {
+            const parsed = this.parseMsg(data2);
+            if (!parsed) {
+              conn.close();
+              return;
             }
-          );
-        };
-        conn.addDataListener(handshakeHandler);
-        conn.addCloseListener(() => {
-          if (session) {
-            log?.info(`connection to ${client()} disconnected`, {
-              clientId: this.clientId,
-              connId: conn.debugId
-            });
-            this.onDisconnect(conn, session);
+            this.handleMsg(parsed);
+          };
+          for (const data2 of buffer) {
+            dataHandler(data2);
           }
-          span.setStatus({ code: import_api2.SpanStatusCode.OK });
-          span.end();
+          conn.removeDataListener(handshakeHandler);
+          conn.addDataListener(dataHandler);
+          buffer.length = 0;
+        }
+      );
+    };
+    conn.addDataListener(handshakeHandler);
+    conn.addCloseListener(() => {
+      if (!session)
+        return;
+      log?.info(`connection to ${client()} disconnected`, {
+        clientId: this.clientId,
+        connId: conn.id
+      });
+      this.onDisconnect(conn, session);
+    });
+    conn.addErrorListener((err) => {
+      conn.telemetry?.span.setStatus({
+        code: import_api3.SpanStatusCode.ERROR,
+        message: "connection error"
+      });
+      if (!session)
+        return;
+      log?.warn(
+        `connection to ${client()} got an error: ${coerceErrorString(err)}`,
+        { clientId: this.clientId, connId: conn.id }
+      );
+    });
+  }
+  async validateHandshakeMetadata(conn, session, rawMetadata, from) {
+    let parsedMetadata = {};
+    if (this.handshakeExtensions) {
+      if (!import_value.Value.Check(this.handshakeExtensions.schema, rawMetadata)) {
+        conn.telemetry?.span.setStatus({
+          code: import_api3.SpanStatusCode.ERROR,
+          message: "malformed handshake meta"
         });
-        conn.addErrorListener((err) => {
-          if (session) {
-            log?.warn(
-              `connection to ${client()} got an error: ${coerceErrorString(
-                err
-              )}`,
-              { clientId: this.clientId, connId: conn.debugId }
-            );
-          }
-          span.setStatus({ code: import_api2.SpanStatusCode.ERROR });
-          span.end();
+        const reason = "received malformed handshake metadata";
+        const responseMsg = handshakeResponseMessage(this.clientId, from, {
+          ok: false,
+          reason
+        });
+        conn.send(this.codec.toBuffer(responseMsg));
+        log?.warn(`received malformed handshake metadata from ${from}`, {
+          clientId: this.clientId,
+          connId: conn.id
         });
+        this.protocolError(ProtocolError.HandshakeFailed, reason);
+        return false;
       }
-    );
+      parsedMetadata = await this.handshakeExtensions.validate(
+        rawMetadata,
+        session
+      );
+      if (parsedMetadata === false) {
+        const reason = "rejected by handshake handler";
+        conn.telemetry?.span.setStatus({
+          code: import_api3.SpanStatusCode.ERROR,
+          message: reason
+        });
+        const responseMsg = handshakeResponseMessage(this.clientId, from, {
+          ok: false,
+          reason
+        });
+        conn.send(this.codec.toBuffer(responseMsg));
+        log?.warn(`rejected handshake from ${from}`, {
+          clientId: this.clientId,
+          connId: conn.id
+        });
+        this.protocolError(ProtocolError.HandshakeFailed, reason);
+        return false;
+      }
+    }
+    return parsedMetadata;
   }
   async receiveHandshakeRequestMessage(data, conn) {
     const parsed = this.parseMsg(data);
     if (!parsed) {
+      conn.telemetry?.span.setStatus({
+        code: import_api3.SpanStatusCode.ERROR,
+        message: "non-transport message"
+      });
       this.protocolError(
         ProtocolError.HandshakeFailed,
         "received non-transport message"
       );
       return false;
     }
-    let activeContext = import_api2.context.active();
-    if (parsed.tracing) {
-      activeContext = import_api2.propagation.extract(activeContext, parsed.tracing);
+    if (!import_value.Value.Check(ControlMessageHandshakeRequestSchema, parsed.payload)) {
+      conn.telemetry?.span.setStatus({
+        code: import_api3.SpanStatusCode.ERROR,
+        message: "invalid handshake request"
+      });
+      const reason = "received invalid handshake msg";
+      const responseMsg2 = handshakeResponseMessage(this.clientId, parsed.from, {
+        ok: false,
+        reason
+      });
+      conn.send(this.codec.toBuffer(responseMsg2));
+      const logData = { ...parsed.payload ?? {}, metadata: "redacted" };
+      log?.warn(reason, {
+        clientId: this.clientId,
+        connId: conn.id,
+        partialTransportMessage: { ...parsed, payload: logData }
+      });
+      this.protocolError(
+        ProtocolError.HandshakeFailed,
+        "invalid handshake request"
+      );
+      return false;
     }
-    return tracing_default.startActiveSpan(
-      "receiveHandshakeRequestMessage",
-      {
-        attributes: {
-          component: "river",
-          "span.kind": "server"
-        },
-        kind: import_api2.SpanKind.SERVER
-      },
-      activeContext,
-      async (span) => {
-        if (!import_value.Value.Check(ControlMessageHandshakeRequestSchema, parsed.payload)) {
-          const reason = "received invalid handshake msg";
-          const responseMsg2 = handshakeResponseMessage(
-            this.clientId,
-            parsed.from,
-            {
-              ok: false,
-              reason
-            }
-          );
-          conn.send(this.codec.toBuffer(responseMsg2));
-          const logData = typeof parsed.payload === "object" ? {
-            ...parsed,
-            payload: { ...parsed.payload, metadata: "redacted" }
-          } : { ...parsed };
-          log?.warn(`${reason}: ${JSON.stringify(logData)}`, {
-            clientId: this.clientId,
-            connId: conn.debugId
-          });
-          this.protocolError(
-            ProtocolError.HandshakeFailed,
-            "invalid handshake request"
-          );
-          span.setStatus({ code: import_api2.SpanStatusCode.ERROR });
-          span.end();
-          return false;
-        }
-        const gotVersion = parsed.payload.protocolVersion;
-        if (gotVersion !== PROTOCOL_VERSION) {
-          const reason = `incorrect version (got: ${gotVersion} wanted ${PROTOCOL_VERSION})`;
-          const responseMsg2 = handshakeResponseMessage(
-            this.clientId,
-            parsed.from,
-            {
-              ok: false,
-              reason
-            }
-          );
-          conn.send(this.codec.toBuffer(responseMsg2));
-          log?.warn(
-            `received handshake msg with incompatible protocol version (got: ${gotVersion}, expected: ${PROTOCOL_VERSION})`,
-            { clientId: this.clientId, connId: conn.debugId }
-          );
-          this.protocolError(ProtocolError.HandshakeFailed, reason);
-          span.setStatus({ code: import_api2.SpanStatusCode.ERROR });
-          span.end();
-          return false;
-        }
-        const { session, isReconnect } = this.getOrCreateSession(
-          parsed.from,
-          conn,
-          parsed.payload.sessionId
-        );
-        let handshakeMetadata;
-        if (this.options.handshake) {
-          if (!import_value.Value.Check(
-            this.options.handshake.requestSchema,
-            parsed.payload.metadata
-          )) {
-            const reason = "received malformed handshake metadata";
-            const responseMsg2 = handshakeResponseMessage(
-              this.clientId,
-              parsed.from,
-              { ok: false, reason }
-            );
-            conn.send(this.codec.toBuffer(responseMsg2));
-            log?.warn(
-              `received malformed handshake metadata from ${parsed.from}`,
-              {
-                clientId: this.clientId,
-                connId: conn.debugId
-              }
-            );
-            this.protocolError(ProtocolError.HandshakeFailed, reason);
-            this.deleteSession(session);
-            span.setStatus({ code: import_api2.SpanStatusCode.ERROR });
-            span.end();
-            return false;
-          }
-          const parsedMetadata = await this.options.handshake.parse(
-            parsed.payload.metadata,
-            session,
-            isReconnect
-          );
-          if (parsedMetadata === false) {
-            const reason = "rejected by server";
-            const responseMsg2 = handshakeResponseMessage(
-              this.clientId,
-              parsed.from,
-              { ok: false, reason }
-            );
-            conn.send(this.codec.toBuffer(responseMsg2));
-            log?.warn(`rejected handshake from ${parsed.from}`, {
-              clientId: this.clientId,
-              connId: conn.debugId
-            });
-            this.protocolError(ProtocolError.HandshakeFailed, reason);
-            this.deleteSession(session);
-            span.setStatus({ code: import_api2.SpanStatusCode.ERROR });
-            span.end();
-            return false;
-          }
-          if (!import_value.Value.Check(this.options.handshake.parsedSchema, parsedMetadata)) {
-            const reason = "failed to parse handshake metadata";
-            const responseMsg2 = handshakeResponseMessage(
-              this.clientId,
-              parsed.from,
-              { ok: false, reason }
-            );
-            conn.send(this.codec.toBuffer(responseMsg2));
-            log?.error(`failed to parse handshake metadata`, {
-              clientId: this.clientId,
-              connId: conn.debugId,
-              tags: ["invariant-violation"]
-            });
-            this.protocolError(ProtocolError.HandshakeFailed, reason);
-            this.deleteSession(session);
-            span.setStatus({ code: import_api2.SpanStatusCode.ERROR });
-            span.end();
-            return false;
-          }
-          handshakeMetadata = parsedMetadata;
-        }
-        handshakeMetadata ??= {};
-        session.metadata = handshakeMetadata;
-        log?.debug(
-          `handshake from ${parsed.from} ok, responding with handshake success`,
-          { clientId: this.clientId, connId: conn.debugId }
-        );
-        const responseMsg = handshakeResponseMessage(
-          this.clientId,
-          parsed.from,
-          {
-            ok: true,
-            sessionId: session.id
-          }
-        );
-        conn.send(this.codec.toBuffer(responseMsg));
-        this.onConnect(conn, parsed.from, session, isReconnect);
-        span.end();
-        return session;
-      }
+    const gotVersion = parsed.payload.protocolVersion;
+    if (gotVersion !== PROTOCOL_VERSION) {
+      conn.telemetry?.span.setStatus({
+        code: import_api3.SpanStatusCode.ERROR,
+        message: "incorrect protocol version"
+      });
+      const reason = `incorrect version (got: ${gotVersion} wanted ${PROTOCOL_VERSION})`;
+      const responseMsg2 = handshakeResponseMessage(this.clientId, parsed.from, {
+        ok: false,
+        reason
+      });
+      conn.send(this.codec.toBuffer(responseMsg2));
+      log?.warn(
+        `received handshake msg with incompatible protocol version (got: ${gotVersion}, expected: ${PROTOCOL_VERSION})`,
+        { clientId: this.clientId, connId: conn.id }
+      );
+      this.protocolError(ProtocolError.HandshakeFailed, reason);
+      return false;
+    }
+    const oldSession = this.sessions.get(parsed.from);
+    const parsedMetadata = await this.validateHandshakeMetadata(
+      conn,
+      oldSession,
+      parsed.payload.metadata,
+      parsed.from
     );
+    if (parsedMetadata === false) {
+      return false;
+    }
+    const { session, isReconnect } = this.getOrCreateSession(
+      parsed.from,
+      conn,
+      parsed.payload.sessionId,
+      parsed.tracing
+    );
+    this.sessionHandshakeMetadata.set(session, parsedMetadata);
+    log?.debug(
+      `handshake from ${parsed.from} ok, responding with handshake success`,
+      { clientId: this.clientId, connId: conn.id }
+    );
+    const responseMsg = handshakeResponseMessage(this.clientId, parsed.from, {
+      ok: true,
+      sessionId: session.id
+    });
+    conn.send(this.codec.toBuffer(responseMsg));
+    this.onConnect(conn, parsed.from, session, isReconnect);
+    return session;
   }
 };