index.js 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173
  1. const crypto = require('crypto')
  2. const os = require('os')
  3. const { promisify } = require('util')
  4. const Redis = require('ioredis')
  5. const {
  6. RedisHealthCheckTimedOut,
  7. RedisHealthCheckWriteError,
  8. RedisHealthCheckVerifyError,
  9. } = require('./Errors')
  10. const HEARTBEAT_TIMEOUT = 2000
  11. // generate unique values for health check
  12. const HOST = os.hostname()
  13. const PID = process.pid
  14. const RND = crypto.randomBytes(4).toString('hex')
  15. let COUNT = 0
  16. function createClient(opts) {
  17. const standardOpts = Object.assign({}, opts)
  18. delete standardOpts.key_schema
  19. if (standardOpts.retry_max_delay == null) {
  20. standardOpts.retry_max_delay = 5000 // ms
  21. }
  22. if (standardOpts.endpoints) {
  23. throw new Error(
  24. '@overleaf/redis-wrapper: redis-sentinel is no longer supported'
  25. )
  26. }
  27. let client
  28. if (standardOpts.cluster) {
  29. delete standardOpts.cluster
  30. client = new Redis.Cluster(opts.cluster, standardOpts)
  31. } else {
  32. client = new Redis(standardOpts)
  33. }
  34. monkeyPatchIoRedisExec(client)
  35. client.healthCheck = (callback) => {
  36. if (callback) {
  37. // callback based invocation
  38. healthCheck(client).then(callback).catch(callback)
  39. } else {
  40. // Promise based invocation
  41. return healthCheck(client)
  42. }
  43. }
  44. return client
  45. }
  46. async function healthCheck(client) {
  47. // check the redis connection by storing and retrieving a unique key/value pair
  48. const uniqueToken = `host=${HOST}:pid=${PID}:random=${RND}:time=${Date.now()}:count=${COUNT++}`
  49. // o-error context
  50. const context = {
  51. uniqueToken,
  52. stage: 'add context for a timeout',
  53. }
  54. await runWithTimeout({
  55. runner: runCheck(client, uniqueToken, context),
  56. timeout: HEARTBEAT_TIMEOUT,
  57. context,
  58. })
  59. }
  60. async function runCheck(client, uniqueToken, context) {
  61. const healthCheckKey = `_redis-wrapper:healthCheckKey:{${uniqueToken}}`
  62. const healthCheckValue = `_redis-wrapper:healthCheckValue:{${uniqueToken}}`
  63. // set the unique key/value pair
  64. context.stage = 'write'
  65. const writeAck = await client
  66. .set(healthCheckKey, healthCheckValue, 'EX', 60)
  67. .catch((err) => {
  68. throw new RedisHealthCheckWriteError('write errored', context, err)
  69. })
  70. if (writeAck !== 'OK') {
  71. context.writeAck = writeAck
  72. throw new RedisHealthCheckWriteError('write failed', context)
  73. }
  74. // check that we can retrieve the unique key/value pair
  75. context.stage = 'verify'
  76. const [roundTrippedHealthCheckValue, deleteAck] = await client
  77. .multi()
  78. .get(healthCheckKey)
  79. .del(healthCheckKey)
  80. .exec()
  81. .catch((err) => {
  82. throw new RedisHealthCheckVerifyError(
  83. 'read/delete errored',
  84. context,
  85. err
  86. )
  87. })
  88. if (roundTrippedHealthCheckValue !== healthCheckValue) {
  89. context.roundTrippedHealthCheckValue = roundTrippedHealthCheckValue
  90. throw new RedisHealthCheckVerifyError('read failed', context)
  91. }
  92. if (deleteAck !== 1) {
  93. context.deleteAck = deleteAck
  94. throw new RedisHealthCheckVerifyError('delete failed', context)
  95. }
  96. }
  97. function unwrapMultiResult(result, callback) {
  98. // ioredis exec returns a results like:
  99. // [ [null, 42], [null, "foo"] ]
  100. // where the first entries in each 2-tuple are
  101. // presumably errors for each individual command,
  102. // and the second entry is the result. We need to transform
  103. // this into the same result as the old redis driver:
  104. // [ 42, "foo" ]
  105. //
  106. // Basically reverse:
  107. // https://github.com/luin/ioredis/blob/v4.17.3/lib/utils/index.ts#L75-L92
  108. const filteredResult = []
  109. for (const [err, value] of result || []) {
  110. if (err) {
  111. return callback(err)
  112. } else {
  113. filteredResult.push(value)
  114. }
  115. }
  116. callback(null, filteredResult)
  117. }
  118. const unwrapMultiResultPromisified = promisify(unwrapMultiResult)
  119. function monkeyPatchIoRedisExec(client) {
  120. const _multi = client.multi
  121. client.multi = function () {
  122. const multi = _multi.apply(client, arguments)
  123. const _exec = multi.exec
  124. multi.exec = (callback) => {
  125. if (callback) {
  126. // callback based invocation
  127. _exec.call(multi, (error, result) => {
  128. // The command can fail all-together due to syntax errors
  129. if (error) return callback(error)
  130. unwrapMultiResult(result, callback)
  131. })
  132. } else {
  133. // Promise based invocation
  134. return _exec.call(multi).then(unwrapMultiResultPromisified)
  135. }
  136. }
  137. return multi
  138. }
  139. }
  140. async function runWithTimeout({ runner, timeout, context }) {
  141. let healthCheckDeadline
  142. await Promise.race([
  143. new Promise((resolve, reject) => {
  144. healthCheckDeadline = setTimeout(() => {
  145. // attach the timeout when hitting the timeout only
  146. context.timeout = timeout
  147. reject(new RedisHealthCheckTimedOut('timeout', context))
  148. }, timeout)
  149. }),
  150. runner.finally(() => clearTimeout(healthCheckDeadline)),
  151. ])
  152. }
  153. module.exports = {
  154. createClient,
  155. }