setup load-testing

Author: yujonglee

.github/workflows/k6.yaml

@@ -0,0 +1,32 @@
+on:
+  workflow_dispatch:
+
+jobs:
+  load-test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: superfly/flyctl-actions/setup-flyctl@master
+
+      - run: flyctl deploy --app hyprnote-api-loadtest --strategy immediate --wait-timeout 120
+        env:
+          FLY_API_TOKEN: ${{ secrets.FLY_API_TOKEN }}
+
+      - run: flyctl scale count 2 --app hyprnote-api-loadtest --yes
+        env:
+          FLY_API_TOKEN: ${{ secrets.FLY_API_TOKEN }}
+
+      - uses: grafana/setup-k6-action@v1
+
+      - uses: grafana/run-k6-action@v1
+        with:
+          path: apps/k6/scripts/websocket/listen.js
+        env:
+          API_URL: wss://hyprnote-api-loadtest.fly.dev
+          AUTH_TOKEN: ${{ secrets.K6_AUTH_TOKEN }}
+
+      - run: flyctl scale count 0 --app hyprnote-api-loadtest --yes
+        if: always()
+        env:
+          FLY_API_TOKEN: ${{ secrets.FLY_API_TOKEN }}

apps/api/src/env.ts

@@ -17,6 +17,7 @@ export const env = createEnv({
     ASSEMBLYAI_API_KEY: z.string().min(1),
     SONIOX_API_KEY: z.string().min(1),
     POSTHOG_API_KEY: z.string().min(1),
+    OVERRIDE_AUTH: z.string().optional(),
   },
   runtimeEnv: Bun.env,
   emptyStringAsUndefined: true,

apps/api/src/index.ts

@@ -11,6 +11,7 @@ import { logger } from "hono/logger";
 
 import { env } from "./env";
 import type { AppBindings } from "./hono-bindings";
+import { loadTestOverride } from "./load-test-auth";
 import { API_TAGS, routes } from "./routes";
 import { sentryMiddleware } from "./sentry/middleware";
 import { verifyStripeWebhook } from "./stripe";
@@ -43,11 +44,11 @@ app.use("*", (c, next) => {
   return corsMiddleware(c, next);
 });
 
-app.use("/chat/completions", requireSupabaseAuth);
+app.use("/chat/completions", loadTestOverride, requireSupabaseAuth);
 app.use("/webhook/stripe", verifyStripeWebhook);
 
 if (env.NODE_ENV !== "development") {
-  app.use("/listen", requireSupabaseAuth);
+  app.use("/listen", loadTestOverride, requireSupabaseAuth);
 }
 
 app.route("/", routes);

apps/api/src/load-test-auth.ts

@@ -0,0 +1,16 @@
+import { createMiddleware } from "hono/factory";
+
+import { env } from "./env";
+
+export const loadTestOverride = createMiddleware<{
+  Variables: { supabaseUserId: string };
+}>(async (c, next) => {
+  if (env.OVERRIDE_AUTH) {
+    const token = c.req.header("Authorization")?.replace("Bearer ", "");
+    if (token === env.OVERRIDE_AUTH) {
+      c.set("supabaseUserId", "load-test-user");
+      return next();
+    }
+  }
+  return next();
+});

apps/api/src/supabase.ts

@@ -16,11 +16,11 @@ export const requireSupabaseAuth = createMiddleware<{
     return c.text("unauthorized", 401);
   }
 
+  const token = authHeader.replace("Bearer ", "");
   const supabaseClient = createClient(env.SUPABASE_URL, env.SUPABASE_ANON_KEY, {
     global: { headers: { Authorization: authHeader } },
   });
 
-  const token = authHeader.replace("Bearer ", "");
   const { data, error } = await supabaseClient.auth.getUser(token);
 
   if (error || !data.user) {

apps/k6/package.json

@@ -0,0 +1,8 @@
+{
+  "name": "@hypr/k6",
+  "private": true,
+  "scripts": {
+    "test:local": "k6 run scripts/websocket/listen.js --env API_URL=ws://localhost:4000",
+    "test:loadtest": "k6 run scripts/websocket/listen.js --env API_URL=wss://hyprnote-api-loadtest.fly.dev"
+  }
+}

apps/k6/scripts/llm.js

@@ -0,0 +1,108 @@
+import { check, sleep } from "k6";
+import http from "k6/http";
+import { Counter, Trend } from "k6/metrics";
+
+const completionRequests = new Counter("completion_requests");
+const completionSuccess = new Counter("completion_success");
+const completionLatency = new Trend("completion_latency");
+const streamingTTFB = new Trend("streaming_ttfb");
+
+export const options = {
+  vus: 5,
+  duration: "30s",
+  thresholds: {
+    completion_success: ["count > 0"],
+    checks: ["rate > 0.9"],
+    completion_latency: ["p(95) < 30000"],
+  },
+};
+
+const API_URL = __ENV.API_URL || "http://localhost:4000";
+const AUTH_TOKEN = __ENV.AUTH_TOKEN || "";
+const STREAMING = __ENV.STREAMING === "true";
+
+const PROMPTS = {
+  short: "What is 2+2?",
+  medium:
+    "Explain the concept of recursion in programming. Include a simple example and mention common use cases.",
+  long: `You are a helpful assistant analyzing a meeting transcript. Here is the context:
+
+The team discussed the following topics:
+1. Q3 roadmap planning - focusing on improving user onboarding experience
+2. Technical debt reduction - specifically around the authentication system
+3. New feature requests from enterprise customers
+4. Performance optimization for the real-time collaboration features
+5. Integration with third-party calendar services
+
+Based on this context, please provide:
+- A summary of the key discussion points
+- Action items that should be tracked
+- Any decisions that were made
+- Recommended follow-up topics for the next meeting
+
+Keep your response concise but comprehensive.`,
+};
+
+const PROMPT_KEYS = Object.keys(PROMPTS);
+
+function getRandomPrompt() {
+  const key = PROMPT_KEYS[Math.floor(Math.random() * PROMPT_KEYS.length)];
+  return { type: key, content: PROMPTS[key] };
+}
+
+function makeRequest(prompt, stream) {
+  const url = `${API_URL}/chat/completions`;
+  const payload = JSON.stringify({
+    messages: [
+      { role: "system", content: "You are a helpful assistant." },
+      { role: "user", content: prompt.content },
+    ],
+    stream,
+    max_tokens: 256,
+  });
+
+  const params = {
+    headers: {
+      "Content-Type": "application/json",
+      ...(AUTH_TOKEN ? { Authorization: `Bearer ${AUTH_TOKEN}` } : {}),
+    },
+    timeout: "60s",
+  };
+
+  const startTime = Date.now();
+  const res = http.post(url, payload, params);
+  const latency = Date.now() - startTime;
+
+  return { res, latency, promptType: prompt.type };
+}
+
+export default function () {
+  const prompt = getRandomPrompt();
+  completionRequests.add(1);
+
+  const { res, latency, promptType } = makeRequest(prompt, STREAMING);
+
+  const isSuccess = check(res, {
+    "status is 200": (r) => r.status === 200,
+    "has response body": (r) => r.body && r.body.length > 0,
+  });
+
+  if (isSuccess) {
+    completionSuccess.add(1);
+    completionLatency.add(latency);
+
+    if (STREAMING) {
+      streamingTTFB.add(res.timings.waiting);
+    }
+
+    console.log(
+      `[${promptType}] ${STREAMING ? "stream" : "sync"} completed in ${latency}ms`,
+    );
+  } else {
+    console.log(
+      `[${promptType}] failed with status ${res.status}: ${res.body}`,
+    );
+  }
+
+  sleep(1);
+}

apps/k6/scripts/stt-live.js

@@ -0,0 +1,112 @@
+import { check, sleep } from "k6";
+import http from "k6/http";
+import { Counter, Trend } from "k6/metrics";
+import ws from "k6/ws";
+
+const wsConnections = new Counter("ws_connections");
+const wsTranscripts = new Counter("ws_transcripts_received");
+const wsConnectionDuration = new Trend("ws_connection_duration");
+const wsFirstTranscriptLatency = new Trend("ws_first_transcript_latency");
+
+export const options = {
+  vus: 10,
+  duration: "30s",
+  thresholds: {
+    ws_connections: ["count > 0"],
+    checks: ["rate > 0.9"],
+  },
+};
+
+const API_URL = __ENV.API_URL || "ws://localhost:4000";
+const AUTH_TOKEN = __ENV.AUTH_TOKEN || "";
+const AUDIO_URL = __ENV.AUDIO_URL || "https://dpgr.am/spacewalk.wav";
+const CHUNK_SIZE = 4096;
+const CHUNK_INTERVAL_MS = 100;
+
+export function setup() {
+  const res = http.get(AUDIO_URL, { responseType: "binary" });
+  check(res, { "audio fetch successful": (r) => r.status === 200 });
+  return { audioData: res.body };
+}
+
+export default function (data) {
+  const url = `${API_URL}/listen?provider=deepgram&language=en&encoding=linear16&sample_rate=16000`;
+  const params = {
+    headers: AUTH_TOKEN ? { Authorization: `Bearer ${AUTH_TOKEN}` } : {},
+  };
+
+  const startTime = Date.now();
+  let firstTranscriptTime = null;
+  let audioSendComplete = false;
+
+  const res = ws.connect(url, params, function (socket) {
+    socket.on("open", () => {
+      wsConnections.add(1);
+
+      const audioBytes = new Uint8Array(data.audioData);
+      let offset = 0;
+
+      socket.setInterval(() => {
+        if (offset < audioBytes.length) {
+          const end = Math.min(offset + CHUNK_SIZE, audioBytes.length);
+          const chunk = audioBytes.slice(offset, end);
+          socket.sendBinary(chunk.buffer);
+          offset = end;
+        } else if (!audioSendComplete) {
+          audioSendComplete = true;
+          socket.send(JSON.stringify({ type: "CloseStream" }));
+        }
+      }, CHUNK_INTERVAL_MS);
+
+      socket.setInterval(() => {
+        socket.send(JSON.stringify({ type: "KeepAlive" }));
+      }, 3000);
+    });
+
+    socket.on("message", (msg) => {
+      try {
+        const response = JSON.parse(msg);
+
+        if (response.type === "Results") {
+          wsTranscripts.add(1);
+
+          if (firstTranscriptTime === null) {
+            firstTranscriptTime = Date.now();
+            wsFirstTranscriptLatency.add(firstTranscriptTime - startTime);
+          }
+
+          const transcript =
+            response.channel?.alternatives?.[0]?.transcript || "";
+          if (transcript && response.is_final) {
+            console.log(`[transcript] ${transcript}`);
+          }
+        } else if (response.type === "Metadata") {
+          socket.close();
+        }
+      } catch (e) {
+        // ignore non-JSON messages
+      }
+    });
+
+    socket.on("error", (e) => {
+      if (e.error() !== "websocket: close sent") {
+        console.log("Error:", e.error());
+      }
+    });
+
+    socket.on("close", () => {
+      wsConnectionDuration.add(Date.now() - startTime);
+    });
+
+    socket.setTimeout(() => {
+      socket.send(JSON.stringify({ type: "CloseStream" }));
+      socket.close();
+    }, 30000);
+  });
+
+  check(res, {
+    "WebSocket upgrade successful": (r) => r && r.status === 101,
+  });
+
+  sleep(1);
+}