From 5b1ed216348f6b9474bddf91208556bbcd34c8a8 Mon Sep 17 00:00:00 2001
From: Ben Lubar <ben.lubar+github@gmail.com>
Date: Fri, 12 Jan 2018 12:32:07 -0600
Subject: [PATCH] [database/mongo] Improve speed of sortedSetRank (#6229)

* [database/mongo] Improve speed of sortedSetRank

* [database/mongo] Fix sortedSetRank to filter by _key
---
 src/database/mongo/sorted.js | 29 ++++++++++++++++++++---------
 test/database/sorted.js      | 30 ++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 9 deletions(-)

diff --git a/src/database/mongo/sorted.js b/src/database/mongo/sorted.js
index 22c7e44196..ca1d53f63b 100644
--- a/src/database/mongo/sorted.js
+++ b/src/database/mongo/sorted.js
@@ -203,25 +203,36 @@ module.exports = function (db, module) {
 	};
 
 	module.sortedSetRank = function (key, value, callback) {
-		getSortedSetRank(module.getSortedSetRange, key, value, callback);
+		getSortedSetRank(false, key, value, callback);
 	};
 
 	module.sortedSetRevRank = function (key, value, callback) {
-		getSortedSetRank(module.getSortedSetRevRange, key, value, callback);
+		getSortedSetRank(true, key, value, callback);
 	};
 
-	function getSortedSetRank(method, key, value, callback) {
+	function getSortedSetRank(reverse, key, value, callback) {
 		if (!key) {
 			return callback();
 		}
 		value = helpers.valueToString(value);
-		method(key, 0, -1, function (err, result) {
-			if (err) {
-				return callback(err);
+		module.sortedSetScore(key, value, function (err, score) {
+			if (err || score === null) {
+				return callback(err, null);
 			}
 
-			var rank = result.indexOf(value);
-			callback(null, rank !== -1 ? rank : null);
+			db.collection('objects').count({
+				$or: [
+					{
+						_key: key,
+						score: reverse ? { $gt: score } : { $lt: score },
+					},
+					{
+						_key: key,
+						score: score,
+						value: reverse ? { $gt: value } : { $lt: value },
+					},
+				],
+			}, function (err, rank) { callback(err, rank); });
 		});
 	}
 
@@ -235,7 +246,7 @@ module.exports = function (db, module) {
 		}
 
 		async.map(data, function (item, next) {
-			getSortedSetRank(module.getSortedSetRange, item.key, item.value, next);
+			getSortedSetRank(false, item.key, item.value, next);
 		}, callback);
 	};
 
diff --git a/test/database/sorted.js b/test/database/sorted.js
index 703d3e93b4..08799513b2 100644
--- a/test/database/sorted.js
+++ b/test/database/sorted.js
@@ -17,6 +17,9 @@ describe('Sorted Set methods', function () {
 			function (next) {
 				db.sortedSetAdd('sortedSetTest3', [2, 4], ['value2', 'value4'], next);
 			},
+			function (next) {
+				db.sortedSetAdd('sortedSetTest4', [1, 1, 2, 3, 5], ['b', 'a', 'd', 'e', 'c'], next);
+			},
 			function (next) {
 				db.sortedSetAdd('sortedSetLex', [0, 0, 0, 0], ['a', 'b', 'c', 'd'], next);
 			},
@@ -305,6 +308,33 @@ describe('Sorted Set methods', function () {
 				done();
 			});
 		});
+
+		it('should return the rank sorted by the score and then the value (a)', function (done) {
+			db.sortedSetRank('sortedSetTest4', 'a', function (err, rank) {
+				assert.equal(err, null);
+				assert.equal(arguments.length, 2);
+				assert.equal(rank, 0);
+				done();
+			});
+		});
+
+		it('should return the rank sorted by the score and then the value (b)', function (done) {
+			db.sortedSetRank('sortedSetTest4', 'b', function (err, rank) {
+				assert.equal(err, null);
+				assert.equal(arguments.length, 2);
+				assert.equal(rank, 1);
+				done();
+			});
+		});
+
+		it('should return the rank sorted by the score and then the value (c)', function (done) {
+			db.sortedSetRank('sortedSetTest4', 'c', function (err, rank) {
+				assert.equal(err, null);
+				assert.equal(arguments.length, 2);
+				assert.equal(rank, 4);
+				done();
+			});
+		});
 	});
 
 	describe('sortedSetRevRank()', function () {