rivetkit 2.0.22-rc.2 → 2.0.23

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (109) hide show
  1. package/dist/schemas/actor-persist/v2.ts +29 -26
  2. package/dist/tsup/{chunk-FLVL7RGH.js → chunk-3I6ZIJVJ.js} +3 -3
  3. package/dist/tsup/{chunk-GXIO5YOT.cjs → chunk-3JYSUFET.cjs} +24 -10
  4. package/dist/tsup/chunk-3JYSUFET.cjs.map +1 -0
  5. package/dist/tsup/{chunk-NDOG6IQ5.cjs → chunk-54DVMQPT.cjs} +6 -6
  6. package/dist/tsup/{chunk-NDOG6IQ5.cjs.map → chunk-54DVMQPT.cjs.map} +1 -1
  7. package/dist/tsup/{chunk-7RUROQAZ.js → chunk-5PKKNNNS.js} +279 -9
  8. package/dist/tsup/chunk-5PKKNNNS.js.map +1 -0
  9. package/dist/tsup/{chunk-F7WVJXPB.js → chunk-5UJQWWO3.js} +2 -2
  10. package/dist/tsup/{chunk-YUBR6XCJ.cjs → chunk-C56XVVV4.cjs} +280 -10
  11. package/dist/tsup/chunk-C56XVVV4.cjs.map +1 -0
  12. package/dist/tsup/{chunk-Q5CAVEKC.cjs → chunk-D6PCH7FR.cjs} +561 -487
  13. package/dist/tsup/chunk-D6PCH7FR.cjs.map +1 -0
  14. package/dist/tsup/{chunk-C4FPCW7T.js → chunk-DLYZKFRY.js} +2 -2
  15. package/dist/tsup/{chunk-AMK3AACS.js → chunk-FTQ62XTN.js} +373 -299
  16. package/dist/tsup/chunk-FTQ62XTN.js.map +1 -0
  17. package/dist/tsup/{chunk-LFP446KS.cjs → chunk-HNYF4T36.cjs} +14 -14
  18. package/dist/tsup/{chunk-LFP446KS.cjs.map → chunk-HNYF4T36.cjs.map} +1 -1
  19. package/dist/tsup/{chunk-5TRXLS6X.cjs → chunk-JMLTKMJ7.cjs} +48 -44
  20. package/dist/tsup/chunk-JMLTKMJ7.cjs.map +1 -0
  21. package/dist/tsup/{chunk-ZY4DKLMT.cjs → chunk-NCUALX2Q.cjs} +3 -3
  22. package/dist/tsup/{chunk-ZY4DKLMT.cjs.map → chunk-NCUALX2Q.cjs.map} +1 -1
  23. package/dist/tsup/{chunk-HLZT5C6A.js → chunk-NOZSCUPQ.js} +99 -50
  24. package/dist/tsup/chunk-NOZSCUPQ.js.map +1 -0
  25. package/dist/tsup/{chunk-CVLO2OOK.js → chunk-PHNIVSG5.js} +19 -5
  26. package/dist/tsup/chunk-PHNIVSG5.js.map +1 -0
  27. package/dist/tsup/{chunk-BHLQTKOD.js → chunk-RUTBXBRR.js} +27 -23
  28. package/dist/tsup/{chunk-BHLQTKOD.js.map → chunk-RUTBXBRR.js.map} +1 -1
  29. package/dist/tsup/{chunk-MQDXPGNE.js → chunk-RVVUS4X6.js} +6 -6
  30. package/dist/tsup/{chunk-UBMUBNS2.cjs → chunk-SN4KWTRA.cjs} +12 -12
  31. package/dist/tsup/{chunk-UBMUBNS2.cjs.map → chunk-SN4KWTRA.cjs.map} +1 -1
  32. package/dist/tsup/{chunk-ZL6NSKF2.cjs → chunk-XSDSNHSE.cjs} +3 -3
  33. package/dist/tsup/{chunk-ZL6NSKF2.cjs.map → chunk-XSDSNHSE.cjs.map} +1 -1
  34. package/dist/tsup/{chunk-YLWF6RFL.cjs → chunk-XYK5PY3B.cjs} +283 -234
  35. package/dist/tsup/chunk-XYK5PY3B.cjs.map +1 -0
  36. package/dist/tsup/{chunk-EJXZYQ3N.js → chunk-YAYNBR37.js} +2 -2
  37. package/dist/tsup/client/mod.cjs +8 -9
  38. package/dist/tsup/client/mod.cjs.map +1 -1
  39. package/dist/tsup/client/mod.d.cts +2 -2
  40. package/dist/tsup/client/mod.d.ts +2 -2
  41. package/dist/tsup/client/mod.js +7 -8
  42. package/dist/tsup/common/log.cjs +2 -3
  43. package/dist/tsup/common/log.cjs.map +1 -1
  44. package/dist/tsup/common/log.js +1 -2
  45. package/dist/tsup/common/websocket.cjs +3 -4
  46. package/dist/tsup/common/websocket.cjs.map +1 -1
  47. package/dist/tsup/common/websocket.js +2 -3
  48. package/dist/tsup/{conn-BYXlxnh0.d.ts → conn-B3Vhbgnd.d.ts} +5 -1
  49. package/dist/tsup/{conn-BiazosE_.d.cts → conn-DJWL3nGx.d.cts} +5 -1
  50. package/dist/tsup/driver-helpers/mod.cjs +4 -5
  51. package/dist/tsup/driver-helpers/mod.cjs.map +1 -1
  52. package/dist/tsup/driver-helpers/mod.d.cts +1 -1
  53. package/dist/tsup/driver-helpers/mod.d.ts +1 -1
  54. package/dist/tsup/driver-helpers/mod.js +3 -4
  55. package/dist/tsup/driver-test-suite/mod.cjs +70 -72
  56. package/dist/tsup/driver-test-suite/mod.cjs.map +1 -1
  57. package/dist/tsup/driver-test-suite/mod.d.cts +1 -1
  58. package/dist/tsup/driver-test-suite/mod.d.ts +1 -1
  59. package/dist/tsup/driver-test-suite/mod.js +11 -13
  60. package/dist/tsup/driver-test-suite/mod.js.map +1 -1
  61. package/dist/tsup/inspector/mod.cjs +5 -6
  62. package/dist/tsup/inspector/mod.cjs.map +1 -1
  63. package/dist/tsup/inspector/mod.d.cts +2 -2
  64. package/dist/tsup/inspector/mod.d.ts +2 -2
  65. package/dist/tsup/inspector/mod.js +4 -5
  66. package/dist/tsup/mod.cjs +9 -10
  67. package/dist/tsup/mod.cjs.map +1 -1
  68. package/dist/tsup/mod.d.cts +4 -4
  69. package/dist/tsup/mod.d.ts +4 -4
  70. package/dist/tsup/mod.js +8 -9
  71. package/dist/tsup/test/mod.cjs +10 -11
  72. package/dist/tsup/test/mod.cjs.map +1 -1
  73. package/dist/tsup/test/mod.d.cts +1 -1
  74. package/dist/tsup/test/mod.d.ts +1 -1
  75. package/dist/tsup/test/mod.js +9 -10
  76. package/dist/tsup/utils.cjs +2 -2
  77. package/dist/tsup/utils.js +1 -1
  78. package/package.json +2 -2
  79. package/src/actor/conn-drivers.ts +0 -32
  80. package/src/actor/conn-socket.ts +2 -0
  81. package/src/actor/conn.ts +13 -12
  82. package/src/actor/instance.ts +164 -36
  83. package/src/actor/persisted.ts +4 -1
  84. package/src/actor/router-endpoints.ts +14 -0
  85. package/src/actor/router.ts +2 -0
  86. package/src/actor/utils.test.ts +48 -0
  87. package/src/actor/utils.ts +23 -0
  88. package/src/drivers/engine/actor-driver.ts +97 -35
  89. package/src/drivers/file-system/manager.ts +4 -0
  90. package/src/schemas/actor-persist/versioned.ts +4 -0
  91. package/src/utils.ts +15 -6
  92. package/dist/tsup/chunk-5N6F5PXD.cjs +0 -269
  93. package/dist/tsup/chunk-5N6F5PXD.cjs.map +0 -1
  94. package/dist/tsup/chunk-5TRXLS6X.cjs.map +0 -1
  95. package/dist/tsup/chunk-7RUROQAZ.js.map +0 -1
  96. package/dist/tsup/chunk-AMK3AACS.js.map +0 -1
  97. package/dist/tsup/chunk-CVLO2OOK.js.map +0 -1
  98. package/dist/tsup/chunk-GXIO5YOT.cjs.map +0 -1
  99. package/dist/tsup/chunk-HLZT5C6A.js.map +0 -1
  100. package/dist/tsup/chunk-Q5CAVEKC.cjs.map +0 -1
  101. package/dist/tsup/chunk-VMFBKBJL.js +0 -269
  102. package/dist/tsup/chunk-VMFBKBJL.js.map +0 -1
  103. package/dist/tsup/chunk-YLWF6RFL.cjs.map +0 -1
  104. package/dist/tsup/chunk-YUBR6XCJ.cjs.map +0 -1
  105. /package/dist/tsup/{chunk-FLVL7RGH.js.map → chunk-3I6ZIJVJ.js.map} +0 -0
  106. /package/dist/tsup/{chunk-F7WVJXPB.js.map → chunk-5UJQWWO3.js.map} +0 -0
  107. /package/dist/tsup/{chunk-C4FPCW7T.js.map → chunk-DLYZKFRY.js.map} +0 -0
  108. /package/dist/tsup/{chunk-MQDXPGNE.js.map → chunk-RVVUS4X6.js.map} +0 -0
  109. /package/dist/tsup/{chunk-EJXZYQ3N.js.map → chunk-YAYNBR37.js.map} +0 -0
@@ -19,6 +19,7 @@ import {
19
19
  bufferToArrayBuffer,
20
20
  EXTRA_ERROR_LOG,
21
21
  getEnvUniversal,
22
+ idToStr,
22
23
  promiseWithResolvers,
23
24
  SinglePromiseQueue,
24
25
  } from "@/utils";
@@ -53,7 +54,7 @@ import type {
53
54
  import { processMessage } from "./protocol/old";
54
55
  import { CachedSerializer } from "./protocol/serde";
55
56
  import { Schedule } from "./schedule";
56
- import { DeadlineError, deadline } from "./utils";
57
+ import { DeadlineError, deadline, isConnStatePath, isStatePath } from "./utils";
57
58
 
58
59
  export const PERSIST_SYMBOL = Symbol("persist");
59
60
 
@@ -244,7 +245,10 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
244
245
  lastSeen: conn.lastSeen,
245
246
  stateEnabled: conn.__stateEnabled,
246
247
  isHibernatable: conn.isHibernatable,
247
- requestId: conn.__socket?.requestId,
248
+ hibernatableRequestId: conn.__persist
249
+ .hibernatableRequestId
250
+ ? idToStr(conn.__persist.hibernatableRequestId)
251
+ : undefined,
248
252
  driver: conn.__driverState
249
253
  ? getConnDriverKindFromState(conn.__driverState)
250
254
  : undefined,
@@ -267,6 +271,7 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
267
271
  const conn = await this.createConn(
268
272
  {
269
273
  requestId: requestId,
274
+ hibernatable: false,
270
275
  driverState: { [ConnDriverKind.HTTP]: {} },
271
276
  },
272
277
  undefined,
@@ -676,6 +681,10 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
676
681
 
677
682
  this.#onPersistSavedPromise?.resolve();
678
683
  } catch (error) {
684
+ this.#rLog.error({
685
+ msg: "error saving persist",
686
+ error: stringifyError(error),
687
+ });
679
688
  this.#onPersistSavedPromise?.reject(error);
680
689
  throw error;
681
690
  }
@@ -694,7 +703,6 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
694
703
  // Set raw persist object
695
704
  this.#persistRaw = target;
696
705
 
697
- // TODO: Only validate this for conn state
698
706
  // TODO: Allow disabling in production
699
707
  // If this can't be proxied, return raw value
700
708
  if (target === null || typeof target !== "object") {
@@ -728,35 +736,46 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
728
736
  _previousValue: any,
729
737
  _applyData: any,
730
738
  ) => {
731
- if (path !== "state" && !path.startsWith("state.")) {
732
- return;
739
+ const actorStatePath = isStatePath(path);
740
+ const connStatePath = isConnStatePath(path);
741
+
742
+ // Validate CBOR serializability for state changes
743
+ if (actorStatePath || connStatePath) {
744
+ let invalidPath = "";
745
+ if (
746
+ !isCborSerializable(
747
+ value,
748
+ (invalidPathPart) => {
749
+ invalidPath = invalidPathPart;
750
+ },
751
+ "",
752
+ )
753
+ ) {
754
+ throw new errors.InvalidStateType({
755
+ path: path + (invalidPath ? `.${invalidPath}` : ""),
756
+ });
757
+ }
733
758
  }
734
759
 
735
- let invalidPath = "";
736
- if (
737
- !isCborSerializable(
738
- value,
739
- (invalidPathPart) => {
740
- invalidPath = invalidPathPart;
741
- },
742
- "",
743
- )
744
- ) {
745
- throw new errors.InvalidStateType({
746
- path: path + (invalidPath ? `.${invalidPath}` : ""),
747
- });
748
- }
760
+ this.#rLog.debug({
761
+ msg: "onChange triggered, setting persistChanged=true",
762
+ path,
763
+ });
749
764
  this.#persistChanged = true;
750
765
 
751
- // Inform the inspector about state changes
752
- this.inspector.emitter.emit(
753
- "stateUpdated",
754
- this.#persist.state,
755
- );
766
+ // Inform the inspector about state changes (only for state path)
767
+ if (actorStatePath) {
768
+ this.inspector.emitter.emit(
769
+ "stateUpdated",
770
+ this.#persist.state,
771
+ );
772
+ }
756
773
 
757
774
  // Call onStateChange if it exists
775
+ //
758
776
  // Skip if we're already inside onStateChange to prevent infinite recursion
759
777
  if (
778
+ actorStatePath &&
760
779
  this.#config.onStateChange &&
761
780
  this.#ready &&
762
781
  !this.#isInOnStateChange
@@ -802,6 +821,8 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
802
821
  this.#rLog.info({
803
822
  msg: "actor restoring",
804
823
  connections: persistData.connections.length,
824
+ hibernatableWebSockets:
825
+ persistData.hibernatableWebSocket.length,
805
826
  });
806
827
 
807
828
  // Set initial state
@@ -1000,6 +1021,74 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
1000
1021
  ): Promise<Conn<S, CP, CS, V, I, DB>> {
1001
1022
  this.#assertReady();
1002
1023
 
1024
+ // Check for hibernatable websocket reconnection
1025
+ if (socket.requestIdBuf && socket.hibernatable) {
1026
+ this.rLog.debug({
1027
+ msg: "checking for hibernatable websocket connection",
1028
+ requestId: socket.requestId,
1029
+ existingConnectionsCount: this.#connections.size,
1030
+ });
1031
+
1032
+ // Find existing connection with matching hibernatableRequestId
1033
+ const existingConn = Array.from(this.#connections.values()).find(
1034
+ (conn) =>
1035
+ conn.__persist.hibernatableRequestId &&
1036
+ arrayBuffersEqual(
1037
+ conn.__persist.hibernatableRequestId,
1038
+ socket.requestIdBuf!,
1039
+ ),
1040
+ );
1041
+
1042
+ if (existingConn) {
1043
+ this.rLog.debug({
1044
+ msg: "reconnecting hibernatable websocket connection",
1045
+ connectionId: existingConn.id,
1046
+ requestId: socket.requestId,
1047
+ });
1048
+
1049
+ // If there's an existing driver state, clean it up without marking as clean disconnect
1050
+ if (existingConn.__driverState) {
1051
+ this.#rLog.warn({
1052
+ msg: "found existing driver state on hibernatable websocket",
1053
+ connectionId: existingConn.id,
1054
+ requestId: socket.requestId,
1055
+ });
1056
+ const driverKind = getConnDriverKindFromState(
1057
+ existingConn.__driverState,
1058
+ );
1059
+ const driver = CONN_DRIVERS[driverKind];
1060
+ if (driver.disconnect) {
1061
+ // Call driver disconnect to clean up directly. Don't use Conn.disconnect since that will remove the connection entirely.
1062
+ driver.disconnect(
1063
+ this,
1064
+ existingConn,
1065
+ (existingConn.__driverState as any)[driverKind],
1066
+ "Reconnecting hibernatable websocket with new driver state",
1067
+ );
1068
+ }
1069
+ }
1070
+
1071
+ // Update with new driver state
1072
+ existingConn.__socket = socket;
1073
+ existingConn.__persist.lastSeen = Date.now();
1074
+
1075
+ // Update sleep timer since connection is now active
1076
+ this.#resetSleepTimer();
1077
+
1078
+ this.inspector.emitter.emit("connectionUpdated");
1079
+
1080
+ // We don't need to send a new init message since this is a
1081
+ // hibernated request that has already been initialized
1082
+
1083
+ return existingConn;
1084
+ } else {
1085
+ this.rLog.debug({
1086
+ msg: "no existing hibernatable connection found, creating new connection",
1087
+ requestId: socket.requestId,
1088
+ });
1089
+ }
1090
+ }
1091
+
1003
1092
  // If connection ID and token are provided, try to reconnect
1004
1093
  if (connectionId && connectionToken) {
1005
1094
  this.rLog.debug({
@@ -1058,14 +1147,12 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
1058
1147
  );
1059
1148
 
1060
1149
  return existingConn;
1150
+ } else {
1151
+ this.rLog.debug({
1152
+ msg: "connection not found or token mismatch, creating new connection",
1153
+ connectionId,
1154
+ });
1061
1155
  }
1062
-
1063
- // If we get here, either connection doesn't exist or token doesn't match
1064
- // Fall through to create new connection with new IDs
1065
- this.rLog.debug({
1066
- msg: "connection not found or token mismatch, creating new connection",
1067
- connectionId,
1068
- });
1069
1156
  }
1070
1157
 
1071
1158
  // Generate new connection ID and token if not provided or if reconnection failed
@@ -1131,6 +1218,19 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
1131
1218
  lastSeen: Date.now(),
1132
1219
  subscriptions: [],
1133
1220
  };
1221
+
1222
+ // Check if this connection is for a hibernatable websocket
1223
+ if (socket.requestIdBuf) {
1224
+ const isHibernatable =
1225
+ this.#persist.hibernatableWebSocket.findIndex((ws) =>
1226
+ arrayBuffersEqual(ws.requestId, socket.requestIdBuf!),
1227
+ ) !== -1;
1228
+
1229
+ if (isHibernatable) {
1230
+ persist.hibernatableRequestId = socket.requestIdBuf;
1231
+ }
1232
+ }
1233
+
1134
1234
  const conn = new Conn<S, CP, CS, V, I, DB>(this, persist);
1135
1235
  conn.__socket = socket;
1136
1236
  this.#connections.set(conn.id, conn);
@@ -1864,6 +1964,13 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
1864
1964
  async saveState(opts: SaveStateOptions) {
1865
1965
  this.#assertReady(opts.allowStoppingState);
1866
1966
 
1967
+ this.#rLog.debug({
1968
+ msg: "saveState called",
1969
+ persistChanged: this.#persistChanged,
1970
+ allowStoppingState: opts.allowStoppingState,
1971
+ immediate: opts.immediate,
1972
+ });
1973
+
1867
1974
  if (this.#persistChanged) {
1868
1975
  if (opts.immediate) {
1869
1976
  // Save immediately
@@ -1920,6 +2027,9 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
1920
2027
  #resetSleepTimer() {
1921
2028
  if (this.#config.options.noSleep || !this.#sleepingSupported) return;
1922
2029
 
2030
+ // Don't sleep if already stopping
2031
+ if (this.#stopCalled) return;
2032
+
1923
2033
  const canSleep = this.#canSleep();
1924
2034
 
1925
2035
  this.#rLog.debug({
@@ -1979,11 +2089,20 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
1979
2089
  * 4. Engine runner will publish EventActorStateUpdate with ActorStateSTop
1980
2090
  **/
1981
2091
  _startSleep() {
2092
+ if (this.#stopCalled) {
2093
+ this.#rLog.debug({
2094
+ msg: "cannot call _startSleep if actor already stopping",
2095
+ });
2096
+ return;
2097
+ }
2098
+
1982
2099
  // IMPORTANT: #sleepCalled should have no effect on the actor's
1983
2100
  // behavior aside from preventing calling _startSleep twice. Wait for
1984
2101
  // `_onStop` before putting in a stopping state.
1985
2102
  if (this.#sleepCalled) {
1986
- this.#rLog.warn({ msg: "already sleeping actor" });
2103
+ this.#rLog.warn({
2104
+ msg: "cannot call _startSleep twice, actor already sleeping",
2105
+ });
1987
2106
  return;
1988
2107
  }
1989
2108
  this.#sleepCalled = true;
@@ -2054,12 +2173,20 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
2054
2173
  }
2055
2174
  }
2056
2175
 
2057
- // Disconnect existing connections
2058
2176
  const promises: Promise<unknown>[] = [];
2177
+
2178
+ // Disconnect existing non-hibernatable connections
2059
2179
  for (const connection of this.#connections.values()) {
2060
- promises.push(connection.disconnect());
2180
+ if (!connection.isHibernatable) {
2181
+ this.#rLog.debug({
2182
+ msg: "disconnecting non-hibernatable connection on actor stop",
2183
+ connId: connection.id,
2184
+ });
2185
+ promises.push(connection.disconnect());
2186
+ }
2061
2187
 
2062
- // TODO: Figure out how to abort HTTP requests on shutdown
2188
+ // TODO: Figure out how to abort HTTP requests on shutdown. This
2189
+ // might already be handled by the engine runner tunnel shutdown.
2063
2190
  }
2064
2191
 
2065
2192
  // Wait for any background tasks to finish, with timeout
@@ -2069,7 +2196,6 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
2069
2196
 
2070
2197
  // Clear timeouts
2071
2198
  if (this.#pendingSaveTimeout) clearTimeout(this.#pendingSaveTimeout);
2072
- if (this.#sleepTimeout) clearTimeout(this.#sleepTimeout);
2073
2199
  if (this.#checkConnLivenessInterval)
2074
2200
  clearInterval(this.#checkConnLivenessInterval);
2075
2201
 
@@ -2149,6 +2275,7 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
2149
2275
  eventName: sub.eventName,
2150
2276
  })),
2151
2277
  lastSeen: BigInt(conn.lastSeen),
2278
+ hibernatableRequestId: conn.hibernatableRequestId ?? null,
2152
2279
  })),
2153
2280
  scheduledEvents: persist.scheduledEvents.map((event) => ({
2154
2281
  eventId: event.eventId,
@@ -2187,6 +2314,7 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
2187
2314
  eventName: sub.eventName,
2188
2315
  })),
2189
2316
  lastSeen: Number(conn.lastSeen),
2317
+ hibernatableRequestId: conn.hibernatableRequestId ?? undefined,
2190
2318
  })),
2191
2319
  scheduledEvents: bareData.scheduledEvents.map((event) => ({
2192
2320
  eventId: event.eventId,
@@ -16,8 +16,11 @@ export interface PersistedConn<CP, CS> {
16
16
  state: CS;
17
17
  subscriptions: PersistedSubscription[];
18
18
 
19
- /** Last time the socket was seen. This is set when disconencted so we can determine when we need to clean this up. */
19
+ /** Last time the socket was seen. This is set when disconnected so we can determine when we need to clean this up. */
20
20
  lastSeen: number;
21
+
22
+ /** Request ID of the hibernatable WebSocket. See PersistedActor.hibernatableWebSocket */
23
+ hibernatableRequestId?: ArrayBuffer;
21
24
  }
22
25
 
23
26
  export interface PersistedSubscription {
@@ -116,6 +116,7 @@ export async function handleWebSocketConnect(
116
116
  encoding: Encoding,
117
117
  parameters: unknown,
118
118
  requestId: string,
119
+ requestIdBuf: ArrayBuffer | undefined,
119
120
  connId: string | undefined,
120
121
  connToken: string | undefined,
121
122
  ): Promise<UpgradeWebSocketArgs> {
@@ -184,9 +185,19 @@ export async function handleWebSocketConnect(
184
185
  actorId,
185
186
  });
186
187
 
188
+ // Check if this is a hibernatable websocket
189
+ const isHibernatable =
190
+ !!requestIdBuf &&
191
+ actor[PERSIST_SYMBOL].hibernatableWebSocket.findIndex(
192
+ (ws) =>
193
+ arrayBuffersEqual(ws.requestId, requestIdBuf),
194
+ ) !== -1;
195
+
187
196
  conn = await actor.createConn(
188
197
  {
189
198
  requestId: requestId,
199
+ requestIdBuf: requestIdBuf,
200
+ hibernatable: isHibernatable,
190
201
  driverState: {
191
202
  [ConnDriverKind.WEBSOCKET]: {
192
203
  encoding,
@@ -365,6 +376,7 @@ export async function handleSseConnect(
365
376
  conn = await actor.createConn(
366
377
  {
367
378
  requestId: requestId,
379
+ hibernatable: false,
368
380
  driverState: {
369
381
  [ConnDriverKind.SSE]: {
370
382
  encoding,
@@ -479,6 +491,7 @@ export async function handleAction(
479
491
  conn = await actor.createConn(
480
492
  {
481
493
  requestId: requestId,
494
+ hibernatable: false,
482
495
  driverState: { [ConnDriverKind.HTTP]: {} },
483
496
  },
484
497
  parameters,
@@ -593,6 +606,7 @@ export async function handleRawWebSocketHandler(
593
606
  path: string,
594
607
  actorDriver: ActorDriver,
595
608
  actorId: string,
609
+ requestIdBuf: ArrayBuffer | undefined,
596
610
  ): Promise<UpgradeWebSocketArgs> {
597
611
  const actor = await actorDriver.loadActor(actorId);
598
612
 
@@ -187,6 +187,7 @@ export function createActorRouter(
187
187
  encoding,
188
188
  connParams,
189
189
  generateConnRequestId(),
190
+ undefined,
190
191
  connIdRaw,
191
192
  connTokenRaw,
192
193
  );
@@ -303,6 +304,7 @@ export function createActorRouter(
303
304
  pathWithQuery,
304
305
  actorDriver,
305
306
  c.env.actorId,
307
+ undefined,
306
308
  );
307
309
  })(c, noopNext());
308
310
  } else {
@@ -0,0 +1,48 @@
1
+ import { describe, expect, test } from "vitest";
2
+ import { isConnStatePath, isStatePath } from "./utils";
3
+
4
+ describe("isStatePath", () => {
5
+ test("matches exact state", () => {
6
+ expect(isStatePath("state")).toBe(true);
7
+ });
8
+
9
+ test("matches nested state paths", () => {
10
+ expect(isStatePath("state.foo")).toBe(true);
11
+ expect(isStatePath("state.foo.bar")).toBe(true);
12
+ });
13
+
14
+ test("does not match other paths", () => {
15
+ expect(isStatePath("connections")).toBe(false);
16
+ expect(isStatePath("stateX")).toBe(false);
17
+ expect(isStatePath("mystate")).toBe(false);
18
+ });
19
+ });
20
+
21
+ describe("isConnStatePath", () => {
22
+ test("matches connection state paths", () => {
23
+ expect(isConnStatePath("connections.0.state")).toBe(true);
24
+ expect(isConnStatePath("connections.123.state")).toBe(true);
25
+ });
26
+
27
+ test("matches nested connection state paths", () => {
28
+ expect(isConnStatePath("connections.0.state.foo")).toBe(true);
29
+ expect(isConnStatePath("connections.5.state.bar.baz")).toBe(true);
30
+ });
31
+
32
+ test("does not match non-state connection paths", () => {
33
+ expect(isConnStatePath("connections.0.params")).toBe(false);
34
+ expect(isConnStatePath("connections.0.token")).toBe(false);
35
+ expect(isConnStatePath("connections.0")).toBe(false);
36
+ });
37
+
38
+ test("does not match other paths", () => {
39
+ expect(isConnStatePath("state")).toBe(false);
40
+ expect(isConnStatePath("connections")).toBe(false);
41
+ expect(isConnStatePath("other.0.state")).toBe(false);
42
+ });
43
+
44
+ test("does not match malformed paths", () => {
45
+ expect(isConnStatePath("connections.state")).toBe(false);
46
+ expect(isConnStatePath("connections.0.stateX")).toBe(false);
47
+ });
48
+ });
@@ -104,3 +104,26 @@ export function generateRandomString(length = 32) {
104
104
  }
105
105
  return result;
106
106
  }
107
+
108
+ /**
109
+ * Checks if a path is an actor state path within the persisted actor data.
110
+ */
111
+ export function isStatePath(path: string): boolean {
112
+ return path === "state" || path.startsWith("state.");
113
+ }
114
+
115
+ /**
116
+ * Checks if a path is a connection state path within the persisted actor data.
117
+ */
118
+ export function isConnStatePath(path: string): boolean {
119
+ if (!path.startsWith("connections.")) {
120
+ return false;
121
+ }
122
+ const stateIndex = path.indexOf(".state", 12); // Start after "connections."
123
+ if (stateIndex === -1) {
124
+ return false;
125
+ }
126
+ const afterState = stateIndex + 6; // ".state".length = 6
127
+ // Check if ".state" is followed by end of string or "."
128
+ return path.length === afterState || path[afterState] === ".";
129
+ }
@@ -111,7 +111,6 @@ export class EngineActorDriver implements ActorDriver {
111
111
  );
112
112
 
113
113
  // Create runner configuration
114
- let hasDisconnected = false;
115
114
  const engineRunnerConfig: EngineRunnerConfig = {
116
115
  version: this.#version,
117
116
  endpoint: getEndpoint(runConfig),
@@ -125,32 +124,9 @@ export class EngineActorDriver implements ActorDriver {
125
124
  },
126
125
  prepopulateActorNames: buildActorNames(registryConfig),
127
126
  onConnected: () => {
128
- if (hasDisconnected) {
129
- logger().info({
130
- msg: "runner reconnected",
131
- namespace: this.#runConfig.namespace,
132
- runnerName: this.#runConfig.runnerName,
133
- });
134
- } else {
135
- logger().debug({
136
- msg: "runner connected",
137
- namespace: this.#runConfig.namespace,
138
- runnerName: this.#runConfig.runnerName,
139
- });
140
- }
141
-
142
127
  this.#runnerStarted.resolve(undefined);
143
128
  },
144
- onDisconnected: (code, reason) => {
145
- logger().warn({
146
- msg: "runner disconnected",
147
- namespace: this.#runConfig.namespace,
148
- runnerName: this.#runConfig.runnerName,
149
- code,
150
- reason,
151
- });
152
- hasDisconnected = true;
153
- },
129
+ onDisconnected: (_code, _reason) => {},
154
130
  onShutdown: () => {
155
131
  this.#runnerStopped.resolve(undefined);
156
132
  this.#isRunnerStopped = true;
@@ -196,20 +172,34 @@ export class EngineActorDriver implements ActorDriver {
196
172
  }
197
173
 
198
174
  // Check for existing WS
199
- const existingWs = handler.actor[
200
- PERSIST_SYMBOL
201
- ].hibernatableWebSocket.find((ws) =>
175
+ const hibernatableArray =
176
+ handler.actor[PERSIST_SYMBOL].hibernatableWebSocket;
177
+ logger().debug({
178
+ msg: "checking hibernatable websockets",
179
+ requestId: idToStr(requestId),
180
+ existingHibernatableWebSockets: hibernatableArray.length,
181
+ });
182
+ const existingWs = hibernatableArray.find((ws) =>
202
183
  arrayBuffersEqual(ws.requestId, requestId),
203
184
  );
204
185
 
205
186
  // Determine configuration for new WS
206
187
  let hibernationConfig: HibernationConfig;
207
188
  if (existingWs) {
189
+ logger().debug({
190
+ msg: "found existing hibernatable websocket",
191
+ requestId: idToStr(requestId),
192
+ lastMsgIndex: existingWs.msgIndex,
193
+ });
208
194
  hibernationConfig = {
209
195
  enabled: true,
210
196
  lastMsgIndex: Number(existingWs.msgIndex),
211
197
  };
212
198
  } else {
199
+ logger().debug({
200
+ msg: "no existing hibernatable websocket found",
201
+ requestId: idToStr(requestId),
202
+ });
213
203
  if (path === PATH_CONNECT_WEBSOCKET) {
214
204
  hibernationConfig = {
215
205
  enabled: true,
@@ -275,12 +265,24 @@ export class EngineActorDriver implements ActorDriver {
275
265
  }
276
266
  }
277
267
 
278
- // Save hibernatable WebSocket
279
- handler.actor[PERSIST_SYMBOL].hibernatableWebSocket.push({
280
- requestId,
281
- lastSeenTimestamp: BigInt(Date.now()),
282
- msgIndex: -1n,
283
- });
268
+ // Save or update hibernatable WebSocket
269
+ if (existingWs) {
270
+ logger().debug({
271
+ msg: "updated existing hibernatable websocket timestamp",
272
+ requestId: idToStr(requestId),
273
+ });
274
+ existingWs.lastSeenTimestamp = BigInt(Date.now());
275
+ } else {
276
+ logger().debug({
277
+ msg: "created new hibernatable websocket entry",
278
+ requestId: idToStr(requestId),
279
+ });
280
+ handler.actor[PERSIST_SYMBOL].hibernatableWebSocket.push({
281
+ requestId,
282
+ lastSeenTimestamp: BigInt(Date.now()),
283
+ msgIndex: -1n,
284
+ });
285
+ }
284
286
 
285
287
  return hibernationConfig;
286
288
  },
@@ -353,6 +355,12 @@ export class EngineActorDriver implements ActorDriver {
353
355
 
354
356
  handler.persistedData = data;
355
357
 
358
+ logger().debug({
359
+ msg: "writing persisted data for actor",
360
+ actorId,
361
+ dataSize: data.byteLength,
362
+ });
363
+
356
364
  await this.#runner.kvPut(actorId, [[KEYS.PERSIST_DATA, data]]);
357
365
  }
358
366
 
@@ -427,6 +435,13 @@ export class EngineActorDriver implements ActorDriver {
427
435
  persistedValue !== null
428
436
  ? persistedValue
429
437
  : serializeEmptyPersistData(input);
438
+
439
+ logger().debug({
440
+ msg: "loaded persisted data for actor",
441
+ actorId,
442
+ dataSize: handler.persistedData?.byteLength,
443
+ wasInStorage: persistedValue !== null,
444
+ });
430
445
  }
431
446
 
432
447
  const name = actorConfig.name as string;
@@ -547,6 +562,7 @@ export class EngineActorDriver implements ActorDriver {
547
562
  encoding,
548
563
  connParams,
549
564
  requestId,
565
+ requestIdBuf,
550
566
  // Extract connId and connToken from protocols if needed
551
567
  undefined,
552
568
  undefined,
@@ -557,6 +573,7 @@ export class EngineActorDriver implements ActorDriver {
557
573
  url.pathname + url.search,
558
574
  this,
559
575
  actorId,
576
+ requestIdBuf,
560
577
  );
561
578
  } else {
562
579
  throw new Error(`Unreachable path: ${url.pathname}`);
@@ -626,7 +643,52 @@ export class EngineActorDriver implements ActorDriver {
626
643
  }
627
644
 
628
645
  async shutdownRunner(immediate: boolean): Promise<void> {
629
- logger().info({ msg: "stopping engine actor driver" });
646
+ logger().info({ msg: "stopping engine actor driver", immediate });
647
+
648
+ // TODO: We need to update the runner to have a draining state so:
649
+ // 1. Send ToServerDraining
650
+ // - This causes Pegboard to stop allocating actors to this runner
651
+ // 2. Pegboard sends ToClientStopActor for all actors on this runner which handles the graceful migration of each actor independently
652
+ // 3. Send ToServerStopping once all actors have successfully stopped
653
+ //
654
+ // What's happening right now is:
655
+ // 1. All actors enter stopped state
656
+ // 2. Actors still respond to requests because only RivetKit knows it's
657
+ // stopping, this causes all requests to issue errors that the actor is
658
+ // stopping. (This will NOT return a 503 bc the runner has no idea the
659
+ // actors are stopping.)
660
+ // 3. Once the last actor stops, then the runner finally stops + actors
661
+ // reschedule
662
+ //
663
+ // This means that:
664
+ // - All actors on this runner are bricked until the slowest _onStop finishes
665
+ // - Guard will not gracefully handle requests bc it's not receiving a 503
666
+ // - Actors can still be scheduled to this runner while the other
667
+ // actors are stopping, meaning that those actors will NOT get _onStop
668
+ // and will potentiall corrupt their state
669
+ //
670
+ // HACK: Stop all actors to allow state to be saved
671
+ // NOTE: _onStop is only supposed to be called by the runner, we're
672
+ // abusing it here
673
+ logger().debug({
674
+ msg: "stopping all actors before shutdown",
675
+ actorCount: this.#actors.size,
676
+ });
677
+ const stopPromises: Promise<void>[] = [];
678
+ for (const [_actorId, handler] of this.#actors.entries()) {
679
+ if (handler.actor) {
680
+ stopPromises.push(
681
+ handler.actor._onStop().catch((err) => {
682
+ handler.actor?.rLog.error({
683
+ msg: "_onStop errored",
684
+ error: stringifyError(err),
685
+ });
686
+ }),
687
+ );
688
+ }
689
+ }
690
+ await Promise.all(stopPromises);
691
+ logger().debug({ msg: "all actors stopped" });
630
692
 
631
693
  // Clear the ack flush interval
632
694
  if (this.#wsAckFlushInterval) {