objectrocket · keshavmagge · May 23, 2016 · May 19, 2016 · May 19, 2016 · May 20, 2016
diff --git a/changelog.rst b/changelog.rst
@@ -12,6 +12,7 @@ changelog
 - Implement the instance bound ACLs interface.
 - Implement the auth._verify method for token verification.
 - Allowing instances to be extensible.
+- sharded mongodb instance stats are fetched using a pool of threads (one thread per shard)
 
 0.3.x
 -----

diff --git a/objectrocket/bases.py b/objectrocket/bases.py
@@ -284,6 +284,7 @@ def _service_url(self):
         """The service specific URL of this instance object."""
         return self._client._url + '{}/{}/'.format(self.service, self.name)
 
+
 ###########
 # Mixins. #
 ###########

diff --git a/objectrocket/instances/mongodb.py b/objectrocket/instances/mongodb.py
@@ -2,10 +2,13 @@
 import datetime
 import json
 import logging
+import time
 
 import pymongo
 import requests
 
+from concurrent import futures
+
 from objectrocket import bases
 from objectrocket import util
 
@@ -98,20 +101,116 @@ def shards(self, add_shard=False):
 
         return response.json()
 
+    def get_aggregate_database_stats(self):
+        return requests.get(self._service_url + 'aggregate_database_stats/',
+                            **self._instances._default_request_kwargs).json()['data']
+
     @property
     @util.token_auto_auth
     def new_relic_stats(self):
         """
         Get stats for this instance.
         """
         if self._new_relic_stats is None:
-            response = requests.get('{}{}'.format(self._url,
-                                    'new-relic-stats'),
-                                    **self._instances._default_request_kwargs)
-            self._new_relic_stats = json.loads(response.content).get(
-                'data') if response.status_code == 200 else {}
+            # if this is a sharded instance, fetch shard stats in parallel
+            if self.type == 'mongodb_sharded':
+                shards = [Shard(self.name, self._service_url + 'shards/',
+                                self._client, shard_doc)
+                          for shard_doc in self.shards().get('data')]
+                fs = []
+                with futures.ThreadPoolExecutor(len(shards)) as executor:
+                    for shard in shards:
+                        fs.append(executor.submit(shard.get_shard_stats))
+                    futures.wait(fs, timeout=None, return_when=futures.ALL_COMPLETED)
+                stats_this_second = self._rollup_shard_stats_to_instance_stats(
+                    {shard.name: future.result() for (shard, future) in zip(shards, fs)})
+                # power nap
+                time.sleep(1)
+                # fetch again
+                fs = []
+                with futures.ThreadPoolExecutor(len(shards)) as executor:
+                    for shard in shards:
+                        fs.append(executor.submit(shard.get_shard_stats))
+                    futures.wait(fs, timeout=None, return_when=futures.ALL_COMPLETED)
+                stats_next_second = self._rollup_shard_stats_to_instance_stats(
+                    {shard.name: future.result() for (shard, future) in zip(shards, fs)})
+                self._new_relic_stats = self._compile_new_relic_stats(stats_this_second, stats_next_second)
+            else:
+                # fetch stats like we did before (by hitting new_relic_stats API resource)
+                response = requests.get('{}{}'.format(self._url,
+                                        'new-relic-stats'),
+                                        **self._instances._default_request_kwargs)
+                self._new_relic_stats = json.loads(response.content).get(
+                    'data') if response.status_code == 200 else {}
         return self._new_relic_stats
 
+    def _rollup_shard_stats_to_instance_stats(self, shard_stats):
+        """
+        roll up all shard stats to instance level stats
+
+        :param shard_stats: dict of {shard_name: shard level stats}
+        """
+        instance_stats = {}
+        opcounters_per_node = []
+
+        # aggregate replication_lag
+        instance_stats['replication_lag'] = max(map(lambda s: s['replication_lag'], shard_stats.values()))
+
+        aggregate_server_statistics = {}
+        for shard_name, stats in shard_stats.items():
+            for statistic_key in stats.get('shard_stats'):
+                if statistic_key != 'connections' and statistic_key in aggregate_server_statistics:
+                    aggregate_server_statistics[statistic_key] = util.sum_values(aggregate_server_statistics[statistic_key],
+                                                                                 stats.get('shard_stats')[statistic_key])
+                else:
+                    aggregate_server_statistics[statistic_key] = stats.get('shard_stats')[statistic_key]
+
+            # aggregate per_node_stats into opcounters_per_node
+            opcounters_per_node.append({shard_name: {member: node_stats['opcounters']
+                                                     for member, node_stats in stats.get('per_node_stats').items()}})
+
+        instance_stats['opcounters_per_node'] = opcounters_per_node
+        instance_stats['aggregate_server_statistics'] = aggregate_server_statistics
+        return instance_stats
+
+    def _compile_new_relic_stats(self, stats_this_second, stats_next_second):
+        """
+        from instance 'stats_this_second' and instance 'stats_next_second', compute some per
+        second stats metrics and other aggregated metrics
+
+        :param dict stats_this_second:
+        :param dict stats_next_second:
+        :return: compiled instance stats that has metrics
+
+        {'opcounters_per_node_per_second': {...},
+         'server_statistics_per_second': {...},
+         'aggregate_server_statistics': {...},
+         'replication_lag': 0.0,
+         'aggregate_database_statistics': {}
+         }
+        """
+        server_statistics_per_second = {}
+        opcounters_per_node_per_second = []
+        for subdoc in ["opcounters", "network"]:
+            first_doc = stats_this_second['aggregate_server_statistics'][subdoc]
+            second_doc = stats_next_second['aggregate_server_statistics'][subdoc]
+            keys = set(first_doc.keys()) | set(second_doc.keys())
+            server_statistics_per_second[subdoc] = {key: int(second_doc[key]) - int(first_doc[key]) for key in keys}
+        for node1, node2 in zip(stats_this_second['opcounters_per_node'], stats_next_second['opcounters_per_node']):
+            node_opcounters_per_second = {}
+            for repl, members in node2.items():
+                node_opcounters_per_second[repl] = {}
+                for member, ops in members.items():
+                    node_opcounters_per_second[repl][member] = {}
+                    for op, count in ops.items():
+                        node_opcounters_per_second[repl][member][op] = count - node1[repl][member][op]
+            opcounters_per_node_per_second.append(node_opcounters_per_second)
+
+        return {'opcounters_per_node_per_second': opcounters_per_node_per_second,
+                'server_statistics_per_second': server_statistics_per_second,
+                'aggregate_server_statistics': stats_next_second.get('aggregate_server_statistics'),
+                'replication_lag': stats_next_second.get('replication_lag'),
+                'aggregate_database_statistics': self.get_aggregate_database_stats()}
 
     @property
     def ssl_connect_string(self):
@@ -174,3 +273,72 @@ def _get_connection(self, ssl):
             connect_string = self.ssl_connect_string
 
         return pymongo.MongoClient(connect_string)
+
+
+class Shard(bases.Extensible):
+    """An ObjectRocket MongoDB instance shard.
+
+    :param dict instance_name: Name of the instance the shard belongs to
+    :param string stats_base_url: Base url to fetch information and stats for this shard.
+    :param objectrocket.client.Client or_client: handle to talk to OR API
+    :param dict shard_document: a dictionary representing a mongodb shard
+    """
+
+    def __init__(self, instance_name, stats_base_url, or_client, shard_document):
+        self._instance_name = instance_name
+        self._shardstr = shard_document['shardstr']
+        self._plan = shard_document['plan']
+        self._id = shard_document['id']
+        self._name = shard_document['name']
+        self._stats_base_url = stats_base_url
+        self._client = or_client
+
+    @property
+    def instance_name(self):
+        """
+        :return: name of parent instance
+        """
+        return self._instance_name
+
+    @property
+    def shard_string(self):
+        """
+        :return: shard string
+        """
+        return self._shardstr
+
+    @property
+    def plan(self):
+        """
+        :return: Objectrocket plan that the parent instance is on
+        """
+        return self._plan
+
+    @property
+    def name(self):
+        """
+        :return: shard's name
+        """
+        return self._name
+
+    @property
+    def id(self):
+        """
+        :return: shard's unique ID
+        """
+        return self._id
+
+    def get_shard_stats(self):
+        """
+        :return: get stats for this mongodb shard
+        """
+        return requests.get(self._stats_url, params={'include_stats': True},
+                            headers={'X-Auth-Token': self._client.auth._token}
+                            ).json()['data']['stats']
+
+    @property
+    def _stats_url(self):
+        """
+        :return: Objectrocket API endpoint to send shard stats request to
+        """
+        return '%s%s/' % (self._stats_base_url, self.name)
diff --git a/objectrocket/util.py b/objectrocket/util.py
@@ -1,5 +1,6 @@
 """Utility code for the objectrocket package."""
 import functools
+import datetime
 
 from six import create_bound_method
 
@@ -57,3 +58,33 @@ def wrapper(self, *args, **kwargs):
 
     # TODO(TheDodd): match func call signature and docs.
     return wrapper
+
+
+def sum_values(value1, value2):
+    # function borrowed from core, keeping original comment intact
+    # TODO: Kill this method with fire, it's the only way to be sure
+    if value1 is None:
+        return value2
+    if value2 is None:
+        return value1
+
+    number_types = (int, float, long, bool)
+
+    if isinstance(value1, number_types) and isinstance(value2, number_types):
+        return int(value1 + value2)
+
+    # make sure the entries are of the same type
+    if not (isinstance(value1, value2.__class__) or isinstance(value2, value1.__class__)):
+        message = ("Entry %s type %s and entry %s type %s are not of the same type"
+                   % (str(value1), str(type(value1)), str(value2), str(type(value2))))
+        raise TypeError(message)
+
+    if isinstance(value1, list):
+        return list(set(value1 + value2))
+    elif isinstance(value1, str) or isinstance(value1, datetime.datetime) or isinstance(value1, unicode):
+        return value1
+    elif isinstance(value1, dict):
+        keys = set(value1.iterkeys()) | set(value2.iterkeys())
+        return dict((key, sum_values(value1.get(key), value2.get(key))) for key in keys)
+    else:
+        return value1 + value2
diff --git a/requirements/prod.txt b/requirements/prod.txt
@@ -4,3 +4,4 @@ redis>=2.0
 requests>=2.0.0
 six<2.0
 stevedore
+futures==3.0.5
diff --git a/setup.py b/setup.py
@@ -4,7 +4,7 @@
 from setuptools import find_packages
 from setuptools import setup
 
-VERSION = ('0', '4', '5')
+VERSION = ('0', '4', '6')
 __version__ = '.'.join(VERSION)
 
 with open('README.md') as f:

diff --git a/tests/conftest.py b/tests/conftest.py
@@ -11,6 +11,8 @@
 from objectrocket import instances
 from objectrocket import constants
 
+from .mock_stats import (SHARD_STATS_THIS_SECOND, INSTANCE_STATS_THIS_SECOND,
+                         INSTANCE_STATS_NEXT_SECOND, NEW_RELIC_STATS)
 
 #####################
 # Generic fixtures. #
@@ -152,3 +154,27 @@ def patched_requests_map(request):
     patches['instances.mongodb'] = mocked.start()
 
     return patches
+
+
+@pytest.fixture
+def mock_shard_stats():
+    """Mock shard stats."""
+    return SHARD_STATS_THIS_SECOND
+
+
+@pytest.fixture
+def mock_instance_stats_this_second():
+    """Mock shard stats."""
+    return INSTANCE_STATS_THIS_SECOND
+
+
+@pytest.fixture
+def mock_instance_stats_next_second():
+    """Mock shard stats."""
+    return INSTANCE_STATS_NEXT_SECOND
+
+
+@pytest.fixture
+def mock_new_relic_stats():
+    """Mock shard stats."""
+    return NEW_RELIC_STATS
diff --git a/tests/instances/test_instances.py b/tests/instances/test_instances.py
@@ -1,9 +1,13 @@
 """Tests for the objectrocket.instances module."""
+import sys
 import responses
 
+import pytest
+
 from objectrocket.instances.mongodb import MongodbInstance
 from objectrocket.acls import Acl
 
+from ..utils import comparable_dictionaries
 
 ##########################################
 # Tests for Instances private interface. #
@@ -104,16 +108,35 @@ def test_instance_acls_create_makes_expected_call(mongodb_sharded_instance, acl)
 
 
 @responses.activate
-def test_instance_new_relic_stats(mongodb_sharded_instance):
-    inst = mongodb_sharded_instance
+def test_instance_new_relic_stats(mongodb_replicaset_instance,):
+    inst = mongodb_replicaset_instance
     expected_url = 'http://localhost:5050/v2/instances/{}/new-relic-stats'.format(inst.name)
     responses.add(
         responses.GET,
         expected_url,
-        status=200,
-        body='{}',
+            status=200,
+        body={},
         content_type="application/json",
     )
 
-    assert hasattr(mongodb_sharded_instance, 'new_relic_stats')
-    assert mongodb_sharded_instance.new_relic_stats == {}
+    assert hasattr(mongodb_replicaset_instance, 'new_relic_stats')
+    assert mongodb_replicaset_instance.new_relic_stats == {}
+
+
+@pytest.mark.skipif(sys.version_info[0] >= 3, reason='long and unincode use in objectrocket.utils.sum_values should be cleaned up before running this test on python 3')
+def test_instance_rollup_shard_stats_to_instance_stats(mongodb_sharded_instance,
+                                                       mock_shard_stats,
+                                                       mock_instance_stats_this_second):
+
+    rolled_up_instance_stats = mongodb_sharded_instance._rollup_shard_stats_to_instance_stats(mock_shard_stats)
+    # we check just the keys because some rolled up values won't be the same every time even if the same set
+    # of shard stats were to be rolled up to instance level.
+    assert comparable_dictionaries(rolled_up_instance_stats, mock_instance_stats_this_second)
+
+
+def test_compile_new_relic_stats(mongodb_sharded_instance, mock_instance_stats_this_second,
+                                 mock_instance_stats_next_second, mock_new_relic_stats):
+
+    new_relic_stats = mongodb_sharded_instance._compile_new_relic_stats(
+        mock_instance_stats_this_second, mock_instance_stats_next_second)
+    assert comparable_dictionaries(new_relic_stats, mock_new_relic_stats)