Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
73 changes: 69 additions & 4 deletions src/handwritten/commands/drive/realtime.ts
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ export type DriveRealtimeMessage =
| { type: "library_changed"; cursor?: string; path?: string; origin_client_id?: string }
| { type: "resync_required"; cursor?: string; reason?: string }
| { type: "error"; code?: string; message?: string }
| { type: "pong" }
| { type: "unknown"; message_type?: string }

export interface DriveRealtimeConnectorInit {
Expand All @@ -17,7 +18,14 @@ export type DriveRealtimeConnector = (url: URL, handlers: {
open: () => void
message: (data: string) => void
close: (error?: unknown) => void
}, init?: DriveRealtimeConnectorInit) => { close: () => void }
}, init?: DriveRealtimeConnectorInit) => { close: () => void; send?: (data: string) => void }

// A half-open socket (NAT/edge dropped the connection without a FIN) looks
// exactly like a quiet library: without keepalive the client waits forever and
// misses every remote change. Ping the server and tear the socket down when
// nothing comes back, so the normal reconnect path takes over.
export const REALTIME_PING_INTERVAL_MS = 30_000
export const REALTIME_PONG_TIMEOUT_MS = 10_000

type DriveRealtimeHandlers = Parameters<DriveRealtimeSource["start"]>[0]

Expand All @@ -38,19 +46,65 @@ export function createDriveRealtimeSource(args: {
const cancelTimeout = args.clearTimeout ?? clearTimeout
let currentRealtime = { ...args.realtime }
let handlers: DriveRealtimeHandlers | undefined
let activeSocket: { close: () => void } | undefined
let activeSocket: { close: () => void; send?: (data: string) => void } | undefined
let reconnectTimer: ReturnType<typeof setTimeout> | undefined
let reconnectDelayMs = 1000
let stopped = false
let authFailed = false
let connectionId = 0
let pingTimer: ReturnType<typeof setTimeout> | undefined
let pongTimer: ReturnType<typeof setTimeout> | undefined

function clearReconnectTimer(): void {
if (reconnectTimer === undefined) return
cancelTimeout(reconnectTimer)
reconnectTimer = undefined
}

function clearKeepaliveTimers(): void {
if (pingTimer !== undefined) {
cancelTimeout(pingTimer)
pingTimer = undefined
}
if (pongTimer !== undefined) {
cancelTimeout(pongTimer)
pongTimer = undefined
}
}

function schedulePing(id: number): void {
const send = activeSocket?.send
if (send === undefined) return
pingTimer = scheduleTimeout(() => {
pingTimer = undefined
if (id !== connectionId || stopped || authFailed) return
try {
send(JSON.stringify({ type: "ping" }))
} catch (error) {
closeConnection(id, error)
return
}
pongTimer = scheduleTimeout(() => {
pongTimer = undefined
closeConnection(
id,
Object.assign(new Error("no traffic within realtime keepalive timeout"), {
code: "PING_TIMEOUT",
}),
)
}, REALTIME_PONG_TIMEOUT_MS)
}, REALTIME_PING_INTERVAL_MS)
}

function markAlive(id: number): void {
if (id !== connectionId) return
if (pongTimer !== undefined) {
cancelTimeout(pongTimer)
pongTimer = undefined
}
if (pingTimer === undefined) schedulePing(id)
}

async function persist(next: DriveRealtimeState): Promise<void> {
currentRealtime = next
await args.writeRealtimeState(next)
Expand Down Expand Up @@ -88,11 +142,13 @@ export function createDriveRealtimeSource(args: {
open() {
if (id !== connectionId || stopped || authFailed) return
reconnectDelayMs = 1000
schedulePing(id)
void persistBestEffort({ ...currentRealtime, last_connected_at: driveIsoTimestamp(clock) })
.then(() => handlers?.onConnected())
},
message(data) {
if (id !== connectionId || stopped || authFailed) return
markAlive(id)
void handleMessage(data).catch((error) => handlers?.onWarning?.(redactedRealtimeError(error)))
},
close(error) {
Expand All @@ -107,6 +163,7 @@ export function createDriveRealtimeSource(args: {
const socket = activeSocket
activeSocket = undefined
clearReconnectTimer()
clearKeepaliveTimers()
socket?.close()
if (isRealtimeAuthError(error)) {
authFailed = true
Expand Down Expand Up @@ -166,12 +223,16 @@ export function createDriveRealtimeSource(args: {
}
return
}
if (message.type === "pong") {
return
}
if (message.type === "error") {
const error = message.message ?? message.code ?? "realtime error"
if (isRealtimeAuthError(error) || isRealtimeAuthError(message.code)) {
authFailed = true
connectionId += 1
clearReconnectTimer()
clearKeepaliveTimers()
handlers?.onAuthFailed(redactedRealtimeError(error))
activeSocket?.close()
activeSocket = undefined
Expand All @@ -193,6 +254,7 @@ export function createDriveRealtimeSource(args: {
async close() {
stopped = true
clearReconnectTimer()
clearKeepaliveTimers()
activeSocket?.close()
activeSocket = undefined
},
Expand Down Expand Up @@ -257,6 +319,9 @@ export function parseDriveRealtimeMessage(raw: string): DriveRealtimeMessage {
...(typeof value.message === "string" ? { message: redactedRealtimeError(value.message) } : {}),
}
}
if (messageType === "pong") {
return { type: "pong" }
}
return {
type: "unknown",
...(messageType === undefined ? {} : { message_type: messageType }),
Expand Down Expand Up @@ -321,7 +386,7 @@ function isRealtimeAuthError(error: unknown): boolean {
return /\b(401|403|auth|authorization|unauthorized|forbidden)\b/i.test(String(error))
}

function nativeWebSocketConnector(url: URL, handlers: Parameters<DriveRealtimeConnector>[1], init?: DriveRealtimeConnectorInit): { close: () => void } {
function nativeWebSocketConnector(url: URL, handlers: Parameters<DriveRealtimeConnector>[1], init?: DriveRealtimeConnectorInit): { close: () => void; send: (data: string) => void } {
const WebSocketWithInit = WebSocket as unknown as {
new (url: string | URL, init?: DriveRealtimeConnectorInit): WebSocket
}
Expand All @@ -342,7 +407,7 @@ function nativeWebSocketConnector(url: URL, handlers: Parameters<DriveRealtimeCo
closeOnce(pendingError)
}
})
return { close: () => ws.close() }
return { close: () => ws.close(), send: (data: string) => ws.send(data) }
}

function webSocketCloseError(event: CloseEvent): Error | undefined {
Expand Down
4 changes: 4 additions & 0 deletions src/handwritten/commands/drive/watch.ts
Original file line number Diff line number Diff line change
Expand Up @@ -233,6 +233,7 @@ export async function runDriveWatch(root: string, options: DriveWatchOptions = {
?.start({
onConnected() {
emit({ kind: "drive_realtime_connected", library_id: state.library_id })
dbg.log("realtime_connected", {})
},
onEvent(event) {
emit(realtimeEvent(event))
Expand All @@ -245,12 +246,15 @@ export async function runDriveWatch(root: string, options: DriveWatchOptions = {
},
onReconnect(delayMs, error) {
emit({ kind: "drive_realtime_reconnecting", delay_ms: delayMs, error })
dbg.log("realtime_reconnecting", { delay_ms: delayMs, error })
},
onAuthFailed(error) {
emit({ kind: "drive_realtime_auth_failed", error: error ?? "auth failed" })
dbg.log("realtime_auth_failed", { error: error ?? "auth failed" })
},
onWarning(warning) {
emit({ kind: "drive_realtime_warning", warning })
dbg.log("realtime_warning", { warning })
},
})
.catch(stopWithError)
Expand Down
67 changes: 66 additions & 1 deletion test/handwritten/drive/realtime.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@ import {
createDriveRealtimeSource,
parseDriveRealtimeMessage,
redactedRealtimeError,
REALTIME_PING_INTERVAL_MS,
REALTIME_PONG_TIMEOUT_MS,
type DriveRealtimeConnector,
} from "../../../src/handwritten/commands/drive/realtime.js"
import type { DriveRealtimeState } from "../../../src/handwritten/commands/drive/state.js"
Expand All @@ -19,23 +21,26 @@ function fakeConnector(): DriveRealtimeConnector & {
handlers: ConnectorHandlers
init: ConnectorInit
close: ReturnType<typeof vi.fn>
send: ReturnType<typeof vi.fn>
}>
} {
const connections: Array<{
url: URL
handlers: ConnectorHandlers
init: ConnectorInit
close: ReturnType<typeof vi.fn>
send: ReturnType<typeof vi.fn>
}> = []
const connect: DriveRealtimeConnector = (url, handlers, init) => {
const connection = {
url,
handlers,
init,
close: vi.fn(),
send: vi.fn(),
}
connections.push(connection)
return { close: connection.close }
return { close: connection.close, send: connection.send }
}
return Object.assign(connect, { connections })
}
Expand Down Expand Up @@ -518,6 +523,66 @@ describe("drive realtime helpers", () => {
expect(connect.connections).toHaveLength(2)
})

it("parses pong messages and handles them silently", async () => {
expect(parseDriveRealtimeMessage(JSON.stringify({ type: "pong" }))).toEqual({ type: "pong" })

const connect = fakeConnector()
const events: unknown[] = []
const source = createDriveRealtimeSource(sourceArgs({ connect }))

await source.start(realtimeHandlers(events))
connect.connections[0]?.handlers.message(JSON.stringify({ type: "pong" }))
await flushPromises()

expect(events).toEqual([])
})

it("sends keepalive pings and reconnects when the server stops responding", async () => {
const connect = fakeConnector()
const events: unknown[] = []
const source = createDriveRealtimeSource(sourceArgs({ connect }))

await source.start(realtimeHandlers(events))
connect.connections[0]?.handlers.open()
await flushPromises()

await vi.advanceTimersByTimeAsync(REALTIME_PING_INTERVAL_MS)
expect(connect.connections[0]?.send).toHaveBeenCalledWith(JSON.stringify({ type: "ping" }))

// No pong (or any traffic) within the timeout: the half-open socket must
// be torn down and reconnected instead of waiting forever.
await vi.advanceTimersByTimeAsync(REALTIME_PONG_TIMEOUT_MS)
expect(connect.connections[0]?.close).toHaveBeenCalled()
expect(events).toContainEqual({ reconnect: 1000, error: "realtime error (PING_TIMEOUT)" })

await vi.advanceTimersByTimeAsync(1000)
expect(connect.connections).toHaveLength(2)
})

it("treats any incoming message as keepalive liveness", async () => {
const connect = fakeConnector()
const events: unknown[] = []
const source = createDriveRealtimeSource(sourceArgs({ connect }))

await source.start(realtimeHandlers(events))
connect.connections[0]?.handlers.open()
await flushPromises()

await vi.advanceTimersByTimeAsync(REALTIME_PING_INTERVAL_MS)
expect(connect.connections[0]?.send).toHaveBeenCalledTimes(1)

connect.connections[0]?.handlers.message(JSON.stringify({ type: "pong" }))
await flushPromises()

await vi.advanceTimersByTimeAsync(REALTIME_PONG_TIMEOUT_MS)
expect(connect.connections[0]?.close).not.toHaveBeenCalled()
expect(events).not.toContainEqual({ reconnect: 1000, error: "realtime error (PING_TIMEOUT)" })

// The ping loop keeps going after each round of liveness.
await vi.advanceTimersByTimeAsync(REALTIME_PING_INTERVAL_MS)
expect(connect.connections[0]?.send).toHaveBeenCalledTimes(2)
})

it("resolves headers from a provider on every connection attempt", async () => {
const connect = fakeConnector()
let token = 0
Expand Down
Loading