[DE-583] Implement cluster rebalancing

Author: pluma

CHANGELOG.md

@@ -31,6 +31,8 @@ This driver uses semantic versioning:
 
 - Added missing attributes to `QueryInfo` and `MultiExplainResult.stats` types (DE-607)
 
+- Added cluster rebalancing methods to `Database` (DE-583)
+
 ## [8.3.1] - 2023-06-05
 
 ### Changed

src/database.ts

@@ -661,6 +661,176 @@ export type QueryInfo = {
   stream: boolean;
 };
 
+/**
+ * Information about a cluster imbalance.
+ */
+export type ClusterImbalanceInfo = {
+  /**
+   * Information about the leader imbalance.
+   */
+  leader: {
+    /**
+     * The weight of leader shards per DB-Server. A leader has a weight of 1 by default but it is higher if collections can only be moved together because of `distributeShardsLike`.
+     */
+    weightUsed: number[];
+    /**
+     * The ideal weight of leader shards per DB-Server.
+     */
+    targetWeight: number[];
+    /**
+     * The number of leader shards per DB-Server.
+     */
+    numberShards: number[];
+    /**
+     * The measure of the leader shard distribution. The higher the number, the worse the distribution.
+     */
+    leaderDupl: number[];
+    /**
+     * The sum of all weights.
+     */
+    totalWeight: number;
+    /**
+     * The measure of the total imbalance. A high value indicates a high imbalance.
+     */
+    imbalance: number;
+    /**
+     * The sum of shards, counting leader shards only.
+     */
+    totalShards: number;
+  };
+  /**
+   * Information about the shard imbalance.
+   */
+  shards: {
+    /**
+     * The size of shards per DB-Server.
+     */
+    sizeUsed: number[];
+    /**
+     * The ideal size of shards per DB-Server.
+     */
+    targetSize: number[];
+    /**
+     * The number of leader and follower shards per DB-Server.
+     */
+    numberShards: number[];
+    /**
+     * The sum of the sizes.
+     */
+    totalUsed: number;
+    /**
+     * The sum of shards, counting leader and follower shards.
+     */
+    totalShards: number;
+    /**
+     * The sum of system collection shards, counting leader shards only.
+     */
+    totalShardsFromSystemCollections: number;
+    /**
+     * The measure of the total imbalance. A high value indicates a high imbalance.
+     */
+    imbalance: number;
+  };
+};
+
+/**
+ * Information about the current state of the cluster imbalance.
+ */
+export type ClusterRebalanceState = ClusterImbalanceInfo & {
+  /**
+   * The number of pending move shard operations.
+   */
+  pendingMoveShards: number;
+  /**
+   * The number of planned move shard operations.
+   */
+  todoMoveShards: number;
+};
+
+/**
+ * Options for rebalancing the cluster.
+ */
+export type ClusterRebalanceOptions = {
+  /**
+   * Maximum number of moves to be computed.
+   *
+   * Default: `1000`
+   */
+  maximumNumberOfMoves?: number;
+  /**
+   * Allow leader changes without moving data.
+   *
+   * Default: `true`
+   */
+  leaderChanges?: boolean;
+  /**
+   * Allow moving leaders.
+   *
+   * Default: `false`
+   */
+  moveLeaders?: boolean;
+  /**
+   * Allow moving followers.
+   *
+   * Default: `false`
+   */
+  moveFollowers?: boolean;
+  /**
+   * Ignore system collections in the rebalance plan.
+   *
+   * Default: `false`
+   */
+  excludeSystemCollections?: boolean;
+  /**
+   * Default: `256**6`
+   */
+  piFactor?: number;
+  /**
+   * A list of database names to exclude from the analysis.
+   *
+   * Default: `[]`
+   */
+  databasesExcluded?: string[];
+};
+
+export type ClusterRebalanceMove = {
+  /**
+   * The server name from which to move.
+   */
+  from: string;
+  /**
+   * The ID of the destination server.
+   */
+  to: string;
+  /**
+   * Shard ID of the shard to be moved.
+   */
+  shard: string;
+  /**
+   * Collection ID of the collection the shard belongs to.
+   */
+  collection: number;
+  /**
+   * True if this is a leader move shard operation.
+   */
+  isLeader: boolean;
+};
+
+export type ClusterRebalanceResult = {
+  /**
+   * Imbalance before the suggested move shard operations are applied.
+   */
+  imbalanceBefore: ClusterImbalanceInfo;
+  /**
+   * Expected imbalance after the suggested move shard operations are applied.
+   */
+  imbalanceAfter: ClusterImbalanceInfo;
+  /**
+   * Suggested move shard operations.
+   */
+  moves: ClusterRebalanceMove[];
+};
+
 /**
  * Database user to create with a database.
  */
@@ -1843,6 +2013,108 @@ export class Database {
   }
   //#endregion
 
+  //#region rebalancing
+  /**
+   * Computes the current cluster imbalance.
+   *
+   * @example
+   * ```js
+   * const db = new Database();
+   * const imbalance = await db.getClusterImbalance();
+   * ```
+   */
+  getClusterImbalance(): Promise<ClusterRebalanceState> {
+    return this.request(
+      { path: "/_admin/cluster/rebalance" },
+      (res) => res.body.result
+    );
+  }
+
+  /**
+   * Computes a set of move shard operations to rebalance the cluster.
+   *
+   * @example
+   * ```js
+   * const db = new Database();
+   * const result = await db.computerClusterRebalance({
+   *   moveLeaders: true,
+   *   moveFollowers: true
+   * });
+   * if (result.moves.length) {
+   *   await db.executeClusterRebalance(result.moves);
+   * }
+   * ```
+   */
+  computeClusterRebalance(
+    opts: ClusterRebalanceOptions
+  ): Promise<ClusterRebalanceResult> {
+    return this.request(
+      {
+        method: "POST",
+        path: "/_admin/cluster/rebalance",
+        body: {
+          version: 1,
+          ...opts,
+        },
+      },
+      (res) => res.body.result
+    );
+  }
+
+  /**
+   * Executes the given cluster move shard operations.
+   *
+   * @example
+   * ```js
+   * const db = new Database();
+   * const result = await db.computerClusterRebalance({
+   *   moveLeaders: true,
+   *   moveFollowers: true
+   * });
+   * if (result.moves.length) {
+   *   await db.executeClusterRebalance(result.moves);
+   * }
+   * ```
+   */
+  executeClusterRebalance(moves: ClusterRebalanceMove[]): Promise<unknown> {
+    return this.request({
+      method: "POST",
+      path: "/_admin/cluster/rebalance/execute",
+      body: {
+        version: 1,
+        moves,
+      },
+    });
+  }
+
+  /**
+   * Computes a set of move shard operations to rebalance the cluster and
+   * executes them.
+   *
+   * @example
+   * ```js
+   * const db = new Database();
+   * const result = await db.rebalanceCluster({
+   *   moveLeaders: true,
+   *   moveFollowers: true
+   * });
+   * // The cluster is now rebalanced.
+   * ```
+   */
+  rebalanceCluster(
+    opts: ClusterRebalanceOptions
+  ): Promise<ClusterRebalanceResult> {
+    return this.request({
+      method: "PUT",
+      path: "/_admin/cluster/rebalance",
+      body: {
+        version: 1,
+        ...opts,
+      },
+    });
+  }
+  //#endregion
+
   //#region databases
   /**
    * Creates a new `Database` instance for the given `databaseName` that