Use a state machine to manage session state transitions #220

jackyzha0 · 2024-06-25T00:36:49Z

Why

Keeping track of what handlers are registered when and knowing when to cleanup/register new callbacks and fields in a session's lifecycle is really difficult. Let's encode valid states via a state machine and make illegal states unrepresentable.

This PR also addresses two other problems that are heavily tied to the session lifecycle:

Message order synchronization during handshake: we used to synchronize nextExpectedSeq in the handshake which is useful to detect desyncs in server to client direction (that is, if the server was about to send a message that would cause the client to receive an out-of-order message, we should just close the session there instead of continuing). However, we never did this check in the other direction.
Relying on browser timers: we used to rely on a client-side setInterval to send heartbeat messages to the server and the server would kill the client if it didn't send heartbeats frequently enough. Turns out that Chromium heavily throttles intervals when backgrounded (source: https://source.chromium.org/chromium/chromium/src/+/main:third_party/blink/renderer/platform/scheduler/main_thread/page_scheduler_impl.h;l=54?q=kDefaultThrottledWakeUpInterval&ss=chromium%2Fchromium%2Fsrc) so a setInterval for, say 3s, could take up to 60s to actually fire. This meant connections were getting killed far too frequently. However, if we make the timers server authoritative, the client can always respond to server heartbeats as data-listener-based wakeups are not throttled :)

What changed

Implement a state machine that supports async-transitions via a callback-based system. We use a proxy to ensure that session states we've transitioned from cannot be used again.
- Each state registers what to do on state exit and on close (exiting from the state machine entirely, e.g. closing the session object) so we have consistent cleanups that don't leak. This has been the source of quite a few bugs.
Made heartbeats server authoritative.
- Clients no longer schedule their own heartbeats. Instead, it responds directly to any server-sent heartbeats.
- Servers are responsible for closing a connection on too many heartbeat misses.
- This is ok even in the case of phantom disconnects as closing the underlying TCP conn will still propagate in the case of a WebSocket transport (and UDS cannot have phantom disconnects)
Update the handshake messages
- No longer sends reconnect, this can be inferred a combination of nextExpectedSeq and nextSentSeq
- Now relies on sessionId to generate the ID on the session on the server side so the IDs of the client session and the server session match
- Now sends nextSentSeq. We do some more comparison here to check for mismatched session state before calling the handshake ok. Specifically, there are two cases where we cannot accept a session:
  - the client is in the future (client.nextSentSeq > server.ack)
  - server is in the future (server.seq > client.nextExpectedSeq)
Distinguish retriable handshake failures from fatal ones.
- Only a session state mismatch is retriable (i.e. kill the session and start over), everything else should be 'fatal' and stop retrying.
Testing infra changes
- Made constructing a transport not automatically connect by default. We had cases where we registered listeners after connecting and would miss some events.
- No longer expose connections to the transport (it doesn't manage these! it shouldn't be able to introspect connections). I pulled these out into test helpers numberOfConnections(transport) and closeAllConnections(transport) which covers our existing use cases.

Versioning

Breaking protocol change
Breaking ts/js API change

masad-frost · 2024-06-25T17:46:22Z

transport/server.ts

-    if (!parsed.payload.expectedSessionState) {
-      // TODO: remove once we have upgraded all clients.
-      ({ session, isTransparentReconnect } = this.getOrCreateSession({
-        to: parsed.from,
-        conn,
-        sessionId: parsed.payload.sessionId,
-        propagationCtx: parsed.tracing,
-      }));
-    } else if (parsed.payload.expectedSessionState.reconnect) {


I guess python-client is deployed with this change replit/river-python#32

yeah i figured we've marinated long enough

masad-frost

Makes sense at face-value, I would've thought that the transitions are automatically handled within the sessions (e.g. the state machine automatically goes from handshaking to connected when the handshaking session gets a successful handshake), but I guess the transport layer eeds to do some verification in some cases and act as a controller. I think this is fine, but that was my intuition.

transport/sessionStateMachine/SessionConnected.ts

masad-frost · 2024-06-25T18:00:16Z

transport/sessionStateMachine/common.ts

+ * │  server-side    ▼                                 │ connection drop
+ * └───────────────► 3. SessionConnected   ────────────┘ heartbeat misses
+ */
+abstract class StateMachineState {


is there a reason to split CommonSession and this?

i wanted to split out the state machine machinery and the session machinery

transport/sessionStateMachine/common.ts

transport/sessionStateMachine/stateMachine.test.ts

transport/sessionStateMachine/SessionConnecting.ts

transport/sessionStateMachine/SessionConnected.ts

transport/sessionStateMachine/SessionHandshaking.ts

transport/sessionStateMachine/SessionPendingIdentification.ts

…ect automatically

Monkatraz

not a very thorough look, mainly nitpicks/bikeshedding and questions

transport/sessionStateMachine/transitions.ts

transport/sessionStateMachine/common.ts

transport/sessionStateMachine/SessionPendingIdentification.ts

router/server.ts

transport/connection.ts

Monkatraz · 2024-07-05T21:54:11Z

transport/server.ts

+    // 2. cliest is reconnecting to an existing session but we don't have it
+    //    reject this handshake, there's nothing we can do to salvage it


maybe not for this PR, but should we have a mechanism for negotiating creating a new session instead? or is the client just expected to retry and request a new session explicitly?

i wanted to keep the handshake fixed-length (e.g. 2-way handshake for now), the intended way is to have the client retry these (we now distinguish retriable and fatal handshake responses so thats the goal there)

jackyzha0 · 2024-07-08T18:06:34Z

__tests__/fixtures/cleanup.ts

 export async function advanceFakeTimersByDisconnectGrace() {
-  for (let i = 0; i < testingSessionOptions.heartbeatsUntilDead; i++) {
-    // wait for heartbeat interval to elapse
-    await vi.runOnlyPendingTimersAsync();


turns out runOnlyPendingTimersAsync introduced incorrect fake timer behaviour so I did my best to remove it

lhchavez

round 1/2 done (pausing for lunch and continuing afterwards)

lhchavez · 2024-07-08T18:15:07Z

__tests__/cleanup.test.ts

      await advanceFakeTimersBySessionGrace();
-
-      expect(clientTransport.sessions.size).toEqual(0);
+      await waitFor(() => expect(numberOfConnections(clientTransport)).toBe(0));


consider using toEqual here for consistency (also primitives don't have identity-matching defined). same elsewhere.

i believe vitest + jest both allow toBe on primitives https://vitest.dev/api/expect.html#tobe

(it just allowed me to shorten it into a single line because toEqual pushes it to 3 😭)

rofl. i guess we should have consistency even if that causes prettier to do stuff that ain't as pretty.

__tests__/disconnects.test.ts

router/server.ts

transport/sessionStateMachine/transitions.ts

lhchavez · 2024-07-08T18:55:09Z

transport/sessionStateMachine/transitions.ts

+ */
+export const SessionStateGraph = {
+  entrypoints: {
+    NoConnection(


might be fun to mention that this is the client entrypoint and the WaitingForHandshake is the server entrypoint. in other words, these two entrypoints don't need to have the same name as the state they represent.

transport/sessionStateMachine/transitions.ts

lhchavez

go forth! please prioritize any comments about correctness / invariants. everything else can be done later.

transport/connection.ts

transport/transport.ts

lhchavez · 2024-07-08T19:53:25Z

transport/transport.ts

@@ -614,4 +200,106 @@ export abstract class Transport<ConnType extends Connection> {
  getStatus(): TransportStatus {
    return this.status;
  }
+
+  protected updateSession<S extends Session<ConnType>>(session: S): S {
+    const activeSession = this.sessions.get(session.to);


hmm i thought that we were going to simplify things so that session.to was always identical to session.id (in other words, each client creates a session id upfront and then all participants refer to that as the id for fewer invariants, and we lose the distinction between client and session IDs, maybe even connection IDs too).

not yet, i wanted to push that to a future PR (removing transport client ids entirely in favour of session ids)

lhchavez · 2024-07-08T21:07:53Z

transport/sessionStateMachine/SessionWaitingForHandshake.ts

+    this.conn.addCloseListener(listeners.onConnectionClosed);
+  }
+
+  onHandshakeData = (msg: Uint8Array) => {


how are we guaranteeing that once this is called, the handshake timeout will not fire?

great question, right now we are relying on the transport to listen for the onHandshake and then transition out of this state (which clears the handshake timeout)

i'll add a comment

other state transitions don't clear the other listeners as soon as an even is emitted so this is consistent with it (the listener of the event is expected to do the transition which does cleanup)

lhchavez · 2024-07-08T21:10:57Z

transport/sessionStateMachine/SessionHandshaking.ts

+ * - Handshaking -> NoConnection (on close)
+ * - Handshaking -> Connected (on handshake)
+ */
+export class SessionHandshaking<


if this state is only valid for clients (since the corresponding server-side change is SessionWaitingForHandshake) should we add a Client to this name and a Server to the other state?

transport/sessionStateMachine/SessionConnecting.ts

lhchavez · 2024-07-08T21:14:54Z

transport/sessionStateMachine/SessionConnected.ts

+    clearInterval(this.heartbeatHandle);
+    this.heartbeatHandle = undefined;


Suggested change

clearInterval(this.heartbeatHandle);

this.heartbeatHandle = undefined;

if (this.heartbeatHandle) {

clearInterval(this.heartbeatHandle);

this.heartbeatHandle = undefined;

}

idempotency ftw

isn't clearInterval already idempotent?

that's putting a lot of trust that node / random browsers will never, ever for any circumstance accidentally reuse an interval ID.

lhchavez · 2024-07-08T21:16:48Z

transport/sessionStateMachine/SessionConnected.ts

+
+    // if we are not actively heartbeating, we are in passive
+    // heartbeat mode and should send a response to the ack
+    if (!this.isActivelyHeartbeating) {


is there a way to create a new state for this? it feels like a new state.

hm we could but i dont think it would add much value, this to me seems like a parameter of a state (like a substate?)

lhchavez · 2024-07-08T21:29:09Z

transport/client.ts

-        this.protocolError(ProtocolError.RetriesExceeded, errMsg);
-        return;
-      }
+    // check budget


is it worth having the "sleeping before reconnecting" state be distinct from "no connection"?

same as above, i thought about this as well but not sure if it gets us anything (behaviorally this is identical to SessionConnecting)

## Why We removed the client being required to send heartbeat on interval due to browser timer throttling in #220. That makes sense, since we don't want the server to expect anything from the client that relies on client timers, making the timer server authoritative makes sense. That said, this entirely removed our ability to detect bad servers or phantom disconnects on the client. ## What changed Add a client timer to expect a heartbeat from the server every `options.heartbeatsUntilDead * options.heartbeatIntervalMs` , if we don't see a heartbeat, we close the connection. This is fine in the context of browser throttling timers since in the worst case of throttling we're holding on a connection for too long, and in the best case we can detect bad connections. ## Versioning - [ ] Breaking protocol change - [ ] Breaking ts/js API change Internal change.  --------- Co-authored-by: Jacky Zhao <j.zhao2k19@gmail.com>

jackyzha0 added 4 commits June 24, 2024 17:36

checkpoint state machine work

327e8ca

class-based approach

fbc6680

more tests

f0f00f4

complete test suite

a0999c7

masad-frost reviewed Jun 25, 2024

View reviewed changes

add timeout related callbacks, use helpers in test suite

520d4ef

masad-frost approved these changes Jun 25, 2024

View reviewed changes

jackyzha0 added 14 commits June 25, 2024 13:13

heartbeats, more tests

e65b180

standardize connection handler handling

e93fad5

move bestEffortClose to method

7debdac

address more review comments

6fdde29

oops fix the proxy

55fad09

checkpoint converting client + server to use state machine

5c60e49

finish server impl

295705b

fix all the type errors

27687a7

fix some tests

dda67dc

checkpoint tests

34d27eb

get tests to all pass

b0f6e54

fix up the tests

22e1757

rebase off main

10488c9

add extra connect statements now that the transport setup doesnt conn…

94631cf

…ect automatically

jackyzha0 changed the title ~~checkpoint state machine work~~ Use a state machine to manage session state transitions Jul 5, 2024

jackyzha0 added 2 commits July 5, 2024 13:53

update protocol docs

ae2cd47

fmt skull emoji

9d9812c

jackyzha0 marked this pull request as ready for review July 5, 2024 21:26

jackyzha0 requested a review from a team as a code owner July 5, 2024 21:26

jackyzha0 requested review from Monkatraz and lhchavez and removed request for a team July 5, 2024 21:26

fix circular import for chunk opt

c613aad

Monkatraz reviewed Jul 5, 2024

View reviewed changes

jackyzha0 added 2 commits July 5, 2024 15:35

rename SessionStateMachine -> SessionStateGraph

3c3595c

rename PendingIdentification -> WaitingForHandshake

dfd4fdf

jackyzha0 commented Jul 8, 2024

View reviewed changes

lhchavez reviewed Jul 8, 2024

View reviewed changes

lhchavez approved these changes Jul 8, 2024

View reviewed changes

jackyzha0 added 2 commits July 8, 2024 15:26

address review comments

164bb57

format lol

11e34e8

jackyzha0 merged commit 3d14dab into main Jul 8, 2024
4 checks passed

jackyzha0 deleted the jackyzha0/state-machine branch July 8, 2024 22:39

jackyzha0 mentioned this pull request Jul 9, 2024

update river impl replit/river-python#40

Merged

masad-frost mentioned this pull request Jul 11, 2024

Bring back client expecting heartbeats #231

Merged

2 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Use a state machine to manage session state transitions #220

Use a state machine to manage session state transitions #220

jackyzha0 commented Jun 25, 2024 •

edited

Loading

masad-frost Jun 25, 2024 •

edited

Loading

jackyzha0 Jun 25, 2024

masad-frost left a comment

masad-frost Jun 25, 2024

jackyzha0 Jun 25, 2024

Monkatraz left a comment

Monkatraz Jul 5, 2024

jackyzha0 Jul 5, 2024

jackyzha0 Jul 8, 2024

lhchavez left a comment

lhchavez Jul 8, 2024

jackyzha0 Jul 8, 2024

lhchavez Jul 8, 2024

lhchavez Jul 8, 2024

lhchavez left a comment

lhchavez Jul 8, 2024

jackyzha0 Jul 8, 2024 •

edited

Loading

lhchavez Jul 8, 2024

jackyzha0 Jul 8, 2024

jackyzha0 Jul 8, 2024

jackyzha0 Jul 8, 2024

lhchavez Jul 8, 2024

lhchavez Jul 8, 2024

jackyzha0 Jul 8, 2024

lhchavez Jul 9, 2024

lhchavez Jul 8, 2024

jackyzha0 Jul 8, 2024

lhchavez Jul 8, 2024

jackyzha0 Jul 8, 2024

		// 2. cliest is reconnecting to an existing session but we don't have it
		// reject this handshake, there's nothing we can do to salvage it

		clearInterval(this.heartbeatHandle);
		this.heartbeatHandle = undefined;

Use a state machine to manage session state transitions #220

Use a state machine to manage session state transitions #220

Conversation

jackyzha0 commented Jun 25, 2024 • edited Loading

Why

What changed

Versioning

masad-frost Jun 25, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

masad-frost left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Monkatraz left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

lhchavez left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

lhchavez left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jackyzha0 Jul 8, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jackyzha0 commented Jun 25, 2024 •

edited

Loading

masad-frost Jun 25, 2024 •

edited

Loading

jackyzha0 Jul 8, 2024 •

edited

Loading