collect-system/apps/server/src/socket/collaboration/yjs.server.ts

118 lines
3.3 KiB
TypeScript
Executable File

import { Injectable } from '@nestjs/common';
import { WebSocketType, WSClient } from '../types';
import { BaseWebSocketServer } from '../base/base-websocket-server';
import { encoding } from 'lib0';
import {
YMessageType,
writeSyncStep1,
encodeAwarenessUpdate,
} from '@nice/common';
import {
getYDoc,
closeConn,
WSSharedDoc,
messageListener,
send,
} from './ws-shared-doc';
@Injectable()
export class YjsServer extends BaseWebSocketServer {
public get serverType(): WebSocketType {
return WebSocketType.YJS;
}
public override handleConnection(connection: WSClient): void {
super.handleConnection(connection);
try {
connection.binaryType = 'arraybuffer';
const doc = this.initializeDocument(connection, connection.roomId, true);
this.setupConnectionHandlers(connection, doc);
this.sendInitialSync(connection, doc);
} catch (error: any) {
this.logger.error(
`Error in handleNewConnection: ${error.message}`,
error.stack,
);
connection.close();
}
}
private initializeDocument(conn: WSClient, docName: string, gc: boolean) {
const doc = getYDoc(docName, gc);
doc.conns.set(conn, new Set());
return doc;
}
private setupConnectionHandlers(
connection: WSClient,
doc: WSSharedDoc,
): void {
connection.on('message', (message: ArrayBuffer) => {
this.handleMessage(connection, doc, message);
});
connection.on('close', () => {
this.handleClose(doc, connection);
});
connection.on('error', (error) => {
this.logger.error(
`WebSocket error for doc ${doc.name}: ${error.message}`,
error.stack,
);
closeConn(doc, connection);
this.logger.warn(
`Connection closed due to error for doc: ${doc.name}. Remaining connections: ${doc.conns.size}`,
);
});
}
private handleClose(doc: WSSharedDoc, connection: WSClient): void {
try {
closeConn(doc, connection);
} catch (error: any) {
this.logger.error(
`Error closing connection: ${error.message}`,
error.stack,
);
}
}
private handleMessage(
connection: WSClient,
doc: WSSharedDoc,
message: ArrayBuffer,
): void {
try {
messageListener(connection, doc, new Uint8Array(message));
} catch (error: any) {
this.logger.error(
`Error handling message: ${error.message}`,
error.stack,
);
}
}
private sendInitialSync(connection: WSClient, doc: any): void {
this.sendSyncStep1(connection, doc);
this.sendAwarenessStates(connection, doc);
}
private sendSyncStep1(connection: WSClient, doc: any): void {
const encoder = encoding.createEncoder();
encoding.writeVarUint(encoder, YMessageType.Sync);
writeSyncStep1(encoder, doc);
send(doc, connection, encoding.toUint8Array(encoder));
}
private sendAwarenessStates(connection: WSClient, doc: WSSharedDoc): void {
const awarenessStates = doc.awareness.getStates();
if (awarenessStates.size > 0) {
const encoder = encoding.createEncoder();
encoding.writeVarUint(encoder, YMessageType.Awareness);
encoding.writeVarUint8Array(
encoder,
encodeAwarenessUpdate(
doc.awareness,
Array.from(awarenessStates.keys()),
),
);
send(doc, connection, encoding.toUint8Array(encoder));
}
}
}