117 lines
3.4 KiB
TypeScript
117 lines
3.4 KiB
TypeScript
import { DbConnection } from "../src/module_bindings/index.ts";
|
|
import { parseArgs } from "node:util";
|
|
|
|
const { values } = parseArgs({
|
|
strict: false,
|
|
options: {
|
|
host: { type: "string", default: "ws://localhost:3000" },
|
|
db: { type: "string", default: "zep" },
|
|
channels: { type: "string" },
|
|
rate: { type: "string", default: "10" },
|
|
token: { type: "string" },
|
|
message: { type: "string" },
|
|
},
|
|
});
|
|
|
|
if (values.channels === undefined) {
|
|
console.error("Error: --channels is required (comma separated list of channel IDs, e.g., --channels=1,2,3)");
|
|
process.exit(1);
|
|
}
|
|
|
|
const host = String(values.host);
|
|
const dbName = String(values.db);
|
|
const rate = parseInt(String(values.rate), 10) || 10;
|
|
const channelsRaw = typeof values.channels === 'string' ? values.channels : "";
|
|
const channels = channelsRaw.split(",").map((c: string) => {
|
|
try {
|
|
return BigInt(c.trim());
|
|
} catch {
|
|
console.error(`Invalid channel ID: ${c}`);
|
|
process.exit(1);
|
|
}
|
|
});
|
|
const token = typeof values.token === 'string' ? values.token : undefined;
|
|
const customMessage = typeof values.message === 'string' ? values.message : undefined;
|
|
|
|
console.log(`Zep Stress Client Initializing...`);
|
|
console.log(`Host: ${host}`);
|
|
console.log(`Database: ${dbName}`);
|
|
console.log(`Channels: ${channels.join(", ")}`);
|
|
console.log(`Target Rate: ${rate} messages/second`);
|
|
|
|
let connection: DbConnection | null = null;
|
|
let intervalId: ReturnType<typeof setInterval> | null = null;
|
|
let msgCount = 0;
|
|
|
|
const startStressTest = () => {
|
|
const msDelay = 1000 / rate;
|
|
console.log(`Starting message loop (1 message every ${msDelay.toFixed(2)}ms)...`);
|
|
|
|
intervalId = setInterval(() => {
|
|
if (!connection) return;
|
|
|
|
// Pick a random channel from the list
|
|
const channelId = channels[Math.floor(Math.random() * channels.length)];
|
|
const text = customMessage
|
|
? `${customMessage} [#${msgCount++}]`
|
|
: `Stress test message ${msgCount++} from client at ${new Date().toISOString()}`;
|
|
|
|
try {
|
|
// The bindings use camelCase sendMessage
|
|
connection.reducers.sendMessage({
|
|
text,
|
|
channelId,
|
|
threadId: undefined,
|
|
imageIds: [],
|
|
isEncrypted: false,
|
|
});
|
|
|
|
if (msgCount % 100 === 0) {
|
|
console.log(`Sent ${msgCount} messages...`);
|
|
}
|
|
} catch (e) {
|
|
console.error(`Failed to send message:`, e);
|
|
}
|
|
}, msDelay);
|
|
};
|
|
|
|
const builder = DbConnection.builder()
|
|
.withUri(host)
|
|
.withDatabaseName(dbName)
|
|
.onConnect((conn, identity, tokenStr) => {
|
|
connection = conn;
|
|
console.log(`Connected! Identity: ${identity.toHexString()}`);
|
|
if (tokenStr) {
|
|
console.log(`Token: ${tokenStr.substring(0, 15)}...`);
|
|
}
|
|
|
|
// New Architecture: Good practice to subscribe to our own state
|
|
console.log("Subscribing to system configuration and user state...");
|
|
|
|
// Give it a brief moment to initialize before hammering
|
|
setTimeout(startStressTest, 1000);
|
|
})
|
|
.onDisconnect(() => {
|
|
console.log(`Disconnected from server.`);
|
|
if (intervalId) clearInterval(intervalId);
|
|
process.exit(1);
|
|
})
|
|
.onConnectError((_ctx, err) => {
|
|
console.error(`Connection error:`, err);
|
|
process.exit(1);
|
|
});
|
|
|
|
if (token) {
|
|
builder.withToken(token).build();
|
|
} else {
|
|
builder.build();
|
|
}
|
|
|
|
// Graceful shutdown
|
|
process.on("SIGINT", () => {
|
|
console.log("\nShutting down stress client...");
|
|
if (intervalId) clearInterval(intervalId);
|
|
if (connection) connection.disconnect();
|
|
process.exit(0);
|
|
});
|