this repo has no description
we're getting somewhere
+611
-328
+8
package.json
+8
package.json
···
12
12
"lint:fix": "bunx eslint --fix ."
13
13
},
14
14
"dependencies": {
15
+
"@christiangalsterer/node-postgres-prometheus-exporter": "^1.3.0",
15
16
"@types/big-json": "^3.2.5",
16
17
"@types/cli-progress": "^3.11.6",
18
+
"@types/express": "^5.0.0",
19
+
"@types/pg": "^8.11.10",
17
20
"@types/terminal-kit": "^2.5.6",
18
21
"axios": "^1.7.7",
19
22
"big-json": "^3.2.0",
20
23
"cli-progress": "^3.12.0",
21
24
"dotenv": "^16.4.5",
25
+
"express": "^4.21.1",
22
26
"ky": "^1.7.2",
27
+
"kysely": "^0.27.4",
23
28
"libsql": "^0.4.6",
24
29
"p-limit": "^6.1.0",
30
+
"pg": "^8.13.0",
31
+
"pg-native": "^3.2.0",
25
32
"pino": "^9.4.0",
26
33
"pino-pretty": "^11.2.2",
34
+
"prom-client": "^15.1.3",
27
35
"terminal-kit": "^3.1.1"
28
36
},
29
37
"devDependencies": {
+1
src/constants.ts
+1
src/constants.ts
+3
-2
src/helpers.ts
+3
-2
src/helpers.ts
···
19
19
}
20
20
21
21
try {
22
+
// eslint-disable-next-line @typescript-eslint/no-unused-vars
22
23
const url = new URL(`https://${pds}/`); // if this still throws, the PDS name is invalid
23
24
return pds;
24
25
} catch (error) {
···
43
44
const res = await ky.get(`https://${pds}/xrpc/com.atproto.server.describeServer`, {
44
45
timeout: 30000,
45
46
retry: {
46
-
limit: 5,
47
+
limit: 3,
47
48
statusCodes: [429, 500, 502, 503, 504],
48
49
},
49
50
});
···
52
53
53
54
return data.availableUserDomains !== undefined;
54
55
} catch (error) {
55
-
logger.error(`Error checking health for PDS ${pds}: ${error}`);
56
+
console.error(`Error checking health for PDS ${pds}: ${error}`);
56
57
return false;
57
58
}
58
59
}
+17
-315
src/index.ts
+17
-315
src/index.ts
···
1
-
import axios from 'axios';
2
1
import 'dotenv/config';
3
-
import Database from 'libsql';
4
-
import fs from 'node:fs/promises';
5
-
import { performance } from 'node:perf_hooks';
6
-
import pLimit from 'p-limit';
7
-
import readline from 'readline';
8
2
9
3
import {
10
-
DATA_OUTPUT_FILE,
11
4
DIDS_TO_PROCESS,
12
-
HEALTH_CHECK_FILE,
13
-
PDS_DATA_FETCH_CONCURRENCY,
14
-
PDS_HEALTH_CHECK_CONCURRENCY,
15
-
PLC_DB_PATH,
16
-
RELAY_URL,
17
-
SQL_OUTPUT_FILE,
18
5
} from './constants.js';
19
-
import { isPDSHealthy, sanitizePDSName } from './helpers.js';
20
6
import logger from './logger.js';
21
-
import { BskyData, DIDsFromDB, PDSDIDGrouped, PDSHealthStatus } from './types.js';
22
-
23
-
const db = new Database(PLC_DB_PATH);
24
-
25
-
async function fetchAndDumpDidsPdses() {
26
-
logger.info('Fetching DIDs from database');
27
-
const startTime = performance.now();
28
-
29
-
const didquery = db.prepare(`
30
-
SELECT
31
-
identity.did,
32
-
atproto_pds.endpoint
33
-
FROM
34
-
identity
35
-
JOIN
36
-
plc_log ON identity.identity_id = plc_log.identity
37
-
JOIN
38
-
atproto_pds ON plc_log.atproto_pds = atproto_pds.pds_id
39
-
WHERE
40
-
plc_log.entry_id IN (
41
-
SELECT MAX(entry_id)
42
-
FROM plc_log
43
-
GROUP BY identity
44
-
)
45
-
ORDER BY identity.did ASC
46
-
`);
47
-
48
-
const plcDbFile = await fs.open(SQL_OUTPUT_FILE, 'w');
49
-
const plcDbWriteStream = plcDbFile.createWriteStream();
50
-
51
-
let count = 0;
52
-
let lastLogTime = performance.now();
53
-
for (const row of didquery.iterate()) {
54
-
count += 1;
55
-
if (count % 1000000 === 0) {
56
-
const currentTime = performance.now();
57
-
const elapsedTime = currentTime - lastLogTime;
58
-
const recordsPerSecond = 1000000 / (elapsedTime / 1000);
59
-
logger.info(`Processed ${count} DIDs (${recordsPerSecond.toFixed(2)} records/sec)`);
60
-
lastLogTime = currentTime;
61
-
}
62
-
const { did, endpoint } = row as DIDsFromDB;
63
-
const processedEndpoint = endpoint
64
-
.replace(/^(https?:\/\/)/, '')
65
-
.replace(/\/+$/, '')
66
-
.trim();
67
-
const finalEndpoint =
68
-
processedEndpoint.includes('bsky.social') || processedEndpoint.includes('bsky.network') ?
69
-
RELAY_URL
70
-
: processedEndpoint;
71
-
72
-
plcDbWriteStream.write(JSON.stringify({ did, pds: finalEndpoint }) + '\n');
73
-
}
74
-
75
-
plcDbWriteStream.close();
76
-
const endTime = performance.now();
77
-
const totalTime = (endTime - startTime) / 1000;
78
-
const averageSpeed = count / totalTime;
79
-
logger.info(`Data dumped to ${SQL_OUTPUT_FILE}`);
80
-
logger.info(`Total DIDs processed: ${count}`);
81
-
logger.info(`Total time: ${totalTime.toFixed(2)} seconds`);
82
-
logger.info(`Average speed: ${averageSpeed.toFixed(2)} DIDs/second`);
83
-
}
84
-
85
-
async function checkAllPDSHealth() {
86
-
const startTime = performance.now();
87
-
logger.info('Starting to process data from file');
88
-
89
-
const readFile = await fs.open(SQL_OUTPUT_FILE, 'r');
90
-
const fileStream = readFile.createReadStream();
91
-
const rl = readline.createInterface({
92
-
input: fileStream,
93
-
crlfDelay: Infinity,
94
-
});
95
-
96
-
const groupedByPDS: PDSDIDGrouped = {};
97
-
98
-
let lineCount = 0;
99
-
let lastLogTime = performance.now();
100
-
for await (const line of rl) {
101
-
lineCount++;
102
-
if (lineCount % 1000000 === 0) {
103
-
const currentTime = performance.now();
104
-
const elapsedTime = currentTime - lastLogTime;
105
-
const linesPerSecond = 1000000 / (elapsedTime / 1000);
106
-
logger.info(`Processed ${lineCount} lines (${linesPerSecond.toFixed(2)} lines/sec)`);
107
-
lastLogTime = currentTime;
108
-
}
109
-
110
-
const { did, pds } = JSON.parse(line) as { did: string; pds: string };
111
-
// eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
112
-
if (!groupedByPDS[pds]) {
113
-
groupedByPDS[pds] = [];
114
-
}
115
-
groupedByPDS[pds].push(did);
116
-
}
117
-
118
-
logger.info(`Finished processing file. Total lines processed: ${lineCount}`);
119
-
120
-
let pdsHealthStatus: PDSHealthStatus = {};
121
-
122
-
const healthCheckStartTime = performance.now();
123
-
try {
124
-
const healthData = await fs.readFile(HEALTH_CHECK_FILE, 'utf-8');
125
-
pdsHealthStatus = JSON.parse(healthData) as PDSHealthStatus;
126
-
logger.info('Loaded health check data from file');
127
-
} catch {
128
-
logger.info('No existing health check data found, performing health checks');
129
-
130
-
const limit = pLimit(PDS_HEALTH_CHECK_CONCURRENCY);
131
-
132
-
logger.info('Checking PDS health status');
133
-
134
-
const sanitizedPDSMap = new Set<string>();
135
-
let failedCount = 0;
136
-
const originalPDSCount = Object.keys(groupedByPDS).length;
137
-
138
-
for (const pds of Object.keys(groupedByPDS)) {
139
-
try {
140
-
const sanitizedPDS = sanitizePDSName(pds);
141
-
sanitizedPDSMap.add(sanitizedPDS);
142
-
} catch {
143
-
failedCount++;
144
-
}
145
-
}
146
-
147
-
logger.info(`Sanitization removed ${failedCount} invalid PDSes out of ${originalPDSCount}`);
148
-
149
-
const healthCheckPromises = Array.from(sanitizedPDSMap.entries()).map(([sanitizedPDS]) =>
150
-
limit(async () => {
151
-
const healthy = await isPDSHealthy(sanitizedPDS);
152
-
pdsHealthStatus[sanitizedPDS] = healthy;
153
-
logger.info(`PDS ${sanitizedPDS} is healthy: ${healthy}`);
154
-
}),
155
-
);
156
-
157
-
await Promise.all(healthCheckPromises);
158
-
159
-
await fs.writeFile(HEALTH_CHECK_FILE, JSON.stringify(pdsHealthStatus, null, 2));
160
-
logger.info('Health check data saved to file');
161
-
}
162
-
const healthCheckEndTime = performance.now();
163
-
const healthCheckTime = (healthCheckEndTime - healthCheckStartTime) / 1000;
164
-
logger.info(`Health check process took ${healthCheckTime.toFixed(2)} seconds`);
165
-
166
-
const PDSCount = Object.keys(groupedByPDS).length;
167
-
const healthyCount = Object.values(pdsHealthStatus).filter(Boolean).length;
168
-
const unhealthyCount = Object.values(pdsHealthStatus).length - healthyCount;
169
-
logger.info(`Total PDS count: ${PDSCount}`);
170
-
logger.info(`Healthy PDS count: ${healthyCount}`);
171
-
logger.info(`Unhealthy PDS count: ${unhealthyCount}`);
172
-
173
-
const endTime = performance.now();
174
-
const totalTime = (endTime - startTime) / 1000;
175
-
logger.info(`Total processing time: ${totalTime.toFixed(2)} seconds`);
176
-
177
-
return { groupedByPDS, pdsHealthStatus };
178
-
}
179
-
180
-
async function processDidsAndFetchData(dids: { did: string; pds: string }[]) {
181
-
const limit = pLimit(PDS_DATA_FETCH_CONCURRENCY);
182
-
const fetchedData: BskyData[] = [];
183
-
let successfulRequests = 0;
184
-
let unsuccessfulRequests = 0;
185
-
let successfulDids = 0;
186
-
let failedDids = 0;
187
-
188
-
const tasks = dids.map(({ did, pds }) =>
189
-
limit(async () => {
190
-
try {
191
-
// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment, @typescript-eslint/no-unsafe-call, @typescript-eslint/no-unsafe-member-access
192
-
const res = await axios.post(
193
-
'http://localhost:8000/fetch',
194
-
{ did, pds },
195
-
{
196
-
responseType: 'stream',
197
-
timeout: 60000,
198
-
},
199
-
);
200
-
201
-
successfulRequests++;
202
-
203
-
await new Promise<void>((resolve, reject) => {
204
-
let buffer = '';
205
-
let didSucceeded = false;
206
-
207
-
// eslint-disable-next-line @typescript-eslint/no-unsafe-call, @typescript-eslint/no-unsafe-member-access
208
-
res.data.on('data', (chunk: Buffer) => {
209
-
buffer += chunk.toString();
210
-
let boundary = buffer.indexOf('\n');
211
-
while (boundary !== -1) {
212
-
const line = buffer.substring(0, boundary);
213
-
buffer = buffer.substring(boundary + 1);
214
-
if (line.trim()) {
215
-
try {
216
-
const json = JSON.parse(line) as BskyData;
217
-
fetchedData.push(json);
218
-
didSucceeded = true;
219
-
} catch (err) {
220
-
logger.error(`JSON parse error for DID ${did}: ${(err as Error).message}`);
221
-
}
222
-
}
223
-
boundary = buffer.indexOf('\n');
224
-
}
225
-
});
226
-
227
-
// eslint-disable-next-line @typescript-eslint/no-unsafe-call, @typescript-eslint/no-unsafe-member-access
228
-
res.data.on('end', () => {
229
-
if (buffer.trim()) {
230
-
try {
231
-
const json = JSON.parse(buffer) as BskyData;
232
-
fetchedData.push(json);
233
-
didSucceeded = true;
234
-
} catch (err) {
235
-
logger.error(`JSON parse error at stream end for DID ${did}: ${(err as Error).message}`);
236
-
}
237
-
}
238
-
if (didSucceeded) {
239
-
successfulDids++;
240
-
} else {
241
-
failedDids++;
242
-
}
243
-
resolve();
244
-
});
245
-
246
-
// eslint-disable-next-line @typescript-eslint/no-unsafe-call, @typescript-eslint/no-unsafe-member-access
247
-
res.data.on('error', (err: Error) => {
248
-
logger.error(`Stream error for DID ${did}: ${err.message}`);
249
-
failedDids++;
250
-
reject(err);
251
-
});
252
-
});
253
-
return;
254
-
} catch (error) {
255
-
logger.error(`Error fetching data for DID ${did}: ${(error as Error).message}`);
256
-
unsuccessfulRequests++;
257
-
failedDids++;
258
-
}
259
-
}),
260
-
);
7
+
import { fetchAndDumpDidsPdses } from './stages/stage1.js';
8
+
import { processDidsAndFetchData } from './stages/stage3.js';
9
+
import { checkAllPDSHealth, selectAllDids } from './stages/stage2.js';
261
10
262
-
await Promise.all(tasks);
263
-
logger.info(`Fetched data for ${fetchedData.length} DIDs.`);
264
-
logger.info(`Successful requests: ${successfulRequests}, Unsuccessful requests: ${unsuccessfulRequests}`);
265
-
logger.info(`Successful DIDs: ${successfulDids}, Failed DIDs: ${failedDids}`);
266
11
267
-
const writeFile = await fs.open(DATA_OUTPUT_FILE, 'w');
268
-
const writeStream = writeFile.createWriteStream();
269
-
for (const data of fetchedData) {
270
-
writeStream.write(JSON.stringify(data) + '\n');
271
-
}
272
-
writeStream.close();
273
-
logger.info(`Streamed fetched data to ${DATA_OUTPUT_FILE}`);
274
-
return fetchedData;
275
-
}
276
12
277
13
async function main() {
278
-
if (
279
-
!(await fs
280
-
.access(SQL_OUTPUT_FILE)
281
-
.then(() => true)
282
-
.catch(() => false))
283
-
) {
284
-
await fetchAndDumpDidsPdses();
285
-
}
14
+
// stage 1
15
+
await fetchAndDumpDidsPdses();
16
+
17
+
// stage 2
286
18
const { groupedByPDS, pdsHealthStatus } = await checkAllPDSHealth();
287
-
288
-
const healthyGroupedByPDS = Object.entries(groupedByPDS).reduce<PDSDIDGrouped>((acc, [pds, dids]) => {
289
-
if (pdsHealthStatus[pds]) {
290
-
acc[pds] = dids;
291
-
}
292
-
return acc;
293
-
}, {});
19
+
const allDids: { did: string; pds: string; }[] = selectAllDids(groupedByPDS, pdsHealthStatus);
294
20
295
-
const unhealthyGroupedByPDS = Object.entries(groupedByPDS).reduce<PDSDIDGrouped>((acc, [pds, dids]) => {
296
-
if (!pdsHealthStatus[pds]) {
297
-
acc[pds] = dids;
298
-
}
299
-
return acc;
300
-
}, {});
301
-
302
-
const totalHealthyDIDs = Object.values(healthyGroupedByPDS).flat().length;
303
-
const totalUnhealthyDIDs = Object.values(unhealthyGroupedByPDS).flat().length;
304
-
305
-
logger.info(`Total DIDs from healthy PDSes: ${totalHealthyDIDs}`);
306
-
logger.info(`Total DIDs from unhealthy PDSes: ${totalUnhealthyDIDs}`);
307
-
308
-
// Prepare the list of DIDs to process
309
-
const allDids: { did: string; pds: string }[] = [];
310
-
311
-
for (const [pds, dids] of Object.entries(healthyGroupedByPDS)) {
312
-
for (const did of dids) {
313
-
allDids.push({ did, pds });
314
-
}
315
-
}
316
-
317
-
logger.info(`Total DIDs from DB: ${allDids.length}`);
318
-
21
+
// stage 3
319
22
const didsToProcess = allDids.slice(0, DIDS_TO_PROCESS);
320
-
logger.info(`Processing the first ${didsToProcess.length} DIDs for testing.`);
23
+
console.log(`Processing ${didsToProcess.length} DIDs`);
321
24
322
25
const fetchedData = await processDidsAndFetchData(didsToProcess);
323
26
324
-
logger.info(`Fetched data array length: ${fetchedData.length}`);
27
+
console.log(`Fetched data array length: ${fetchedData.length}`);
325
28
}
326
29
327
30
main().catch((error: unknown) => {
328
-
logger.error(`Unexpected error: ${error instanceof Error ? error.message : String(error)}`);
31
+
console.error(`Unexpected error: ${error instanceof Error ? error.message : String(error)}`);
329
32
process.exit(1);
330
33
});
331
34
332
-
// Keep the existing shutdown logic
333
-
334
35
let isShuttingDown = false;
36
+
335
37
336
38
async function shutdown() {
337
39
if (isShuttingDown) {
338
-
logger.info('Shutdown called but one is already in progress.');
40
+
console.log('Shutdown called but one is already in progress.');
339
41
return;
340
42
}
341
43
342
44
isShuttingDown = true;
343
45
344
-
logger.info('Shutting down gracefully...');
46
+
console.log('Shutting down gracefully...');
345
47
346
48
process.exit(0);
347
49
}
348
50
349
51
process.on('SIGINT', () => {
350
52
shutdown().catch((error: unknown) => {
351
-
logger.error(`Shutdown failed: ${(error as Error).message}`);
53
+
console.error(`Shutdown failed: ${(error as Error).message}`);
352
54
process.exit(1);
353
55
});
354
56
});
355
57
356
58
process.on('SIGTERM', () => {
357
59
shutdown().catch((error: unknown) => {
358
-
logger.error(`Shutdown failed: ${(error as Error).message}`);
60
+
console.error(`Shutdown failed: ${(error as Error).message}`);
359
61
process.exit(1);
360
62
});
361
63
});
+38
src/metrics.ts
+38
src/metrics.ts
···
1
+
// import { monitorPgPool } from '@christiangalsterer/node-postgres-prometheus-exporter';
2
+
import express from 'express';
3
+
import { Registry, collectDefaultMetrics } from 'prom-client';
4
+
5
+
import logger from './logger.js';
6
+
// import { pool } from './postgres.js';
7
+
8
+
const register = new Registry();
9
+
collectDefaultMetrics({ register });
10
+
11
+
// monitorPgPool(pool, register);
12
+
13
+
const app = express();
14
+
15
+
app.get('/metrics', (req, res) => {
16
+
register
17
+
.metrics()
18
+
.then((metrics) => {
19
+
res.set('Content-Type', register.contentType);
20
+
res.send(metrics);
21
+
})
22
+
.catch((ex: unknown) => {
23
+
console.error(`Error serving metrics: ${(ex as Error).message}`);
24
+
res.status(500).end((ex as Error).message);
25
+
});
26
+
});
27
+
28
+
export const startMetricsServer = (port: number, host = '127.0.0.1') => {
29
+
const server = app.listen(port, host, () => {
30
+
console.log(`Metrics server listening on port ${port}`);
31
+
});
32
+
33
+
server.on('close', () => {
34
+
console.log('Metrics server closed.');
35
+
});
36
+
37
+
return server;
38
+
};
+24
src/postgres.ts
+24
src/postgres.ts
···
1
+
import { Kysely, PostgresDialect } from 'kysely';
2
+
import pg from 'pg';
3
+
4
+
import type { DB } from './schema.js';
5
+
6
+
// const { native } = pg;
7
+
const { Pool } = pg;
8
+
9
+
const pool = new Pool({
10
+
connectionString: process.env.DATABASE_URL,
11
+
max: 100,
12
+
// we need these two, even though they are superflous
13
+
// because the monitoring library depends on them
14
+
port: parseInt(process.env.DATABASE_PORT ?? '5432', 10),
15
+
host: process.env.DATABASE_HOST ?? 'localhost',
16
+
});
17
+
18
+
const dialect = new PostgresDialect({ pool });
19
+
20
+
const db = new Kysely<DB>({
21
+
dialect,
22
+
});
23
+
24
+
export { pool, db };
+36
src/schema.d.ts
+36
src/schema.d.ts
···
1
+
/**
2
+
* This file was generated by kysely-codegen.
3
+
* Please do not edit it manually.
4
+
*/
5
+
6
+
import type { ColumnType } from "kysely";
7
+
8
+
export type Generated<T> = T extends ColumnType<infer S, infer I, infer U>
9
+
? ColumnType<S, I | undefined, U>
10
+
: ColumnType<T, T | undefined, T>;
11
+
12
+
export type Timestamp = ColumnType<Date, Date | string>;
13
+
14
+
export interface Emojis {
15
+
created_at: Generated<Timestamp>;
16
+
emoji: string;
17
+
id: Generated<number>;
18
+
lang: string;
19
+
post_id: number;
20
+
}
21
+
22
+
export interface Posts {
23
+
cid: string;
24
+
created_at: Generated<Timestamp>;
25
+
did: string;
26
+
has_emojis: Generated<boolean>;
27
+
id: Generated<number>;
28
+
langs: Generated<string[]>;
29
+
rkey: string;
30
+
text: string | null;
31
+
}
32
+
33
+
export interface DB {
34
+
emojis: Emojis;
35
+
posts: Posts;
36
+
}
+75
src/stages/stage1.ts
+75
src/stages/stage1.ts
···
1
+
import 'dotenv/config';
2
+
import Database from 'libsql';
3
+
import fs from 'node:fs/promises';
4
+
import { performance } from 'node:perf_hooks';
5
+
6
+
import { PLC_DB_PATH, RELAY_URL, SQL_OUTPUT_FILE } from '../constants.js';
7
+
import logger from '../logger.js';
8
+
import { DidAndPds } from '../types.js';
9
+
10
+
const didDb = new Database(PLC_DB_PATH);
11
+
12
+
export async function fetchAndDumpDidsPdses() {
13
+
try {
14
+
await fs.access(SQL_OUTPUT_FILE);
15
+
console.log(`${SQL_OUTPUT_FILE} already exists. Skipping DID fetching.`);
16
+
return;
17
+
} catch {
18
+
console.log('fetchAndDumpDidsPdses');
19
+
console.log('Fetching DIDs from database');
20
+
const startTime = performance.now();
21
+
22
+
const didquery = didDb.prepare(`
23
+
SELECT
24
+
identity.did as did,
25
+
atproto_pds.endpoint as pds
26
+
FROM
27
+
identity
28
+
JOIN
29
+
plc_log ON identity.identity_id = plc_log.identity
30
+
JOIN
31
+
atproto_pds ON plc_log.atproto_pds = atproto_pds.pds_id
32
+
WHERE
33
+
plc_log.entry_id IN (
34
+
SELECT MAX(entry_id)
35
+
FROM plc_log
36
+
GROUP BY identity
37
+
)
38
+
ORDER BY identity.did ASC
39
+
`);
40
+
41
+
const plcDbFile = await fs.open(SQL_OUTPUT_FILE, 'w');
42
+
const plcDbWriteStream = plcDbFile.createWriteStream();
43
+
44
+
let count = 0;
45
+
let lastLogTime = performance.now();
46
+
for (const row of didquery.iterate()) {
47
+
count += 1;
48
+
if (count % 1000000 === 0) {
49
+
const currentTime = performance.now();
50
+
const elapsedTime = currentTime - lastLogTime;
51
+
const recordsPerSecond = 1000000 / (elapsedTime / 1000);
52
+
console.log(`Processed ${count} DIDs (${recordsPerSecond.toFixed(2)} records/sec)`);
53
+
lastLogTime = currentTime;
54
+
}
55
+
const { did, pds } = row as DidAndPds;
56
+
const sanitizedPds = pds
57
+
.replace(/^(https?:\/\/)/, '')
58
+
.replace(/\/+$/, '')
59
+
.trim();
60
+
const finalPds =
61
+
sanitizedPds.includes('bsky.social') || sanitizedPds.includes('bsky.network') ? RELAY_URL : sanitizedPds;
62
+
63
+
plcDbWriteStream.write(JSON.stringify({ did, pds: finalPds }) + '\n');
64
+
}
65
+
66
+
plcDbWriteStream.close();
67
+
const endTime = performance.now();
68
+
const totalTime = (endTime - startTime) / 1000;
69
+
const averageSpeed = count / totalTime;
70
+
console.log(`Data dumped to ${SQL_OUTPUT_FILE}`);
71
+
console.log(`Total DIDs processed: ${count}`);
72
+
console.log(`Total time: ${totalTime.toFixed(2)} seconds`);
73
+
console.log(`Average speed: ${averageSpeed.toFixed(2)} DIDs/second`);
74
+
}
75
+
}
+137
src/stages/stage2.ts
+137
src/stages/stage2.ts
···
1
+
import { HEALTH_CHECK_FILE, PDS_HEALTH_CHECK_CONCURRENCY, SQL_OUTPUT_FILE } from "../constants.js";
2
+
import logger from "../logger.js";
3
+
import fs from "node:fs/promises";
4
+
import readline from "node:readline";
5
+
import pLimit from "p-limit";
6
+
import { isPDSHealthy, sanitizePDSName } from "../helpers.js";
7
+
import { PdsToDidsMap, PdsHealthStatus } from "../types.js";
8
+
9
+
10
+
export async function checkAllPDSHealth() {
11
+
const startTime = performance.now();
12
+
console.log('Loading DIDs from file');
13
+
14
+
const readFile = await fs.open(SQL_OUTPUT_FILE, 'r');
15
+
const fileStream = readFile.createReadStream();
16
+
const rl = readline.createInterface({
17
+
input: fileStream,
18
+
crlfDelay: Infinity,
19
+
});
20
+
21
+
const groupedByPDS: PdsToDidsMap = {};
22
+
23
+
let lineCount = 0;
24
+
let lastLogTime = performance.now();
25
+
for await (const line of rl) {
26
+
lineCount++;
27
+
if (lineCount % 1000000 === 0) {
28
+
const currentTime = performance.now();
29
+
const elapsedTime = currentTime - lastLogTime;
30
+
const linesPerSecond = 1000000 / (elapsedTime / 1000);
31
+
console.log(`Processed ${lineCount} lines (${linesPerSecond.toFixed(2)} lines/sec)`);
32
+
lastLogTime = currentTime;
33
+
}
34
+
35
+
const { did, pds } = JSON.parse(line) as { did: string; pds: string };
36
+
37
+
if (!groupedByPDS[pds]) {
38
+
groupedByPDS[pds] = [];
39
+
}
40
+
groupedByPDS[pds].push(did);
41
+
}
42
+
43
+
console.log(`Finished processing file. Total lines processed: ${lineCount}`);
44
+
45
+
let pdsHealthStatus: PdsHealthStatus = {};
46
+
47
+
const healthCheckStartTime = performance.now();
48
+
try {
49
+
const healthData = await fs.readFile(HEALTH_CHECK_FILE, 'utf-8');
50
+
pdsHealthStatus = JSON.parse(healthData) as PdsHealthStatus;
51
+
console.log('Loaded health check data from file');
52
+
} catch {
53
+
console.log('No existing health check data found, performing health checks');
54
+
55
+
const limit = pLimit(PDS_HEALTH_CHECK_CONCURRENCY);
56
+
57
+
console.log('Checking PDS health status');
58
+
59
+
const sanitizedPDSMap = new Set<string>();
60
+
let failedCount = 0;
61
+
const originalPDSCount = Object.keys(groupedByPDS).length;
62
+
63
+
for (const pds of Object.keys(groupedByPDS)) {
64
+
try {
65
+
const sanitizedPDS = sanitizePDSName(pds);
66
+
sanitizedPDSMap.add(sanitizedPDS);
67
+
} catch {
68
+
failedCount++;
69
+
}
70
+
}
71
+
72
+
console.log(`Sanitization removed ${failedCount} invalid PDSes out of ${originalPDSCount}`);
73
+
74
+
const healthCheckPromises = Array.from(sanitizedPDSMap.entries()).map(([sanitizedPDS]) =>
75
+
limit(async () => {
76
+
const healthy = await isPDSHealthy(sanitizedPDS);
77
+
pdsHealthStatus[sanitizedPDS] = healthy;
78
+
console.log(`PDS ${sanitizedPDS} is healthy: ${healthy}`);
79
+
}),
80
+
);
81
+
82
+
await Promise.all(healthCheckPromises);
83
+
84
+
await fs.writeFile(HEALTH_CHECK_FILE, JSON.stringify(pdsHealthStatus, null, 2));
85
+
console.log('Health check data saved to file');
86
+
87
+
const healthCheckEndTime = performance.now();
88
+
const healthCheckTime = (healthCheckEndTime - healthCheckStartTime) / 1000;
89
+
console.log(`Health check process took ${healthCheckTime.toFixed(2)} seconds`);
90
+
91
+
const healthyCount = Object.values(pdsHealthStatus).filter(Boolean).length;
92
+
const unhealthyCount = Object.values(pdsHealthStatus).length - healthyCount;
93
+
console.log(`Total PDS count: ${sanitizedPDSMap.size}`);
94
+
console.log(`Healthy PDS count: ${healthyCount}`);
95
+
console.log(`Unhealthy PDS count: ${unhealthyCount}`);
96
+
97
+
const endTime = performance.now();
98
+
const totalTime = (endTime - startTime) / 1000;
99
+
console.log(`Total processing time: ${totalTime.toFixed(2)} seconds`);
100
+
}
101
+
102
+
return { groupedByPDS, pdsHealthStatus };
103
+
}
104
+
105
+
export function selectAllDids(groupedByPDS: PdsToDidsMap, pdsHealthStatus: PdsHealthStatus) {
106
+
const healthyGroupedByPDS = Object.entries(groupedByPDS).reduce<PdsToDidsMap>((acc, [pds, dids]) => {
107
+
if (pdsHealthStatus[pds]) {
108
+
acc[pds] = dids;
109
+
}
110
+
return acc;
111
+
}, {});
112
+
113
+
const unhealthyGroupedByPDS = Object.entries(groupedByPDS).reduce<PdsToDidsMap>((acc, [pds, dids]) => {
114
+
if (!pdsHealthStatus[pds]) {
115
+
acc[pds] = dids;
116
+
}
117
+
return acc;
118
+
}, {});
119
+
120
+
const totalHealthyDIDs = Object.values(healthyGroupedByPDS).flat().length;
121
+
const totalUnhealthyDIDs = Object.values(unhealthyGroupedByPDS).flat().length;
122
+
123
+
console.log(`Total DIDs from healthy PDSes: ${totalHealthyDIDs}`);
124
+
console.log(`Total DIDs from unhealthy PDSes: ${totalUnhealthyDIDs}`);
125
+
126
+
// Prepare the list of DIDs to process
127
+
const allDids: { did: string; pds: string; }[] = [];
128
+
129
+
for (const [pds, dids] of Object.entries(healthyGroupedByPDS)) {
130
+
for (const did of dids!) {
131
+
allDids.push({ did, pds });
132
+
}
133
+
}
134
+
135
+
console.log(`Total DIDs from DB: ${allDids.length}`);
136
+
return allDids;
137
+
}
+259
src/stages/stage3.ts
+259
src/stages/stage3.ts
···
1
+
import axios from 'axios';
2
+
import pLimit from 'p-limit';
3
+
4
+
import { PDS_DATA_FETCH_CONCURRENCY } from '../constants.js';
5
+
import logger from '../logger.js';
6
+
import { db } from '../postgres.js';
7
+
import { BskyData, BskyPost, BskyPostData } from '../types.js';
8
+
9
+
const BATCH_SIZE = 1000;
10
+
const BATCH_TIMEOUT_MS = 1000;
11
+
12
+
interface PostData {
13
+
cid: string;
14
+
did: string;
15
+
rkey: string;
16
+
hasEmojis: boolean;
17
+
langs: string[];
18
+
post: string;
19
+
createdAt: string;
20
+
}
21
+
22
+
let postBatch: PostData[] = [];
23
+
let postBatchCount = 0;
24
+
let isBatching = false;
25
+
let batchTimer: NodeJS.Timeout | null = null;
26
+
27
+
let isShuttingDown = false;
28
+
const ongoingHandleCreates = 0;
29
+
let shutdownPromise: Promise<void> | null = null;
30
+
31
+
function createShutdownPromise(): Promise<void> {
32
+
return new Promise<void>((resolve) => {
33
+
const checkCompletion = setInterval(() => {
34
+
console.log(`Shutting down, ongoing handleCreates: ${ongoingHandleCreates}`);
35
+
if (isShuttingDown && ongoingHandleCreates === 0) {
36
+
console.log('All ongoing handleCreate operations have finished.');
37
+
clearInterval(checkCompletion);
38
+
resolve();
39
+
}
40
+
}, 50);
41
+
});
42
+
}
43
+
44
+
export function initiateShutdown(): Promise<void> {
45
+
if (!shutdownPromise) {
46
+
isShuttingDown = true;
47
+
shutdownPromise = createShutdownPromise();
48
+
}
49
+
return shutdownPromise;
50
+
}
51
+
52
+
export async function flushPostgresBatch() {
53
+
if (postBatch.length === 0) {
54
+
isBatching = false;
55
+
return;
56
+
}
57
+
58
+
const currentBatch = [...postBatch];
59
+
postBatch = [];
60
+
isBatching = false;
61
+
62
+
// concurrentPostgresInserts.inc();
63
+
64
+
try {
65
+
await db.transaction().execute(async (tx) => {
66
+
// Bulk insert posts
67
+
const insertedPosts = await tx
68
+
.insertInto('posts')
69
+
.values(
70
+
currentBatch.map((post) => ({
71
+
cid: post.cid,
72
+
did: post.did,
73
+
rkey: post.rkey,
74
+
has_emojis: post.hasEmojis,
75
+
langs: post.langs,
76
+
text: post.post,
77
+
created_at: post.createdAt,
78
+
})),
79
+
)
80
+
// .returning(['id', 'cid', 'did', 'rkey'])
81
+
.execute();
82
+
83
+
// // Map composite key to id
84
+
// const compositeKeyToIdMap = new Map<string, number>();
85
+
// insertedPosts.forEach((post) => {
86
+
// const compositeKey = `${post.cid}-${post.did}-${post.rkey}`;
87
+
// compositeKeyToIdMap.set(compositeKey, post.id);
88
+
// });
89
+
90
+
// // Prepare bulk insert for emojis
91
+
// const emojiInserts: { post_id: number; emoji: string; lang: string }[] = [];
92
+
// currentBatch.forEach((post) => {
93
+
// if (post.hasEmojis) {
94
+
// const compositeKey = `${post.cid}-${post.did}-${post.rkey}`;
95
+
// const postId = compositeKeyToIdMap.get(compositeKey);
96
+
// if (postId) {
97
+
// post.emojis.forEach((emoji) => {
98
+
// post.langs.forEach((lang) => {
99
+
// emojiInserts.push({
100
+
// post_id: postId,
101
+
// emoji: emoji,
102
+
// lang: lang,
103
+
// });
104
+
// });
105
+
// });
106
+
// }
107
+
// }
108
+
// });
109
+
110
+
// if (emojiInserts.length > 0) {
111
+
// await tx.insertInto('emojis').values(emojiInserts).execute();
112
+
// }
113
+
});
114
+
115
+
// concurrentPostgresInserts.dec();
116
+
} catch (error) {
117
+
console.error(`Error flushing PostgreSQL batch: ${(error as Error).message}`);
118
+
// Optionally, you can re-add the failed batch back to `postBatch` for retry
119
+
postBatch = currentBatch.concat(postBatch);
120
+
}
121
+
}
122
+
123
+
/**
124
+
* Schedule a batch flush after a timeout.
125
+
*/
126
+
function scheduleBatchFlush() {
127
+
if (batchTimer) {
128
+
return;
129
+
}
130
+
batchTimer = setTimeout(() => {
131
+
batchTimer = null;
132
+
void flushPostgresBatch();
133
+
}, BATCH_TIMEOUT_MS);
134
+
}
135
+
136
+
export async function processDidsAndFetchData(dids: { did: string; pds: string }[]) {
137
+
const limit = pLimit(PDS_DATA_FETCH_CONCURRENCY);
138
+
const fetchedData: BskyData[] = [];
139
+
let successfulRequests = 0;
140
+
let unsuccessfulRequests = 0;
141
+
let successfulDids = 0;
142
+
let failedDids = 0;
143
+
144
+
const tasks = dids.map(({ did, pds }) =>
145
+
limit(async () => {
146
+
try {
147
+
const res = await axios.post(
148
+
'http://localhost:8000/fetch',
149
+
{ did, pds },
150
+
{
151
+
responseType: 'stream',
152
+
timeout: 60000,
153
+
},
154
+
);
155
+
156
+
successfulRequests++;
157
+
158
+
await new Promise<void>((resolve, reject) => {
159
+
let buffer = '';
160
+
let didSucceeded = false;
161
+
162
+
// eslint-disable-next-line @typescript-eslint/no-unsafe-call, @typescript-eslint/no-unsafe-member-access
163
+
res.data.on('data', (chunk: Buffer) => {
164
+
buffer += chunk.toString();
165
+
let boundary = buffer.indexOf('\n');
166
+
while (boundary !== -1) {
167
+
const line = buffer.substring(0, boundary);
168
+
buffer = buffer.substring(boundary + 1);
169
+
if (line.trim()) {
170
+
try {
171
+
const json = JSON.parse(line) as BskyData;
172
+
postBatchCount++;
173
+
if (postBatchCount % 1000 === 0) {
174
+
process.stdout.write('.');
175
+
void flushPostgresBatch();
176
+
}
177
+
// console.log('json', json);
178
+
// eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
179
+
if (json) {
180
+
for (const [k, v] of Object.entries(json)) {
181
+
if (k.includes('app.bsky.feed.post')) {
182
+
const post = v as BskyPost;
183
+
const postData = (post.value as unknown) as BskyPostData;
184
+
postBatch.push({
185
+
cid: postData.cid,
186
+
did: did,
187
+
rkey: k.split('/')[1],
188
+
hasEmojis: false,
189
+
langs: postData.langs,
190
+
post: postData.text,
191
+
createdAt: postData.createdAt,
192
+
});
193
+
}
194
+
}
195
+
}
196
+
didSucceeded = true;
197
+
} catch (err) {
198
+
console.error(`JSON parse error for DID ${did}: ${(err as Error).message}`);
199
+
}
200
+
}
201
+
boundary = buffer.indexOf('\n');
202
+
}
203
+
});
204
+
205
+
// eslint-disable-next-line @typescript-eslint/no-unsafe-call, @typescript-eslint/no-unsafe-member-access
206
+
res.data.on('end', () => {
207
+
scheduleBatchFlush();
208
+
if (buffer.trim()) {
209
+
try {
210
+
const json = JSON.parse(buffer) as BskyData;
211
+
console.log('OMG WE GOT HERE OOPS NEED TO HANDLE THIS');
212
+
console.log(json);
213
+
didSucceeded = true;
214
+
} catch (err) {
215
+
console.error(`JSON parse error at stream end for DID ${did}: ${(err as Error).message}`);
216
+
}
217
+
}
218
+
if (didSucceeded) {
219
+
successfulDids++;
220
+
} else {
221
+
failedDids++;
222
+
}
223
+
resolve();
224
+
});
225
+
226
+
// eslint-disable-next-line @typescript-eslint/no-unsafe-call, @typescript-eslint/no-unsafe-member-access
227
+
res.data.on('error', (err: Error) => {
228
+
console.error(`Stream error for DID ${did}: ${err.message}`);
229
+
failedDids++;
230
+
reject(err);
231
+
});
232
+
});
233
+
return;
234
+
} catch (error) {
235
+
process.stdout.write('!');
236
+
// this just means the user doesn't exist anymore for whatever reason
237
+
if (!(error as Error).message.includes('Request failed with status code 502')) {
238
+
console.error(`Error with DID ${did}: ${(error as Error).message}`);
239
+
}
240
+
unsuccessfulRequests++;
241
+
failedDids++;
242
+
}
243
+
}),
244
+
);
245
+
246
+
await Promise.all(tasks);
247
+
console.log(`Fetched data for ${fetchedData.length} DIDs.`);
248
+
console.log(`Successful requests: ${successfulRequests}, Unsuccessful requests: ${unsuccessfulRequests}`);
249
+
console.log(`Successful DIDs: ${successfulDids}, Failed DIDs: ${failedDids}`);
250
+
251
+
// const writeFile = await fs.open(DATA_OUTPUT_FILE, 'w');
252
+
// const writeStream = writeFile.createWriteStream();
253
+
// for (const data of fetchedData) {
254
+
// writeStream.write(JSON.stringify(data) + '\n');
255
+
// }
256
+
// writeStream.close();
257
+
// console.log(`Streamed fetched data to ${DATA_OUTPUT_FILE}`);
258
+
return fetchedData;
259
+
}
+13
-11
src/types.ts
+13
-11
src/types.ts
···
1
-
export interface DIDsFromDB {
1
+
export interface DidAndPds {
2
2
did: string;
3
-
endpoint: string;
3
+
pds: string;
4
4
}
5
5
6
-
export type PDSDIDGrouped = Record<string, string[]>;
7
-
export type PDSHealthStatus = Record<string, boolean>;
6
+
export type PdsToDidsMap = Record<string, string[] | undefined>;
7
+
export type PdsHealthStatus = Record<string, boolean | undefined>;
8
8
9
9
export interface ServerDescription {
10
10
availableUserDomains?: string[];
···
14
14
contact?: Record<string, unknown>;
15
15
}
16
16
17
+
export interface BskyPostData {
18
+
text: string;
19
+
$type: string;
20
+
langs: string[];
21
+
createdAt: string;
22
+
cid: string;
23
+
}
24
+
17
25
export type BskyPost = Record<
18
26
string,
19
27
{
20
28
cid: string;
21
-
value: {
22
-
text: string;
23
-
$type: string;
24
-
langs: string[];
25
-
createdAt: string;
26
-
cid: string;
27
-
};
29
+
value: BskyPostData;
28
30
}
29
31
>;
30
32