426 lines
13 KiB
TypeScript
426 lines
13 KiB
TypeScript
import type { WebSocket } from "ws";
|
|
import type { ControlMessage, ControlMessageType } from "@media-cluster/protocol";
|
|
import { createMessage } from "@media-cluster/protocol";
|
|
import type { CreatePlaybackSessionPayload } from "@media-cluster/shared-types";
|
|
import { prisma } from "../database/client";
|
|
import { authenticateNode } from "../nodes/routes";
|
|
import { LibrarySyncService } from "../media/sync";
|
|
import { MetadataService } from "../metadata/service";
|
|
import type { Config } from "../config";
|
|
import type {
|
|
FullLibrarySyncPayload,
|
|
HeartbeatPayload,
|
|
HelloPayload,
|
|
LibraryEventPayload,
|
|
PlaybackSessionEndedPayload,
|
|
CreatePlaybackSessionAckPayload,
|
|
} from "@media-cluster/protocol";
|
|
|
|
interface NodeConnection {
|
|
nodeId: string;
|
|
socket: WebSocket;
|
|
lastHeartbeat: Date;
|
|
}
|
|
|
|
interface PendingAck {
|
|
resolve: (ok: boolean) => void;
|
|
timer: ReturnType<typeof setTimeout>;
|
|
}
|
|
|
|
interface SyncState {
|
|
syncId: string;
|
|
fileIds: Set<string>;
|
|
}
|
|
|
|
class NodeConnectionManager {
|
|
private connections = new Map<string, NodeConnection>();
|
|
private librarySync: LibrarySyncService | null = null;
|
|
private config: Config | null = null;
|
|
private offlineCheckInterval: ReturnType<typeof setInterval> | null = null;
|
|
private pendingSessionAcks = new Map<string, PendingAck>();
|
|
private syncState = new Map<string, SyncState>();
|
|
private syncQueues = new Map<string, Promise<void>>();
|
|
|
|
init(config: Config): void {
|
|
this.config = config;
|
|
const metadata = new MetadataService(config.TMDB_API_KEY);
|
|
this.librarySync = new LibrarySyncService(metadata);
|
|
|
|
this.offlineCheckInterval = setInterval(() => {
|
|
void this.checkOfflineNodes();
|
|
}, 30_000);
|
|
}
|
|
|
|
async handleConnection(socket: WebSocket): Promise<void> {
|
|
let nodeId: string | null = null;
|
|
let authenticated = false;
|
|
|
|
const authTimeout = setTimeout(() => {
|
|
if (!authenticated) socket.close(4001, "Authentication timeout");
|
|
}, 30_000);
|
|
|
|
socket.on("message", async (data) => {
|
|
try {
|
|
const message = JSON.parse(data.toString()) as ControlMessage;
|
|
await this.handleMessage(socket, message, (id) => {
|
|
nodeId = id;
|
|
authenticated = true;
|
|
clearTimeout(authTimeout);
|
|
});
|
|
} catch (err) {
|
|
console.error("WebSocket message error:", err);
|
|
this.send(socket, createMessage("ERROR", { message: "Invalid message" }));
|
|
}
|
|
});
|
|
|
|
socket.on("close", () => {
|
|
if (nodeId) {
|
|
this.connections.delete(nodeId);
|
|
this.syncState.delete(nodeId);
|
|
void prisma.node.update({
|
|
where: { id: nodeId },
|
|
data: { status: "OFFLINE" },
|
|
});
|
|
void prisma.mediaFile.updateMany({
|
|
where: { nodeId },
|
|
data: { available: false },
|
|
});
|
|
console.log(`Node disconnected: ${nodeId}`);
|
|
}
|
|
});
|
|
}
|
|
|
|
private async handleMessage(
|
|
socket: WebSocket,
|
|
message: ControlMessage,
|
|
onAuth: (nodeId: string) => void
|
|
): Promise<void> {
|
|
if (message.protocolVersion !== 1) {
|
|
this.send(
|
|
socket,
|
|
createMessage("ERROR", {
|
|
code: "UNSUPPORTED_PROTOCOL",
|
|
message: `Unsupported protocol version: ${message.protocolVersion}`,
|
|
})
|
|
);
|
|
socket.close(4002, "Unsupported protocol");
|
|
return;
|
|
}
|
|
|
|
switch (message.type) {
|
|
case "HELLO":
|
|
await this.handleHello(socket, message.payload as HelloPayload, onAuth);
|
|
break;
|
|
case "HEARTBEAT":
|
|
await this.handleHeartbeat(message.payload as HeartbeatPayload);
|
|
this.send(socket, createMessage("HEARTBEAT_ACK", {}));
|
|
break;
|
|
case "LIBRARY_EVENT":
|
|
await this.handleLibraryEvent(
|
|
message.payload as LibraryEventPayload,
|
|
socket
|
|
);
|
|
break;
|
|
case "FULL_LIBRARY_SYNC":
|
|
await this.handleFullSync(message.payload as FullLibrarySyncPayload, socket);
|
|
break;
|
|
case "PLAYBACK_SESSION_ENDED":
|
|
await this.handlePlaybackEnded(message.payload as PlaybackSessionEndedPayload);
|
|
break;
|
|
case "CREATE_PLAYBACK_SESSION_ACK":
|
|
this.handlePlaybackSessionAck(message.payload as CreatePlaybackSessionAckPayload);
|
|
break;
|
|
default:
|
|
this.send(socket, createMessage("ERROR", { message: `Unknown type: ${message.type}` }));
|
|
}
|
|
}
|
|
|
|
private async handleHello(
|
|
socket: WebSocket,
|
|
payload: HelloPayload,
|
|
onAuth: (nodeId: string) => void
|
|
): Promise<void> {
|
|
const valid = await authenticateNode(payload.nodeId, payload.apiKey);
|
|
if (!valid) {
|
|
this.send(
|
|
socket,
|
|
createMessage("HELLO_ACK", {
|
|
accepted: false,
|
|
reason: "Invalid credentials",
|
|
heartbeatIntervalSeconds: 30,
|
|
fullScanIntervalSeconds: 21600,
|
|
})
|
|
);
|
|
socket.close(4003, "Authentication failed");
|
|
return;
|
|
}
|
|
|
|
this.syncState.delete(payload.nodeId);
|
|
this.connections.set(payload.nodeId, {
|
|
nodeId: payload.nodeId,
|
|
socket,
|
|
lastHeartbeat: new Date(),
|
|
});
|
|
onAuth(payload.nodeId);
|
|
|
|
await prisma.node.update({
|
|
where: { id: payload.nodeId },
|
|
data: {
|
|
status: "ONLINE",
|
|
lastSeen: new Date(),
|
|
version: payload.softwareVersion,
|
|
architecture: payload.architecture,
|
|
hostname: payload.hostname,
|
|
publicStreamUrl: payload.publicStreamUrl,
|
|
},
|
|
});
|
|
|
|
// Restore availability for previously indexed files until full sync arrives
|
|
await prisma.mediaFile.updateMany({
|
|
where: { nodeId: payload.nodeId },
|
|
data: { available: true },
|
|
});
|
|
|
|
console.log(`Node connected: ${payload.nodeId} (${payload.hostname})`);
|
|
|
|
this.send(
|
|
socket,
|
|
createMessage("HELLO_ACK", {
|
|
accepted: true,
|
|
heartbeatIntervalSeconds: this.config?.NODE_HEARTBEAT_INTERVAL_SECONDS ?? 30,
|
|
fullScanIntervalSeconds: 21600,
|
|
})
|
|
);
|
|
}
|
|
|
|
private async handleHeartbeat(payload: HeartbeatPayload): Promise<void> {
|
|
const conn = this.connections.get(payload.nodeId);
|
|
if (conn) conn.lastHeartbeat = new Date();
|
|
|
|
await prisma.node.update({
|
|
where: { id: payload.nodeId },
|
|
data: {
|
|
status: "ONLINE",
|
|
lastSeen: new Date(),
|
|
totalStorage: BigInt(payload.storageTotalBytes),
|
|
freeStorage: BigInt(payload.storageFreeBytes),
|
|
activeStreams: payload.activeStreams,
|
|
currentBandwidth: BigInt(payload.currentBytesPerSec),
|
|
libraryFileCount: payload.libraryFileCount,
|
|
version: payload.softwareVersion,
|
|
},
|
|
});
|
|
}
|
|
|
|
private async handleLibraryEvent(
|
|
payload: LibraryEventPayload,
|
|
socket: WebSocket
|
|
): Promise<void> {
|
|
const conn = [...this.connections.values()].find((c) => c.socket === socket);
|
|
if (!conn || !this.librarySync) return;
|
|
|
|
// Don't request another full sync while one is already in progress
|
|
if (this.syncState.has(conn.nodeId)) {
|
|
await this.librarySync.processEvents(conn.nodeId, payload.events);
|
|
this.send(socket, createMessage("LIBRARY_EVENT_ACK", { accepted: true }));
|
|
return;
|
|
}
|
|
|
|
const node = await prisma.node.findUnique({ where: { id: conn.nodeId } });
|
|
if (!node) return;
|
|
|
|
const expectedRevision = node.lastRevision + 1;
|
|
const minRevision = Math.min(...payload.events.map((e) => e.nodeRevision));
|
|
|
|
if (payload.events.length > 0 && minRevision > expectedRevision) {
|
|
this.send(socket, createMessage("FULL_LIBRARY_SYNC", { full: true }));
|
|
return;
|
|
}
|
|
|
|
await this.librarySync.processEvents(conn.nodeId, payload.events);
|
|
this.send(socket, createMessage("LIBRARY_EVENT_ACK", { accepted: true }));
|
|
}
|
|
|
|
private async handleFullSync(
|
|
payload: FullLibrarySyncPayload,
|
|
socket: WebSocket
|
|
): Promise<void> {
|
|
const conn = [...this.connections.values()].find((c) => c.socket === socket);
|
|
if (!conn || !this.librarySync) return;
|
|
|
|
const syncId = payload.syncId ?? `legacy-${Date.now()}`;
|
|
const isLast = payload.isLast ?? true;
|
|
const batchIndex = payload.batchIndex ?? 0;
|
|
|
|
let state = this.syncState.get(conn.nodeId);
|
|
if (!state || state.syncId !== syncId) {
|
|
state = { syncId, fileIds: new Set() };
|
|
this.syncState.set(conn.nodeId, state);
|
|
}
|
|
|
|
// ACK immediately so the node WS stays free for playback / next batches.
|
|
// Metadata matching can take seconds per batch and must not block the socket.
|
|
this.send(
|
|
socket,
|
|
createMessage("FULL_LIBRARY_SYNC_ACK", {
|
|
accepted: true,
|
|
syncId: payload.syncId,
|
|
batchIndex: payload.batchIndex,
|
|
})
|
|
);
|
|
|
|
const prev = this.syncQueues.get(conn.nodeId) ?? Promise.resolve();
|
|
const job = prev
|
|
.then(async () => {
|
|
if (!this.librarySync) return;
|
|
await this.librarySync.processFullSync(
|
|
conn.nodeId,
|
|
payload.nodeRevision,
|
|
payload.files,
|
|
{ isLast, fileIdsSoFar: state!.fileIds }
|
|
);
|
|
|
|
if (isLast) {
|
|
console.log(
|
|
`Full library sync complete for ${conn.nodeId}: ${state!.fileIds.size} files` +
|
|
(payload.batchCount != null ? ` (${payload.batchCount} batches)` : "")
|
|
);
|
|
this.syncState.delete(conn.nodeId);
|
|
} else if (batchIndex % 10 === 0) {
|
|
console.log(
|
|
`Full library sync progress for ${conn.nodeId}: batch ${(batchIndex ?? 0) + 1}/${payload.batchCount ?? "?"}`
|
|
);
|
|
}
|
|
})
|
|
.catch((err) => {
|
|
console.error(`Full sync failed for ${conn.nodeId}:`, err);
|
|
this.syncState.delete(conn.nodeId);
|
|
});
|
|
this.syncQueues.set(conn.nodeId, job);
|
|
}
|
|
|
|
private handlePlaybackSessionAck(payload: CreatePlaybackSessionAckPayload): void {
|
|
const pending = this.pendingSessionAcks.get(payload.sessionId);
|
|
if (!pending) return;
|
|
clearTimeout(pending.timer);
|
|
this.pendingSessionAcks.delete(payload.sessionId);
|
|
pending.resolve(payload.ok);
|
|
}
|
|
|
|
private async handlePlaybackEnded(payload: PlaybackSessionEndedPayload): Promise<void> {
|
|
await prisma.playbackSession.updateMany({
|
|
where: { id: payload.sessionId },
|
|
data: { status: "COMPLETED" },
|
|
});
|
|
}
|
|
|
|
sendToNode(nodeId: string, type: ControlMessageType, payload: unknown): boolean {
|
|
const conn = this.connections.get(nodeId);
|
|
if (!conn || conn.socket.readyState !== 1) return false;
|
|
this.send(conn.socket, createMessage(type, payload));
|
|
return true;
|
|
}
|
|
|
|
sendRescan(nodeId: string, full: boolean): boolean {
|
|
return this.sendToNode(nodeId, "RESCAN", { full });
|
|
}
|
|
|
|
createPlaybackSession(nodeId: string, payload: CreatePlaybackSessionPayload): boolean {
|
|
return this.sendToNode(nodeId, "CREATE_PLAYBACK_SESSION", payload);
|
|
}
|
|
|
|
/** Push session to node and wait for ACK (or timeout). */
|
|
async createPlaybackSessionAcked(
|
|
nodeId: string,
|
|
payload: CreatePlaybackSessionPayload,
|
|
timeoutMs = 5000
|
|
): Promise<boolean> {
|
|
if (!this.isOnline(nodeId)) return false;
|
|
|
|
const acked = new Promise<boolean>((resolve) => {
|
|
const timer = setTimeout(() => {
|
|
// Fail closed: never hand Stremio a node URL without a confirmed session.
|
|
this.pendingSessionAcks.delete(payload.sessionId);
|
|
resolve(false);
|
|
}, timeoutMs);
|
|
this.pendingSessionAcks.set(payload.sessionId, { resolve, timer });
|
|
});
|
|
|
|
const pushed = this.createPlaybackSession(nodeId, payload);
|
|
if (!pushed) {
|
|
const pending = this.pendingSessionAcks.get(payload.sessionId);
|
|
if (pending) {
|
|
clearTimeout(pending.timer);
|
|
this.pendingSessionAcks.delete(payload.sessionId);
|
|
}
|
|
return false;
|
|
}
|
|
|
|
return acked;
|
|
}
|
|
|
|
revokePlaybackSession(nodeId: string, sessionId: string): boolean {
|
|
return this.sendToNode(nodeId, "REVOKE_PLAYBACK_SESSION", { sessionId });
|
|
}
|
|
|
|
disconnect(nodeId: string): void {
|
|
const conn = this.connections.get(nodeId);
|
|
if (conn) {
|
|
conn.socket.close(4004, "Revoked");
|
|
this.connections.delete(nodeId);
|
|
this.syncState.delete(nodeId);
|
|
}
|
|
}
|
|
|
|
isOnline(nodeId: string): boolean {
|
|
const conn = this.connections.get(nodeId);
|
|
return !!conn && conn.socket.readyState === 1;
|
|
}
|
|
|
|
private send(socket: WebSocket, message: ControlMessage): void {
|
|
if (socket.readyState === 1) {
|
|
socket.send(JSON.stringify(message));
|
|
}
|
|
}
|
|
|
|
private async checkOfflineNodes(): Promise<void> {
|
|
if (!this.config) return;
|
|
const threshold = new Date(
|
|
Date.now() - this.config.NODE_OFFLINE_THRESHOLD_SECONDS * 1000
|
|
);
|
|
|
|
const stale = await prisma.node.findMany({
|
|
where: {
|
|
status: "ONLINE",
|
|
lastSeen: { lt: threshold },
|
|
revoked: false,
|
|
},
|
|
select: { id: true },
|
|
});
|
|
|
|
for (const node of stale) {
|
|
if (this.connections.has(node.id)) {
|
|
this.disconnect(node.id);
|
|
}
|
|
await prisma.node.update({
|
|
where: { id: node.id },
|
|
data: { status: "OFFLINE" },
|
|
});
|
|
await prisma.mediaFile.updateMany({
|
|
where: { nodeId: node.id },
|
|
data: { available: false },
|
|
});
|
|
}
|
|
}
|
|
|
|
shutdown(): void {
|
|
if (this.offlineCheckInterval) clearInterval(this.offlineCheckInterval);
|
|
for (const conn of this.connections.values()) {
|
|
conn.socket.close(1001, "Server shutting down");
|
|
}
|
|
this.connections.clear();
|
|
this.syncState.clear();
|
|
}
|
|
}
|
|
|
|
export const nodeConnectionManager = new NodeConnectionManager();
|