[document-updater] wrap doc locks in project lock (#35416)

* [document-updater] wrap doc locks in project lock

* [document-updater] use BlockingProject as redis key prefix

* [redis-wrapper] update comment

GitOrigin-RevId: 0ce50e76386871e6c29481494febc8d6b37e7df5

libraries/redis-wrapper/RedisLocker.js

@@ -13,6 +13,10 @@ const MAX_REDIS_REQUEST_LENGTH = 5000 // 5 seconds
 
 const UNLOCK_SCRIPT =
   'if redis.call("get", KEYS[1]) == ARGV[1] then return redis.call("del", KEYS[1]) else return 0 end'
+// Extend (bump) the TTL of a lock we still own, so a long-running critical
+// section does not let the lock expire from under us.
+const EXTEND_SCRIPT =
+  'if redis.call("get", KEYS[1]) == ARGV[1] then return redis.call("expire", KEYS[1], ARGV[2]) else return 0 end'
 
 module.exports = class RedisLocker {
   /**
@@ -69,6 +73,7 @@ module.exports = class RedisLocker {
     this.promises = {
       checkLock: promisify(this.checkLock.bind(this)),
       getLock: promisify(this.getLock.bind(this)),
+      extendLock: promisify(this.extendLock.bind(this)),
       releaseLock: promisify(this.releaseLock.bind(this)),
 
       // tryLock returns two values: gotLock and lockValue. We need to merge
@@ -198,6 +203,46 @@ module.exports = class RedisLocker {
     })
   }
 
+  /**
+   * Extend (bump) the TTL of a lock we still hold back to the full lock TTL.
+   * Useful for long critical sections (e.g. draining a whole project) that
+   * would otherwise outlive the lock and let a second worker in.
+   *
+   * @param {string} id
+   * @param {string} lockValue
+   * @param {function(Error): void} callback
+   */
+  extendLock(id, lockValue, callback) {
+    if (callback == null) {
+      callback = function () {}
+    }
+    const key = this.getKey(id)
+    this.rclient.eval(
+      EXTEND_SCRIPT,
+      1,
+      key,
+      lockValue,
+      this.LOCK_TTL,
+      (err, result) => {
+        if (err != null) {
+          return callback(err)
+        } else if (result != null && result !== 1) {
+          // A successful extension updates the expiration time on exactly one key.
+          // Anything else means we no longer hold the lock (it expired or was taken).
+          logger.error(
+            { id, key, lockValue, redis_err: err, redis_result: result },
+            'extending lock error'
+          )
+          metrics.inc(this.metricsPrefix + '-extend-error')
+          return callback(new Error('tried to extend a lock we no longer hold'))
+        } else {
+          metrics.inc(this.metricsPrefix + '-extend-success')
+          return callback(null)
+        }
+      }
+    )
+  }
+
   /**
    * @param {string} id
    * @param {string} lockValue

libraries/redis-wrapper/test/unit/src/test.js

@@ -116,6 +116,38 @@ describe('index', function () {
     })
   })
 
+  describe('extendLock', function () {
+    beforeEach(function () {
+      this.rclient = { eval: sinon.stub() }
+      this.locker = new this.RedisLocker({
+        rclient: this.rclient,
+        getKey: id => `lock:${id}`,
+        wrapTimeoutError: (err, id) => err,
+        metricsPrefix: 'test',
+        lockTTLSeconds: 30,
+      })
+    })
+
+    it('should extend the lock TTL when we still own it', function (done) {
+      this.rclient.eval.yields(null, 1)
+      this.locker.extendLock('id-1', 'lock-value', err => {
+        expect(err).to.not.exist
+        this.rclient.eval
+          .calledWith(sinon.match.string, 1, 'lock:id-1', 'lock-value', 30)
+          .should.equal(true)
+        done()
+      })
+    })
+
+    it('should error when we no longer hold the lock', function (done) {
+      this.rclient.eval.yields(null, 0)
+      this.locker.extendLock('id-1', 'lock-value', err => {
+        expect(err).to.exist
+        done()
+      })
+    })
+  })
+
   describe('redis-sentinel', function () {
     it('should throw an error when creating a client', function () {
       const redisSentinelOptions = {