CosmosDB cassandra retry policy

cosmos-retry-policy.js

/**
 * Inspired by https://github.com/Azure-Samples/azure-cosmos-cassandra-java-retry-sample-v4/blob/master/java-examples/src/test/java/com/microsoft/azure/cosmos/cassandra/CosmosRetryPolicy.java
 *
 * The MIT License (MIT)
 *
 * Copyright (c) Jan Schmidle. All rights reserved.
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
 * documentation files (the "Software"), to deal in the Software without restriction, including without limitation the
 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to
 * permit persons to whom the Software is furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all copies or substantial portions of the
 * Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
 * WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
 * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
 * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */


const util = require('util');
const cassandra = require('cassandra-driver');

function getRetryAfterMs(errorMessage) {

  const retryAfterMs = parseInt(
    errorMessage.split(",")
                   .map( s => s.trim().split("=") )
                   .filter( r => r[0] === "RetryAfterMs")[0][1] || -1
    ,10);

  return retryAfterMs;
}

function CosmosRetryPolicy(
  maxRetryCount,
) {
  this.maxRetryCount = maxRetryCount || 20;
}

util.inherits(CosmosRetryPolicy, cassandra.policies.retry.RetryPolicy);

CosmosRetryPolicy.prototype.retryManyTimesOrThrow = function (retryNumber) {
  return (this.maxRetryCount == -1 || retryNumber < this.maxRetryCount) ?
    this.retryResult() :
    this.rethrowResult();
}

CosmosRetryPolicy.prototype.onReadTimeout = function (info, consistency, received, blockFor, isDataPresent) {
  return this.retryManyTimesOrThrow(info.nbRetry);
}

CosmosRetryPolicy.prototype.onUnavailable = function (info, consistency, required, alive) {
  return this.retryManyTimesOrThrow(info.nbRetry);
}

CosmosRetryPolicy.prototype.onWriteTimeout = function (info, consistency, received, blockFor, writeType) {
  return this.retryManyTimesOrThrow(info.nbRetry);
}

CosmosRetryPolicy.prototype.onRequestError = async function (info, consistency, err){

  try {

    if (this.maxRetryCount == -1 || info.nbRetry < this.maxRetryCount) {

      const retryAfterMs = getRetryAfterMs(err.message);

      await (new Promise( r => setTimeout( () => r(), retryAfterMs)));

      return this.retryResult(consistency, true)
    }

    return this.rethrowResult();
  } catch (e) {

    return this.rethrowResult();
  }
}

module.exports = CosmosRetryPolicy;