Database Sharding & Replication

import crypto from "node:crypto";
import pg from "pg";

// --- Consistent Hash Ring ---
class ConsistentHashRing {
  private ring: Map<number, string> = new Map();
  private sortedKeys: number[] = [];
  private virtualNodes: number;

  constructor(nodes: string[], virtualNodes = 150) {
    this.virtualNodes = virtualNodes;
    for (const node of nodes) {
      this.addNode(node);
    }
  }

  private hash(key: string): number {
    const h = crypto.createHash("md5").update(key).digest();
    return h.readUInt32BE(0);
  }

  addNode(node: string): void {
    for (let i = 0; i < this.virtualNodes; i++) {
      const virtualKey = `${node}:v${i}`;
      const hash = this.hash(virtualKey);
      this.ring.set(hash, node);
      this.sortedKeys.push(hash);
    }
    this.sortedKeys.sort((a, b) => a - b);
  }

  removeNode(node: string): void {
    for (let i = 0; i < this.virtualNodes; i++) {
      const virtualKey = `${node}:v${i}`;
      const hash = this.hash(virtualKey);
      this.ring.delete(hash);
      this.sortedKeys = this.sortedKeys.filter((k) => k !== hash);
    }
  }

  getNode(key: string): string {
    if (this.sortedKeys.length === 0) {
      throw new Error("No nodes in the ring");
    }

    const hash = this.hash(key);

    // Binary search for the first node clockwise from the hash
    let low = 0;
    let high = this.sortedKeys.length - 1;

    if (hash > this.sortedKeys[high]) {
      // Wrap around to first node
      return this.ring.get(this.sortedKeys[0])!;
    }

    while (low < high) {
      const mid = (low + high) >>> 1;
      if (this.sortedKeys[mid] < hash) {
        low = mid + 1;
      } else {
        high = mid;
      }
    }

    return this.ring.get(this.sortedKeys[low])!;
  }
}

// --- Shard Manager ---
interface ShardConfig {
  name: string;
  connectionString: string;
}

class ShardManager {
  private shards: Map<string, pg.Pool> = new Map();
  private ring: ConsistentHashRing;
  private readReplicas: Map<string, pg.Pool[]> = new Map();

  constructor(
    shardConfigs: ShardConfig[],
    replicaConfigs?: Record<string, string[]>
  ) {
    // Create connection pools for each shard
    const nodeNames: string[] = [];
    for (const config of shardConfigs) {
      const pool = new pg.Pool({
        connectionString: config.connectionString,
        max: 10,
        idleTimeoutMillis: 30000,
      });
      this.shards.set(config.name, pool);
      nodeNames.push(config.name);
    }

    // Create read replica pools
    if (replicaConfigs) {
      for (const [shard, replicas] of Object.entries(replicaConfigs)) {
        const pools = replicas.map(
          (connStr) => new pg.Pool({ connectionString: connStr, max: 10 })
        );
        this.readReplicas.set(shard, pools);
      }
    }

    this.ring = new ConsistentHashRing(nodeNames);
  }

  // Get the shard for a given key
  getShardName(shardKey: string): string {
    return this.ring.getNode(shardKey);
  }

  // Get write pool (primary)
  getWritePool(shardKey: string): pg.Pool {
    const shardName = this.getShardName(shardKey);
    const pool = this.shards.get(shardName);
    if (!pool) throw new Error(`Shard ${shardName} not found`);
    return pool;
  }

  // Get read pool (replica or primary fallback)
  getReadPool(shardKey: string): pg.Pool {
    const shardName = this.getShardName(shardKey);
    const replicas = this.readReplicas.get(shardName);

    if (replicas && replicas.length > 0) {
      // Round-robin across replicas
      const idx = Math.floor(Math.random() * replicas.length);
      return replicas[idx];
    }

    // Fallback to primary
    return this.getWritePool(shardKey);
  }

  // --- Query helpers ---
  async writeQuery<T>(
    shardKey: string,
    sql: string,
    params: unknown[]
  ): Promise<T[]> {
    const pool = this.getWritePool(shardKey);
    const result = await pool.query(sql, params);
    return result.rows;
  }

  async readQuery<T>(
    shardKey: string,
    sql: string,
    params: unknown[]
  ): Promise<T[]> {
    const pool = this.getReadPool(shardKey);
    const result = await pool.query(sql, params);
    return result.rows;
  }

  // Scatter-gather: query all shards and merge results
  async queryAllShards<T>(sql: string, params: unknown[]): Promise<T[]> {
    const promises = Array.from(this.shards.values()).map((pool) =>
      pool.query(sql, params).then((r) => r.rows as T[])
    );

    const results = await Promise.all(promises);
    return results.flat();
  }

  async close(): Promise<void> {
    for (const pool of this.shards.values()) {
      await pool.end();
    }
    for (const replicas of this.readReplicas.values()) {
      for (const pool of replicas) {
        await pool.end();
      }
    }
  }
}

// --- Usage Example ---
async function main() {
  const manager = new ShardManager(
    [
      { name: "shard-1", connectionString: "postgres://localhost:5432/blog_shard1" },
      { name: "shard-2", connectionString: "postgres://localhost:5433/blog_shard2" },
      { name: "shard-3", connectionString: "postgres://localhost:5434/blog_shard3" },
    ],
    {
      "shard-1": ["postgres://localhost:5435/blog_shard1_replica"],
      "shard-2": ["postgres://localhost:5436/blog_shard2_replica"],
    }
  );

  const userId = "user-12345";

  // Writes go to primary
  await manager.writeQuery(
    userId,
    `INSERT INTO users (id, username, email) VALUES ($1, $2, $3)
     ON CONFLICT (id) DO NOTHING`,
    [userId, "johndoe", "john@example.com"]
  );

  // Reads go to replica
  const user = await manager.readQuery(
    userId,
    "SELECT * FROM users WHERE id = $1",
    [userId]
  );

  console.log(`User on shard: ${manager.getShardName(userId)}`, user);

  // Cross-shard query (scatter-gather)
  const allUsers = await manager.queryAllShards(
    "SELECT id, username FROM users ORDER BY created_at DESC LIMIT 10",
    []
  );
  console.log("Users across all shards:", allUsers.length);

  await manager.close();
}

main().catch(console.error);