feat: restore history from shard and close session if dropped

+44 -2

src/lib/handlers/connect.ts

··· 1 1 import { 2 2 createNewSession, 3 + deleteSession, 3 4 issuedLatticeTokens, 4 5 isValidSession, 5 6 } from "@/lib/sessions"; 6 - import { shardMessageSchema } from "@/lib/types/messages"; 7 + import { 8 + requestHistoryMessageSchema, 9 + shardMessageSchema, 10 + } from "@/lib/types/messages"; 7 11 import type { PreHandler, WsRouteHandler } from "@/lib/types/routes"; 8 12 import { stringToAtUri } from "@/lib/utils/atproto"; 9 - import { sendToChannelClients, storeMessageInShard } from "@/lib/utils/gmstn"; 13 + import { 14 + sendHistoryRequestToShard, 15 + sendToChannelClients, 16 + storeMessageInShard, 17 + } from "@/lib/utils/gmstn"; 10 18 import { 11 19 rawDataToString, 12 20 validateWsMessageType, ··· 105 113 106 114 sendToChannelClients({ channelAtUri, message: shardMessage }); 107 115 storeMessageInShard({ channelAtUri, message: shardMessage }); 116 + break; 117 + } 118 + case "shard/requestHistory": { 119 + const { 120 + success, 121 + error, 122 + data: requestHistoryMessage, 123 + } = requestHistoryMessageSchema.safeParse( 124 + validateTypeResult.data, 125 + ); 126 + if (!success) { 127 + console.error( 128 + "could not parse", 129 + validateTypeResult.data, 130 + "as a valid history request message.", 131 + ); 132 + console.error(z.treeifyError(error)); 133 + return; 134 + } 135 + 136 + const { channel } = requestHistoryMessage; 137 + 138 + const atUriParseResult = stringToAtUri(channel); 139 + if (!atUriParseResult.ok) return; 140 + const { data: channelAtUri } = atUriParseResult; 141 + 142 + sendHistoryRequestToShard({ 143 + channelAtUri, 144 + message: requestHistoryMessage, 145 + }); 108 146 } 109 147 } 148 + }); 149 + 150 + socket.on("close", () => { 151 + deleteSession(sessionInfo); 110 152 }); 111 153 };

+54

src/lib/listeners/shard-history.ts

··· 1 + import { clientSessions } from "@/lib/sessions"; 2 + import { historyMessageSchema } from "@/lib/types/messages"; 3 + import { 4 + rawDataToString, 5 + validateWsMessageType, 6 + } from "@/lib/utils/ws/validate"; 7 + import type WebSocket from "ws"; 8 + import { z } from "zod"; 9 + 10 + export const attachHistoryFromShardListener = (socket: WebSocket) => { 11 + socket.on("message", (rawData) => { 12 + const event = rawDataToString(rawData); 13 + 14 + const data: unknown = JSON.parse(event); 15 + const validateTypeResult = validateWsMessageType(data); 16 + if (!validateTypeResult.ok) return; 17 + 18 + console.log("received", validateTypeResult.data, "from shard") 19 + 20 + const { type: messageType } = validateTypeResult.data; 21 + if (messageType !== "shard/history") return; 22 + const { 23 + success, 24 + error, 25 + data: historyMessage, 26 + } = historyMessageSchema.safeParse(validateTypeResult.data); 27 + if (!success) { 28 + console.error( 29 + "could not parse", 30 + validateTypeResult.data, 31 + "as a valid history message.", 32 + ); 33 + console.error(z.treeifyError(error)); 34 + return; 35 + } 36 + const { forClient: intendedRecipient } = historyMessage; 37 + const clientSessionInfo = clientSessions 38 + .keys() 39 + .find((sessionInfo) => sessionInfo.clientDid === intendedRecipient); 40 + if (!clientSessionInfo) { 41 + console.error("Could not client session info in sessions map."); 42 + return; 43 + } 44 + const clientSocket = clientSessions.get(clientSessionInfo); 45 + if (!clientSocket) { 46 + console.error( 47 + "Could find session info in map but somehow couldn't find socket? This should not happen.", 48 + ); 49 + return; 50 + } 51 + clientSocket.send(JSON.stringify(historyMessage)); 52 + console.log("sent off", historyMessage, "to client") 53 + }); 54 + };

+2

src/lib/types/messages.ts

··· 29 29 type: z.literal("shard/history"), 30 30 messages: z.optional(z.array(shardMessageSchema)), 31 31 channel: z.string(), 32 + forClient: didSchema, 32 33 }) 33 34 .strict(); 34 35 export type HistoryMessage = z.infer<typeof historyMessageSchema>; ··· 37 38 .safeExtend({ 38 39 type: z.literal("shard/requestHistory"), 39 40 channel: z.string(), 41 + requestedBy: didSchema, 40 42 }) 41 43 .strict(); 42 44 export type RequestHistoryMessage = z.infer<typeof requestHistoryMessageSchema>;

+40 -1

src/lib/utils/gmstn.ts

··· 1 + import { attachHistoryFromShardListener } from "@/lib/listeners/shard-history"; 1 2 import { clientSessions } from "@/lib/sessions"; 2 3 import { shardSessions } from "@/lib/state"; 3 4 import type { AtUri, Did } from "@/lib/types/atproto"; 4 5 import type { ShardSessionInfo } from "@/lib/types/handshake"; 5 - import type { ShardMessage } from "@/lib/types/messages"; 6 + import type { RequestHistoryMessage, ShardMessage } from "@/lib/types/messages"; 6 7 import { getEndpointFromDid } from "@/lib/utils/atproto"; 7 8 import WebSocket from "ws"; 8 9 ··· 22 23 endpoint.searchParams.append("token", token); 23 24 const ws = new WebSocket(endpoint); 24 25 shardSessions.set(sessionInfo, ws); 26 + attachHistoryFromShardListener(ws); 25 27 return ws; 26 28 }; 27 29 ··· 93 95 ); 94 96 }); 95 97 }; 98 + 99 + export const sendHistoryRequestToShard = ({ 100 + channelAtUri, 101 + message, 102 + }: { 103 + channelAtUri: AtUri; 104 + message: RequestHistoryMessage; 105 + }) => { 106 + const shardSessionInfo = shardSessions 107 + .keys() 108 + .find((sessionInfo) => 109 + sessionInfo.allowedChannels.some( 110 + (allowedChannel) => allowedChannel.rKey === channelAtUri.rKey, 111 + ), 112 + ); 113 + if (!shardSessionInfo) return; 114 + 115 + const shardSocket = shardSessions.get(shardSessionInfo); 116 + if (!shardSocket) { 117 + console.error( 118 + "Could find session info object in map, but socket could not be retrieved from map. Race condition?", 119 + ); 120 + return; 121 + } 122 + const messageToSendToShard = { 123 + ...message, 124 + }; 125 + if (shardSocket.readyState === WebSocket.OPEN) 126 + shardSocket.send(JSON.stringify(messageToSendToShard)); 127 + 128 + console.log( 129 + "Sent off message", 130 + message, 131 + "to shard located at", 132 + shardSocket.url, 133 + ); 134 + };

Configure Feed

Configure Feed