Firestore 500/50/5 rule illustration with `k6`

generate-firebase-token.py

from google.oauth2 import service_account
from google.auth.transport.requests import Request

# Path to your service account key file
key_path = "[INSERT_PATH_TO_KEY_FILE]"

# Create credentials using the service account key file
credentials = service_account.Credentials.from_service_account_file(
    key_path,
    scopes=["https://www.googleapis.com/auth/cloud-platform"]
)

# Refresh the credentials to get a valid access token
credentials.refresh(Request())

# Get the access token
access_token = credentials.token

print(access_token)

k6_firestore_scaling.js

import http from "k6/http";
import { check, sleep } from "k6";
import { Rate, Counter } from "k6/metrics";
import { SharedArray } from "k6/data";

// Custom metrics
const readSuccess = new Rate("successful_reads");
const errorCounter = new Counter("errors");

// Firestore project details
const projectId = "INSERT_PROJECT_ID";
const collection = "INSERT_COLLECTION";

// Warmup parameters
const initialRPS = 500;
const targetRPS = 1500;
const stablePeriodSeconds = 300; // 5 minutes
const rampPeriodSeconds = 60; // 1 minute
const stageCount = Math.ceil(Math.log(targetRPS / initialRPS) / Math.log(1.5));

// Generate stages
const warmupStages = [];
let currentRPS = initialRPS;

for (let i = 0; i < stageCount; i++) {
  // Stay at current RPS for 5 minutes
  warmupStages.push({
    target: currentRPS,
    duration: `${stablePeriodSeconds}s`,
  });

  // Calculate next RPS
  const nextRPS = Math.min(Math.floor(currentRPS * 1.5), targetRPS);

  // Ramp up to next RPS over 1 minute
  warmupStages.push({
    target: nextRPS,
    duration: `${rampPeriodSeconds}s`,
  });

  currentRPS = nextRPS;

  // If we've reached or exceeded the target RPS, stop adding stages
  if (currentRPS >= targetRPS) break;
}

// Log the stages
console.log("Warmup Stages:");
warmupStages.forEach((stage, index) => {
  console.log(
    `Stage ${index + 1}: Target RPS: ${stage.target}, Duration: ${
      stage.duration
    }`
  );
});

/**
 * Configuration options for the k6 load testing script.
 * 
 * @constant {Object} options - The main configuration object.
 * @property {Object} scenarios - The scenarios configuration object.
 * @property {Object} scenarios.firestore_warmup - The specific scenario configuration.
 * @property {string} scenarios.firestore_warmup.executor - The type of executor to use.
 * @property {number} scenarios.firestore_warmup.startRate - The initial requests per second (RPS).
 * @property {string} scenarios.firestore_warmup.timeUnit - The time unit for the rate.
 * @property {number} scenarios.firestore_warmup.preAllocatedVUs - The number of pre-allocated virtual users.
 * @property {number} scenarios.firestore_warmup.maxVUs - The maximum number of virtual users.
 * @property {Array<Object>} scenarios.firestore_warmup.stages - The stages of the test.
 */
export const options = {
  scenarios: {
    firestore_warmup: {
      executor: "ramping-arrival-rate",
      startRate: initialRPS,
      timeUnit: "1s",
      preAllocatedVUs: 1000,
      maxVUs: 10000,
      stages: warmupStages,
    },
  },
};

// Load document IDs
const documentIds = new SharedArray("document_ids", function () {
  return open("./orders.txt").split("\n").filter(Boolean);
});

/**
 * Sets up the environment by retrieving the authentication token.
 *
 * @returns {Object} An object containing the authentication token.
 * @throws {Error} If the authentication token is not set.
 */
export function setup() {
  const authToken =
    "INSERT_TOKEN";
  if (!authToken) {
    throw new Error("AUTH_TOKEN environment variable is not set");
  }
  return { authToken };
}

export default function (data) {
  const authToken = data.authToken;
  const documentId =
    documentIds[Math.floor(Math.random() * documentIds.length)];

  const url = `https://firestore.googleapis.com/v1/projects/${projectId}/databases/(default)/documents/${collection}/${documentId}`;

  const params = {
    headers: {
      Authorization: `Bearer ${authToken}`,
      "Content-Type": "application/json",
    },
  };

  const response = http.get(url, params);

  const success = check(response, {
    "status is 200": (r) => r.status === 200,
  });

  if (!success) {
    errorCounter.add(1);
    console.error(
      `Error for document ${documentId}: ${response.status} ${response.body}`
    );
  }

  readSuccess.add(success);
}