Skip to content

Commit fee8bf8

Browse files
OxyjunOxyjun
authored
Updating Websocket server example (#24333)
1 parent c1705f8 commit fee8bf8

File tree

1 file changed

+101
-75
lines changed

1 file changed

+101
-75
lines changed

‎src/content/docs/durable-objects/examples/websocket-server.mdx

Lines changed: 101 additions & 75 deletions
Original file line numberDiff line numberDiff line change
@@ -22,87 +22,113 @@ WebSocket connections pin your Durable Object to memory, and so duration charges
2222

2323
<TypeScriptExample>
2424
```ts
25-
import { DurableObject } from "cloudflare:workers";
26-
27-
export interface Env {
28-
WEBSOCKET_SERVER: DurableObjectNamespace<WebSocketServer>;
29-
}
25+
import { DurableObject } from 'cloudflare:workers';
3026

3127
// Worker
3228
export default {
33-
async fetch(request, env, ctx): Promise<Response> {
34-
if (request.url.endsWith("/websocket")) {
35-
// Expect to receive a WebSocket Upgrade request.
36-
// If there is one, accept the request and return a WebSocket Response.
37-
const upgradeHeader = request.headers.get("Upgrade");
38-
if (!upgradeHeader || upgradeHeader !== "websocket") {
39-
return new Response("Durable Object expected Upgrade: websocket", {
40-
status: 426,
41-
});
42-
}
43-
44-
// This example will refer to the same Durable Object,
45-
// since the name "foo" is hardcoded.
46-
let id = env.WEBSOCKET_SERVER.idFromName("foo");
47-
let stub = env.WEBSOCKET_SERVER.get(id);
48-
49-
return stub.fetch(request);
50-
}
51-
52-
return new Response(null, {
53-
status: 400,
54-
statusText: "Bad Request",
55-
headers: {
56-
"Content-Type": "text/plain",
57-
},
58-
});
59-
},
60-
} satisfies ExportedHandler<Env>;
29+
async fetch(request: Request, env: Env, ctx: ExecutionContext): Promise<Response> {
30+
if (request.url.endsWith('/websocket')) {
31+
// Expect to receive a WebSocket Upgrade request.
32+
// If there is one, accept the request and return a WebSocket Response.
33+
const upgradeHeader = request.headers.get('Upgrade');
34+
if (!upgradeHeader || upgradeHeader !== 'websocket') {
35+
return new Response('Worker expected Upgrade: websocket', {
36+
status: 426,
37+
});
38+
}
39+
40+
if (request.method !== 'GET') {
41+
return new Response('Worker expected GET method', {
42+
status: 400,
43+
});
44+
}
45+
46+
// Since we are hard coding the Durable Object ID by providing the constant name 'foo',
47+
// all requests to this Worker will be sent to the same Durable Object instance.
48+
let id = env.WEBSOCKET_SERVER.idFromName('foo');
49+
let stub = env.WEBSOCKET_SERVER.get(id);
50+
51+
return stub.fetch(request);
52+
}
53+
54+
return new Response(
55+
`Supported endpoints:
56+
/websocket: Expects a WebSocket upgrade request`,
57+
{
58+
status: 200,
59+
headers: {
60+
'Content-Type': 'text/plain',
61+
},
62+
}
63+
);
64+
},
65+
};
6166

6267
// Durable Object
6368
export class WebSocketServer extends DurableObject {
64-
currentlyConnectedWebSockets: number;
65-
66-
constructor(ctx: DurableObjectState, env: Env) {
67-
// This is reset whenever the constructor runs because
68-
// regular WebSockets do not survive Durable Object resets.
69-
//
70-
// WebSockets accepted via the Hibernation API can survive
71-
// a certain type of eviction, but we will not cover that here.
72-
super(ctx, env);
73-
this.currentlyConnectedWebSockets = 0;
74-
}
75-
76-
async fetch(request: Request): Promise<Response> {
77-
// Creates two ends of a WebSocket connection.
78-
const webSocketPair = new WebSocketPair();
79-
const [client, server] = Object.values(webSocketPair);
80-
81-
// Calling `accept()` tells the runtime that this WebSocket is to begin terminating
82-
// request within the Durable Object. It has the effect of "accepting" the connection,
83-
// and allowing the WebSocket to send and receive messages.
84-
server.accept();
85-
this.currentlyConnectedWebSockets += 1;
86-
87-
// Upon receiving a message from the client, the server replies with the same message,
88-
// and the total number of connections with the "[Durable Object]: " prefix
89-
server.addEventListener("message", (event: MessageEvent) => {
90-
server.send(
91-
`[Durable Object] currentlyConnectedWebSockets: ${this.currentlyConnectedWebSockets}`,
92-
);
93-
});
94-
95-
// If the client closes the connection, the runtime will close the connection too.
96-
server.addEventListener("close", (cls: CloseEvent) => {
97-
this.currentlyConnectedWebSockets -= 1;
98-
server.close(cls.code, "Durable Object is closing WebSocket");
99-
});
100-
101-
return new Response(null, {
102-
status: 101,
103-
webSocket: client,
104-
});
105-
}
69+
// Keeps track of all WebSocket connections
70+
sessions: Map<WebSocket, { [key: string]: string }>;
71+
72+
constructor(ctx: DurableObjectState, env: Env) {
73+
super(ctx, env);
74+
this.sessions = new Map();
75+
}
76+
77+
async fetch(request: Request): Promise<Response> {
78+
// Creates two ends of a WebSocket connection.
79+
const webSocketPair = new WebSocketPair();
80+
const [client, server] = Object.values(webSocketPair);
81+
82+
// Calling `accept()` tells the runtime that this WebSocket is to begin terminating
83+
// request within the Durable Object. It has the effect of "accepting" the connection,
84+
// and allowing the WebSocket to send and receive messages.
85+
server.accept();
86+
87+
// Generate a random UUID for the session.
88+
const id = crypto.randomUUID();
89+
// Add the WebSocket connection to the map of active sessions.
90+
this.sessions.set(server, { id });
91+
92+
server.addEventListener('message', (event) => {
93+
this.handleWebSocketMessage(server, event.data);
94+
});
95+
96+
// If the client closes the connection, the runtime will close the connection too.
97+
server.addEventListener('close', () => {
98+
this.handleConnectionClose(server);
99+
});
100+
101+
return new Response(null, {
102+
status: 101,
103+
webSocket: client,
104+
});
105+
}
106+
107+
async handleWebSocketMessage(ws: WebSocket, message: string | ArrayBuffer) {
108+
const connection = this.sessions.get(ws)!;
109+
110+
// Reply back with the same message to the connection
111+
ws.send(`[Durable Object] message: ${message}, from: ${connection.id}`);
112+
113+
// Broadcast the message to all the connections,
114+
// except the one that sent the message.
115+
this.sessions.forEach((k, session) => {
116+
if (session !== ws) {
117+
session.send(`[Durable Object] message: ${message}, from: ${connection.id}`);
118+
}
119+
});
120+
121+
// Broadcast the message to all the connections,
122+
// including the one that sent the message.
123+
this.sessions.forEach((k, session) => {
124+
session.send(`[Durable Object] message: ${message}, from: ${connection.id}`);
125+
});
126+
}
127+
128+
async handleConnectionClose(ws: WebSocket) {
129+
this.sessions.delete(ws);
130+
ws.close(1000, 'Durable Object is closing WebSocket');
131+
}
106132
}
107133
```
108134
</TypeScriptExample>

0 commit comments

Comments
 (0)