consistent-hashing.ts

import { createHash, randomBytes } from "crypto";

interface INode<T = any> {
  id: number;
  data: Map<string, T>;
  hash: string;
}

interface IHashRing {
  nodes: INode<any>[];
  addNode(node: INode): void;
  removeNode(node: INode): INode;
  putObject(doc: Document): void;
  getObject(collectionName: string, id: number): Document | undefined;
}

interface Document<T = any> {
  keyName: string;
  collectionName: string;
  data: T;
}

function createNodeHash(): string {
  return randomBytes(4).toString("hex");
}

function createObjectHash(key: string): string {
  return createHash("shake256", { outputLength: 4 }).update(key).digest("hex");
}

function isHashGreaterThan(hash1: string, hash2: string): boolean {
  return parseInt(hash1, 16) > parseInt(hash2, 16);
}

class ServerNode<T> implements INode<T> {
  public hash: string;
  constructor(public id: number, public data = new Map<string, T>()) {
    this.hash = createNodeHash();
  }
}

class HashRing implements IHashRing {
  constructor(public nodes: INode[] = []) {}

  public addNode(node: INode): void {
    let placedIndex = -1;
    let i = -1;

    while (++i < this.nodes.length) {
      const currentNode = this.nodes[i];
      const shouldPlaceNode = !isHashGreaterThan(node.hash, currentNode.hash);

      if (shouldPlaceNode) {
        this.nodes.splice(i, 0, node);
        placedIndex = i;
        break;
      }
    }

    if (placedIndex === -1) {
      this.nodes.push(node);
      const nodeIndexToBeRehashed = 0;
      this._rehash(this.nodes[nodeIndexToBeRehashed]);
    } else {
      this._rehash(this.nodes[placedIndex + 1]);
    }
  }
  public removeNode(node: INode): INode {
    this.nodes.splice(this.nodes.indexOf(node), 1);
    this._rehash(node);

    return node;
  }

  public putObject(doc: Document): void {
    const hash = createObjectHash(
      `${doc.collectionName}:${doc.data[doc.keyName]}`,
    );

    const nodeToPlaceDoc = this._findNode(hash);

    nodeToPlaceDoc.data.set(hash, doc);
  }

  public getObject(collectionName: string, id: number): Document | undefined {
    const hash = createObjectHash(`${collectionName}:${id}`);
    const node = this._findNode(hash);
    const object = node.data.get(hash);
    return object ? { ...object, nodeHash: node.hash } : undefined;
  }

  protected _rehash(nodeToRehash: INode) {
    for (const [hash, objectValue] of nodeToRehash.data.entries()) {
      const newNode = this._findNode(hash);
      const shouldPlaceToDifferentNode = newNode !== nodeToRehash;

      if (shouldPlaceToDifferentNode) {
        newNode.data.set(hash, objectValue);
        nodeToRehash.data.delete(hash);
      }
    }
  }

  protected _findNode(docHash: string): INode {
    let start = 0;
    let end = this.nodes.length;

    let cursor = start + Math.floor((end - start) / 2);

    while (start < end && this.nodes[cursor]) {
      if (isHashGreaterThan(docHash, this.nodes[cursor].hash)) {
        start = cursor + 1;
      } else {
        end = cursor;
      }
      cursor = start + Math.floor((end - start) / 2);
    }

    return this.nodes[cursor < this.nodes.length ? cursor : 0];
  }
}
const hashRing = new HashRing();
const node1 = new ServerNode<any>(1);
const node2 = new ServerNode<any>(2);

hashRing.addNode(node1);
hashRing.addNode(node2);

// insert data
for (let i = 0; i < 1000; i++) {
  hashRing.putObject({
    collectionName: "users",
    keyName: "user_id",
    data: {
      user_id: i + 1,
      name: String.fromCharCode(65 + Math.floor(Math.random() * 26)),
      age: Math.floor(Math.random() * 100 + 18),
      salary: "£100000",
    },
  });
}

console.log(hashRing.getObject("users", 300));
console.log(hashRing.getObject("users", 500));

console.log("\nAdding 2 new server nodes \n");
const node3 = new ServerNode<any>(3);
const node4 = new ServerNode<any>(4);
hashRing.addNode(node3);
hashRing.addNode(node4);

console.log(hashRing.getObject("users", 300));
console.log(hashRing.getObject("users", 500));

console.log("\nRemoving Added server nodes \n");

const removedNode3 = hashRing.removeNode(node3);
const removedNode4 = hashRing.removeNode(node4);

console.log(hashRing.getObject("users", 300));
console.log(hashRing.getObject("users", 500));

console.log("\nRe-adding the removed server nodes \n");
hashRing.addNode(removedNode3);
hashRing.addNode(removedNode4);

console.log(hashRing.getObject("users", 300));
console.log(hashRing.getObject("users", 500));

console.log("\nAdding 10 new server nodes \n");

hashRing.addNode(new ServerNode<any>(5));
hashRing.addNode(new ServerNode<any>(6));
hashRing.addNode(new ServerNode<any>(7));
hashRing.addNode(new ServerNode<any>(8));
hashRing.addNode(new ServerNode<any>(9));
hashRing.addNode(new ServerNode<any>(10));
hashRing.addNode(new ServerNode<any>(11));
hashRing.addNode(new ServerNode<any>(12));
hashRing.addNode(new ServerNode<any>(13));
hashRing.addNode(new ServerNode<any>(14));

console.log(hashRing.getObject("users", 300));
console.log(hashRing.getObject("users", 500));

console.log(
  "Server node hashring hash range",
  hashRing.nodes.map((item) => item.hash),
);