BB-413: backport conductor checkpoint feature

extensions/lifecycle/conductor/LifecycleConductor.js

@@ -18,13 +18,13 @@ const safeJsonParse = require('../util/safeJsonParse');
 const { AccountIdCache } = require('../util/AccountIdCache');
 const { BreakerState, CircuitBreaker } = require('breakbeat').CircuitBreaker;
 const {
-    updateCircuitBreakerConfigForImplicitOutputQueue
+    updateCircuitBreakerConfigForImplicitOutputQueue,
 } = require('../../../lib/CircuitBreaker');
 
 const DEFAULT_CRON_RULE = '* * * * *';
 const DEFAULT_CONCURRENCY = 10;
 const ACCOUNT_SPLITTER = ':';
-const BUCKET_CHECKPOINT_PUSH_NUMBER = 50;
+const BUCKET_CHECKPOINT_PUSH_NUMBER_BUCKETD = 50;
 
 const LIFEYCLE_CONDUCTOR_CLIENT_ID = 'lifecycle:conductor';
 
@@ -315,7 +315,13 @@ class LifecycleConductor {
             return this.listZookeeperBuckets(queue, log, cb);
         }
 
-        return this.listBucketdBuckets(queue, log, cb);
+        return this.restoreBucketCheckpoint((err, marker) => {
+            if (err) {
+                return cb(err);
+            }
+
+            return this.listBucketdBuckets(queue, marker || null, log, cb);
+        });
     }
 
     listZookeeperBuckets(queue, log, cb) {
@@ -361,24 +367,73 @@ class LifecycleConductor {
             });
     }
 
-    listBucketdBuckets(queue, log, cb) {
-        let isTruncated = false;
-        let marker = null;
+    checkpointBucket(bucketEntry, cb) {
+        if (bucketEntry === null) {
+            return process.nextTick(cb);
+        }
+
+        return this._zkClient.setData(
+            this.getBucketProgressZkPath(),
+            Buffer.from(bucketEntry),
+            this.lastSentVersion,
+            (err, stat) => {
+                if (err) {
+                    return cb(err);
+                }
+
+                if (stat) {
+                    this.lastSentVersion = stat.version;
+                }
+
+                this.lastSentId = null;
+
+                return cb();
+            },
+        );
+    }
+
+    restoreBucketCheckpoint(cb) {
+        this._zkClient.getData(this.getBucketProgressZkPath(), (err, data, stat) => {
+            if (err) {
+                return cb(err);
+            }
+
+            const entry = data ? data.toString('ascii') : null;
+            if (stat) {
+                this.lastSentVersion = stat.version;
+            }
+
+            return cb(null, entry);
+        });
+    }
+
+    listBucketdBuckets(queue, initMarker, log, cb) {
+        let isTruncated = true;
+        let marker = initMarker;
         let nEnqueued = 0;
         const start = new Date();
         const retryWrapper = new BackbeatTask();
 
+        this.lastSentId = null;
+        this.lastSentVersion = -1;
+
         async.doWhilst(
             next => {
-                if (queue.length() > this._maxInFlightBatchSize) {
-                    log.info('delaying bucket pull', {
-                        nEnqueuedToDownstream: nEnqueued,
-                        inFlight: queue.length(),
-                        maxInFlight: this._maxInFlightBatchSize,
-                        bucketListingPushRateHz: Math.round(nEnqueued * 1000 / (new Date() - start)),
+                const breakerState = this._circuitBreaker.state;
+                const queueInfo = {
+                    nEnqueuedToDownstream: nEnqueued,
+                    inFlight: queue.length(),
+                    maxInFlight: this._maxInFlightBatchSize,
+                    bucketListingPushRateHz: Math.round(nEnqueued * 1000 / (new Date() - start)),
+                    breakerState,
+                };
+
+                if (queue.length() > this._maxInFlightBatchSize ||
+                    breakerState !== BreakerState.Nominal) {
+                    log.info('delaying bucket pull', queueInfo);
+                    return this.checkpointBucket(this.lastSentId, () => {
+                        setTimeout(next, 10000);
                     });
-
-                    return setTimeout(next, 10000);
                 }
 
                 return retryWrapper.retry({
@@ -400,6 +455,7 @@ class LifecycleConductor {
                                 }
 
                                 isTruncated = result.IsTruncated;
+                                let needCheckpoint = false;
 
                                 result.Contents.forEach(o => {
                                     marker = o.key;
@@ -408,6 +464,11 @@ class LifecycleConductor {
                                     if (!this._isBlacklisted(canonicalId, bucketName)) {
                                         nEnqueued += 1;
                                         queue.push({ canonicalId, bucketName });
+                                        this.lastSentId = o.key;
+                                        if (nEnqueued % BUCKET_CHECKPOINT_PUSH_NUMBER_BUCKETD === 0) {
+                                            needCheckpoint = true;
+                                        }
+
                                     // Optimization:
                                     // If we only blacklist by accounts, and the last bucket is blacklisted
                                     //  we can skip listing buckets until the next account.
@@ -419,6 +480,10 @@ class LifecycleConductor {
                                     }
                                 });
 
+                                if (needCheckpoint) {
+                                    return this.checkpointBucket(marker, done);
+                                }
+
                                 return done();
                             }
                         ),
@@ -427,7 +492,21 @@ class LifecycleConductor {
                 }, next);
             },
             () => isTruncated,
-            err => cb(err, nEnqueued));
+            err => {
+                if (err) {
+                    return cb(err, nEnqueued);
+                }
+
+                // clear last seen bucket from zk
+                return this.checkpointBucket('', err => {
+                    if (err) {
+                        return cb(err);
+                    }
+
+                    return cb(null, nEnqueued);
+                });
+            }
+        );
     }
 
     _controlBacklog(done) {
@@ -521,7 +600,7 @@ class LifecycleConductor {
     }
 
     _setupZookeeperClient(cb) {
-        if (this._bucketSource !== 'zookeeper') {
+        if (!this.needsZookeeper()) {
             process.nextTick(cb);
             return;
         }
@@ -542,6 +621,12 @@ class LifecycleConductor {
         });
     }
 
+    needsZookeeper() {
+        return this._bucketSource === 'zookeeper' || // bucket list stored in zk
+            this._bucketSource === 'mongodb' || // bucket stream checkpoints in zk
+            this._bucketSource === 'bucketd'; // bucket stream checkpoints in zk
+    }
+
     /**
      * Initialize kafka producer and clients
      *