@lpgroup/feathers-mongodb-hooks/src/sessions.js

Hooks for feathers-mongo.

/* eslint-disable no-use-before-define */
/* eslint-disable no-underscore-dangle */
// Read more:
// https://docs.mongodb.com/manual/core/transactions-in-applications/#transactions-retry
// https://www.mongodb.com/blog/post/how-to-select--for-update-inside-mongodb-transactions
// http://mongodb.github.io/node-mongodb-native/3.6/api/ClientSession.html

import { log, executeSequentially } from "@lpgroup/utils";
import { ObjectId } from "mongodb";
import { nanoid } from "nanoid";
import { sleep } from "@lpgroup/feathers-utils";
import { TransactionAborted } from "./exceptions.js";

const { debug, error } = log("database");

// Number of ms to wait before retrying a writeConflict
const writeConflictBase = 50;

const writeConflictJitterMax = 100;

// Number of times to retry to resume a writeConflict.
const writeConflictRetries = 20;

const sessionOptions = {
  // readPreference: { mode: 'primary' },
};

const transactionOptions = {
  readPreference: "primary",
  readConcern: { level: "snapshot" },
  writeConcern: { w: "majority", wtimeout: 1 },
};

let mongoClient = null;
export function setClient(client) {
  mongoClient = client;
}

export function getClient() {
  return mongoClient;
}

let mongoDatabase = null;
export function setDatabase(database) {
  mongoDatabase = database;
}

export function getDatabase() {
  return mongoDatabase;
}

const sessions = {};
export async function startSession() {
  const sessionId = nanoid();
  const session = getClient().startSession(sessionOptions);
  sessions[sessionId] = { id: sessionId, counter: 1, session };

  return { sessionId, session };
}

export function reuseSession(sessionId) {
  if (!hasSessionObject(sessionId)) {
    error(`reuseSession: Mongo session doesn't exist ${sessionId}`);
    throw new TransactionAborted("reuseSession: Mongo session doesn't exist", sessionId);
  }
  getSessionObject(sessionId).counter += 1;
}

function hasSessionObject(sessionId) {
  return !!sessions[sessionId];
}

function getSessionObject(sessionId) {
  if (!sessions[sessionId]) {
    // Maybe a service that reused the sessionId failed and deleted the id.
    throw new TransactionAborted("getSessionObject: Session doesn't exist", sessionId);
  }
  return sessions[sessionId];
}

export function getSession(sessionId) {
  return getSessionObject(sessionId).session;
}

export function getSessionCounter(sessionId) {
  return getSessionObject(sessionId).counter;
}

export async function endSession(params) {
  if (hasSessionObject(params.sessionId)) {
    const session = getSession(params.sessionId);
    // Might be in transaction if endSession is called by errorSession hook
    if (session.inTransaction()) {
      await session.abortTransaction();
    }
    await session.endSession();
    delete sessions[params.sessionId];

    // eslint-disable-next-line no-param-reassign
    delete params.sessionId;
    // eslint-disable-next-line no-param-reassign
    delete params.mongodb.session;
  }
}

async function _commitWithRetry(session, retry = 1) {
  try {
    await session.commitTransaction();
  } catch (err) {
    if (err.hasErrorLabel("UnknownTransactionCommitResult") || err.hasErrorLabel("WriteConflict")) {
      if (retry >= writeConflictRetries) throw err;
      const timeout = _getExponentialTimeoutWithJitter(retry);
      const counter = getCounter("commitRetry");
      debug(
        `UnknownTransactionCommitResult, retrying commit operation. retry: ${retry}:${counter}, timeout: ${timeout}`,
      );
      await sleep(timeout);
      await _commitWithRetry(session, retry + 1);
    } else {
      error("Error during commit", err);
      throw err;
    }
  }
}

async function ensureTransactionCompletion(session, maxRetryCount = 50) {
  // When we are trying to split our operations into multiple transactions
  // Sometimes we are getting an error that the earlier transaction is still in progress
  // To avoid that, we ensure the earlier transaction has finished
  let count = 0;
  while (session.inTransaction()) {
    if (count >= maxRetryCount) {
      break;
    }
    // Adding a delay so that the transaction get be committed
    // eslint-disable-next-line no-await-in-loop, no-promise-executor-return
    await new Promise((r) => setTimeout(r, 100));
    count += 1;
  }
}

export async function endSessionAndCommitTransaction(context) {
  const { params } = context;
  getSessionObject(params.sessionId).counter -= 1;
  if (getSessionObject(params.sessionId).counter === 0) {
    try {
      if (getSession(params.sessionId).inTransaction()) {
        await _commitWithRetry(getSession(params.sessionId));
        debugMsg("commitWithRetry", context);
      }
      ensureTransactionCompletion(getSession(params.sessionId));
    } catch (err) {
      error("abortTransaction", err);
      await endSession(params);
      throw err;
    }
    return endSession(params);
  }
  return undefined;
}

export async function startGetAndLockTransaction(context, collections, retry = 0) {
  const { params } = context;
  try {
    debugMsg("startGetAndLockTransaction 1 ", context);
    await _startTransaction(params.sessionId);
    const fromQuery = _getFilterQueriesFromQuery(context, collections);
    const fromModel = _getFilterQueriesFromModel(context);
    await _lockDocument(context, "updateOne", params.sessionId, fromQuery);
    const result = await _lockDocument(context, "findOneAndUpdate", params.sessionId, fromModel);

    debugMsg("startGetAndLockTransaction 2 ", context);

    if (result.length > 0) return result[0];
  } catch (err) {
    if (err.codeName === "WriteConflict") {
      if (retry >= writeConflictRetries) throw err;
      // eslint-disable-next-line no-param-reassign
      retry += 1;
      const timeout = _getExponentialTimeoutWithJitter(retry);
      const counter = getCounter("writeConflict");
      debugMsg(`Write conflict (lock) retry: ${retry}:${counter}, timeout: ${timeout}`, context);
      await sleep(timeout);
      await getSession(params.sessionId).abortTransaction();
      return startGetAndLockTransaction(context, collections, retry);
    }
    throw err;
  }
  return {};
}

async function _startTransaction(sessionId) {
  if (getSessionCounter(sessionId) === 1) {
    if (getSession(sessionId).inTransaction()) {
      error("Aborting transaction, should it really be started here?", sessionId);
      await getSession(sessionId).abortTransaction();
    }
  }

  if (!getSession(sessionId).inTransaction()) {
    return getSession(sessionId).startTransaction(transactionOptions);
  }
  return undefined;
}

async function _lockDocument(context, operation, sessionId, filterQueries) {
  const lockQuery = { $set: { myLock: new ObjectId() } };
  return executeSequentially(filterQueries, async (v) => {
    debug(
      `${sessionId} _lockDocument filterQueries ${JSON.stringify(
        filterQueries.map((f) => ({
          collection: f?.collection?.collectionName,
          filterQuery: f.filterQuery,
        })),
      )}`,
    );
    const session = getSession(sessionId);

    const result = await v.collection[operation](v.filterQuery, lockQuery, {
      session,
    }).catch((err) => {
      debugMsg(
        `_lockDocument error ${JSON.stringify({
          state: session?.transaction?.state,
          active: session?.transaction?.isActive,
        })}:`,
        context,
      );
      throw err;
    });

    debugMsg(`_lockDocument transactionState: ${session?.transaction?.state}:`, context);

    return result;
  });
}

/**
 * Build array with filterQueries from options.Model set in
 * the feathersjs Service. ie. users.class.js
 * Using context.id that is parsed by feathersjs from the uri.
 */
function _getFilterQueriesFromModel(context) {
  if (context?.service?.options?.Model) {
    const { options } = context.service;
    return [
      {
        collection: options.Model,
        filterQuery: { ...context.params.query, [options.id]: context.id },
      },
    ];
  }
  return [];
}

/**
 * Get filterQuery based on collection parameter on lock-data hook.
 *
 * @param {*} collections [
 *  {
 *    collection: "users",  // Name of mongodb collection
 *    field: "_id",         // Name of mongodb document field
 *    query: "userId"       // Name of query field in Feathersjs.
 *  }]
 */
function _getFilterQueriesFromQuery(context, collections) {
  if (collections && collections.length > 0) {
    const { params } = context;
    const options = context?.service?.options;
    return collections.map((v) => {
      const field = v.field || options.id || "_id";
      const query = params.query[v.query] || context.id;
      return {
        collection: getDatabase().collection(v.collection),
        filterQuery: { [field]: query },
      };
    });
  }
  return [];
}

export function debugMsg(prefix, context, suffix = "") {
  const { sessionId } = context.params;
  let counter = "NA";
  if (hasSessionObject(sessionId)) counter = getSessionCounter(sessionId);
  debug(`${sessionId} ${prefix}(${counter}) ${context.method} ${getUrl(context)} ${suffix} `);
}

function getUrl(context) {
  let url = context.path || "";
  if (context.params.query)
    Object.entries(context.params.query).forEach((key, value) => {
      url = url.replace(`:${key}`, value);
    });
  if (context.id) url += `/${context.id}`;
  return url;
}

function _getExponentialTimeoutWithJitter(retry) {
  const exp = getRandomInt(retry * 0.5) + retry * 0.5;
  const exponent = 1.6 ** exp;
  const rand = getRandomInt(writeConflictJitterMax * retry);
  return Math.floor(writeConflictBase * exponent + rand);
}

function getRandomInt(max) {
  return Math.random() * Math.floor(max);
}

const counters = { writeConflict: 0, commitRetry: 0 };
function getCounter(name) {
  counters[name] += 1;
  return counters[name];
}