Optimize storage (serialization and de-serilization) of very large dictionaries inside MongoDB

conf/st2.conf.sample

@@ -115,6 +115,8 @@ username = None
 connection_retry_max_delay_m = 3
 # ca_certs file contains a set of concatenated CA certificates, which are used to validate certificates passed from MongoDB.
 ssl_ca_certs = None
+# True to use a special implementation of escaped dict field which saves value as JSON to avoid expensive escaping. NOTE: Experimental.
+use_json_dict_field = False
 # Certificate file used to identify the localconnection
 ssl_certfile = None
 # Connection retry backoff max (seconds).
@@ -141,6 +143,8 @@ connection_timeout = 3000
 password = None
 # port of db server
 port = 27017
+# The backend to use for marshalling JSON to the JSONDictField.
+json_dict_field_backend = ujson
 
 [exporter]
 # location of the logging.exporter.conf file

conf/st2.dev.conf

@@ -1,6 +1,8 @@
 # Config used by local development environment (tools/launch.dev.sh)
 [database]
 host = 127.0.0.1
+use_json_dict_field = True
+json_dict_field_backend = ujson
 
 [api]
 # Host and port to bind the API server.

st2common/st2common/config.py

@@ -212,7 +212,14 @@ def register_opts(ignore_errors=False):
             'authentication_mechanism', default=None,
             help='Specifies database authentication mechanisms. '
                  'By default, it use SCRAM-SHA-1 with MongoDB 3.0 and later, '
-                 'MONGODB-CR (MongoDB Challenge Response protocol) for older servers.')
+                 'MONGODB-CR (MongoDB Challenge Response protocol) for older servers.'),
+        cfg.BoolOpt(
+            'use_json_dict_field', default=False,
+            help='True to use a special implementation of escaped dict field which saves '
+            'value as JSON to avoid expensive escaping. NOTE: Experimental.'),
+        cfg.StrOpt(
+            'json_dict_field_backend', default='ujson', choices='cjson, ujson',
+            help='The backend to use for marshalling JSON to the JSONDictField.'),
     ]
 
     do_register_opts(db_opts, 'database', ignore_errors)

st2common/st2common/fields.py

@@ -13,19 +13,28 @@
 # limitations under the License.
 
 from __future__ import absolute_import
+
 import datetime
 import calendar
 
+import six
+
 from mongoengine import LongField
+from mongoengine import BinaryField
+from oslo_config import cfg
 
+from st2common.models.db import stormbase
 from st2common.util import date as date_utils
+from st2common.util import mongoescape
 
 __all__ = [
     'ComplexDateTimeField'
 ]
 
 SECOND_TO_MICROSECONDS = 1000000
 
+from st2common import log as logging
+LOG = logging.getLogger(__name__)
 
 class ComplexDateTimeField(LongField):
     """
@@ -114,3 +123,95 @@ def to_mongo(self, value):
 
     def prepare_query_value(self, op, value):
         return self._convert_from_datetime(value)
+
+
+class JSONDictField(BinaryField):
+    """
+    Custom field types which stores dictionary as JSON serialized strings.
+
+    This is done because storing large objects as JSON serialized strings is much more efficent
+    on the serialize and unserialize paths compared to used EscapedDictField which needs to escape
+    all the special values ($, .).
+
+    Only downside is that to MongoDB those values are plain raw strings which means you can't query
+    on actual dictionary field values. That's not an issue for us, because in places where we use
+    it, we already treat those values more or less as opaque strings.
+
+    # NOTE(Tomaz): I've done bencharmking of ujson and cjson and cjson is more performant on large
+    objects and ujson on smaller ones.
+    """
+    def __init__(self, *args, **kwargs):
+        json_backend = kwargs.pop('json_backend', None)
+
+        if not json_backend:
+            json_backend = cfg.CONF.database.json_dict_field_backend
+
+        if json_backend not in ['ujson', 'cjson']:
+            raise ValueError('Unsupported backend "%s" specified for JSONDictField.' % json_backend)
+
+        super(JSONDictField, self).__init__(*args, **kwargs)
+
+        if json_backend == 'ujson':
+            import ujson
+            self.json_loads = ujson.loads
+            self.json_dumps = ujson.dumps
+        elif json_backend == 'cjson':
+            import cjson
+            self.json_loads = cjson.decode
+            self.json_dumps = cjson.encode
+
+    def to_mongo(self, value):
+        if not isinstance(value, dict):
+            message = 'The value argument must be a dictionary. Type: %s Content: %s'
+            raise ValueError(message % (type(value), str(value)))
+
+        return self.json_dumps(value)
+
+    def to_python(self, value):
+        if isinstance(value, six.text_type) or isinstance(value, six.binary_type):
+            return self.json_loads(value)
+
+        return value
+
+    def validate(self, value):
+        if isinstance(value, dict):
+            value = self.to_mongo(value)
+
+        return super(JSONDictField, self).validate(value)
+
+
+class JSONDictEscapedFieldCompatibilityField(JSONDictField):
+    """
+    Special version of JSONDictField which takes care of compatibility between old EscapedDictField
+    and EscapedDynamicField format and the new one.
+
+    On retrieval, if an old format is detected it's correctly un-serialized and on insertion, we
+    always insert data in a new format.
+    """
+
+    def to_mongo(self, value, use_db_field=True, fields=None):
+        if not cfg.CONF.database.use_json_dict_field:
+            value = mongoescape.escape_chars(value)
+            return super(stormbase.EscapedDynamicField, self).to_mongo(
+                value=value, use_db_field=use_db_field, fields=fields)
+
+        if not isinstance(value, dict):
+            message = 'The value argument must be a dictionary. Type: %s Content: %s'
+            raise ValueError(message % (type(value), str(value)))
+
+        return self.json_dumps(value)
+
+    def to_python(self, value):
+        if not cfg.CONF.database.use_json_dict_field:
+            value = super(stormbase.EscapedDynamicField, self).to_python(value)
+            return mongoescape.unescape_chars(value)
+
+        if isinstance(value, dict):
+            # Old format which used a native dict with escaped special characters
+            value = mongoescape.unescape_chars(value)
+            return value
+
+        if isinstance(value, six.text_type) or isinstance(value, six.binary_type):
+            return self.json_loads(value)
+
+        return value

st2common/st2common/models/api/execution.py

@@ -152,6 +152,10 @@ class ActionExecutionAPI(BaseAPI):
     @classmethod
     def from_model(cls, model, mask_secrets=False):
         doc = cls._from_model(model, mask_secrets=mask_secrets)
+
+        import json
+        doc['result'] = json.loads(doc['result'])
+
         start_timestamp = model.start_timestamp
         start_timestamp_iso = isotime.format(start_timestamp, offset=False)
         doc['start_timestamp'] = start_timestamp_iso

st2common/st2common/models/db/execution.py

@@ -19,6 +19,7 @@
 
 from st2common import log as logging
 from st2common.models.db import stormbase
+from st2common.fields import JSONDictEscapedFieldCompatibilityField
 from st2common.fields import ComplexDateTimeField
 from st2common.util import date as date_utils
 from st2common.util.secrets import get_secret_parameters
@@ -61,7 +62,7 @@ class ActionExecutionDB(stormbase.StormFoundationDB):
     parameters = stormbase.EscapedDynamicField(
         default={},
         help_text='The key-value pairs passed as to the action runner & action.')
-    result = stormbase.EscapedDynamicField(
+    result = JSONDictEscapedFieldCompatibilityField(
         default={},
         help_text='Action defined result.')
     context = me.DictField(

st2common/st2common/models/db/liveaction.py

@@ -22,6 +22,7 @@
 from st2common.models.db import stormbase
 from st2common.models.db.notification import NotificationSchema
 from st2common.fields import ComplexDateTimeField
+from st2common.fields import JSONDictEscapedFieldCompatibilityField
 from st2common.util import date as date_utils
 from st2common.util.secrets import get_secret_parameters
 from st2common.util.secrets import mask_secret_parameters
@@ -56,7 +57,7 @@ class LiveActionDB(stormbase.StormFoundationDB):
     parameters = stormbase.EscapedDynamicField(
         default={},
         help_text='The key-value pairs passed as to the action runner & execution.')
-    result = stormbase.EscapedDynamicField(
+    result = JSONDictEscapedFieldCompatibilityField(
         default={},
         help_text='Action defined result.')
     context = me.DictField(