See if we can be faster with streams

gwicke · gwicke · commit dd86f237f25e · 2014-12-11T22:47:46.000-08:00
diff --git a/PromiseStream.js b/PromiseStream.js
@@ -0,0 +1,63 @@
+"use strict";
+
+function PromiseStream (fn, args, size, maxConcurrency) {
+    this._buf = [];
+    this._fn = fn;
+    this._args = args;
+    this._size = size;
+    this._concurrency = 0;
+    this._maxConcurrency = maxConcurrency || 1;
+    this._waiters = [];
+}
+
+PromiseStream.prototype.next = function () {
+    var self = this;
+    function startRequest () {
+        self._concurrency++;
+        console.log('start', self._concurrency);
+        var arg;
+        if (Array.isArray(self._args) && self._args.length) {
+            arg = self._args.shift();
+        } else {
+            arg = self._args;
+            self._args = undefined;
+        }
+        return self._fn(arg).then(handleResult);
+    }
+
+    function handleResult (res) {
+        console.log('end', self._concurrency);
+        self._concurrency--;
+        if (self._waiters.length) {
+            self._waiters.shift().resolve(res);
+        } else {
+            self._buf.push(res);
+        }
+        if (!self._args) {
+            self._args = res;
+        }
+        if (self._buf.length < self._size) {
+            while (self._concurrency < self._maxConcurrency) {
+                startRequest();
+            }
+        }
+
+    }
+
+    while (self._concurrency < self._maxConcurrency) {
+        startRequest();
+    }
+
+    if (self._buf.length) {
+        return Promise.resolve(self._buf.shift());
+    } else {
+        return new Promise(function(resolve, reject) {
+            self._waiters.push({
+                resolve: resolve,
+                reject: reject
+            });
+        });
+    }
+};
+
+module.exports = PromiseStream;
diff --git a/htmldumper.js b/htmldumper.js
@@ -1,50 +1,55 @@
 "use strict";
 
-// Prefer bluebird promise implementation over es6-shim pulled in by prfun
-if (!global.Promise) {
+if (!global.Promise || !global.promise.promisify) {
     global.Promise = require('bluebird');
 }
-if (!global.Promise.promisify) {
-    global.Promise.promisifyAll = require('bluebird').promisifyAll;
-}
 
 var preq = require('preq');
-var http = require('http');
 var fs = Promise.promisifyAll(require('fs'));
+var PromiseStream = require('./PromiseStream');
 
-// Higher per-host parallelism
-var maxConcurrency = 50;
-http.globalAgent.maxSockets = maxConcurrency;
+// Article dump parallelism
+var maxConcurrency = 200;
 
-function getArticles (apiURL, namespace, next) {
+function getArticles (apiURL, namespace, res) {
+    var next = res.next || '';
+    if (next === 'finished') {
+        // nothing more to do.
+        return Promise.reject('Articles done');
+    }
 
     var url = apiURL + '?action=query&generator=allpages&gapfilterredir=nonredirects'
         + '&gaplimit=500&prop=revisions&gapnamespace='
         + namespace + '&format=json&gapcontinue=' + encodeURIComponent( next );
-    console.log(url);
-        return preq.get(url, { retries: 10 })
-        .then(function(res) {
-            res = res.body;
-            var articles = [];
-            var articleChunk = res.query.pages;
-            Object.keys(articleChunk).forEach( function(key) {
-                var article = articleChunk[key];
-                if ( article.revisions !== undefined ) {
-                    var title = article.title.replace( / /g, '_' );
-                    articles.push([title, article.revisions[0].revid]);
-                }
-            });
-            next = res['query-continue'].allpages.gapcontinue;
-            // XXX
-            //next = 'finished';
-            return { articles: articles, next: next || '' };
-        })
-        .catch(function(e) {
-            console.error('Error in getArticles:', e);
-            throw e;
+    //console.log(url);
+
+    return preq.get(url, { retries: 10 })
+    .then(function(res) {
+        res = res.body;
+        var articles = [];
+        var articleChunk = res.query.pages;
+        Object.keys(articleChunk).forEach( function(key) {
+            var article = articleChunk[key];
+            if ( article.revisions !== undefined ) {
+                var title = article.title.replace( / /g, '_' );
+                articles.push([title, article.revisions[0].revid]);
+            }
         });
+        next = res['query-continue'].allpages.gapcontinue;
+        // XXX
+        //next = 'finished';
+        return {
+            articles: articles,
+            next: next
+        };
+    })
+    .catch(function(e) {
+        console.error('Error in getArticles:', e);
+        throw e;
+    });
 }
 
+
 function dumpArticle (prefix, title, oldid) {
     var dirName = prefix + '/' + encodeURIComponent(title),
         fileName = dirName + '/' + oldid;
@@ -92,33 +97,63 @@ function makeDump (apiURL, prefix, ns) {
         fs.mkdirSync(prefix);
     } catch (e) {}
 
-    function dumpBatch(articleResult) {
-        var articles = articleResult.articles;
-        var next = articleResult.next;
-        Promise.all([
-            // Fetch the next batch of oldids while processing the last one
-            getArticles(apiURL, ns, next),
-
-            Promise.filter(articles, function(article) {
-                var title = article[0];
-                var oldid = article[1];
-                return dumpArticle(prefix, title, oldid)
-                .catch(function(e) {
-                    console.error('Error in makeDump:', title, oldid, e.stack);
-                });
-            }, { concurrency: maxConcurrency })
-        ])
-        .then(function(results){
-            //console.log(results);
-            var articleResult = results[0];
-            if (articleResult.next !== 'finished') {
-                return dumpBatch(articleResult);
+    var articleArgs = {
+        apiURL: apiURL,
+        namespace: ns,
+        next: ''
+    };
+
+    var articleStream = new PromiseStream(getArticles.bind(null, apiURL, ns),
+            {next: ''}, 10);
+    var articles = [];
+    var waiters = [];
+
+    function processArticles (newArticles) {
+        articles = newArticles.articles;
+        while(waiters.length && articles.length) {
+            waiters.pop().resolve(articles.shift());
+        }
+        if (waiters.length) {
+            articleStream.next().then(processArticles);
+        }
+    }
+
+    function getArticle() {
+        if (articles.length) {
+            return Promise.resolve(articles.shift());
+        } else {
+            if (!waiters.length) {
+                articleStream.next().then(processArticles);
             }
+            return new Promise(function(resolve, reject) {
+                waiters.push({resolve: resolve, reject: reject});
+            });
+        }
+    }
+
+    function dumpOne () {
+        return getArticle()
+        .then(function(article) {
+            var title = article[0];
+            var oldid = article[1];
+            return dumpArticle(prefix, title, oldid)
+            .catch(function(e) {
+                console.error('Error in makeDump:', title, oldid, e.stack);
+            });
+        });
+    }
+
+    var dumpStream = new PromiseStream(dumpOne, undefined, 10, maxConcurrency);
+
+    function loop () {
+        return dumpStream.next()
+        .then(loop)
+        .catch(function(e) {
+            console.log(e);
         });
     }
 
-    return getArticles(apiURL, ns, '')
-    .then(dumpBatch);
+    return loop();
 }
 
 if (module.parent === null) {
