From 064689d44221452f434d0c7ea5e2ee752999fe83 Mon Sep 17 00:00:00 2001
From: Stepan Blyshchak <38952541+stepanblyschak@users.noreply.github.com>
Date: Thu, 31 Oct 2019 18:17:29 +0200
Subject: [PATCH] [sonic-cfggen] optimize sonic-cfggen startup (#3658)

* [sonic-cfggen] optimize execution time

a lot of template rendering causes switch to start longer because jinja2
needs to parse them. Introducing RedisBytecodeCache to store parsed buckets of
internal template bytecode to speedup same template rendering during start

* [sonic-cfggen] do lazy regexp compilation to speedup sonic-cfggen

* [sonic-cfggen] address pep8 related comments

Signed-off-by: Stepan Blyschak <stepanb@mellanox.com>
---
 src/sonic-config-engine/lazy_re.py   | 22 ++++++++++++++++++++++
 src/sonic-config-engine/redis_bcc.py | 26 ++++++++++++++++++++++++++
 src/sonic-config-engine/setup.py     |  2 +-
 src/sonic-config-engine/sonic-cfggen | 16 ++++++++++++++--
 4 files changed, 63 insertions(+), 3 deletions(-)
 create mode 100644 src/sonic-config-engine/lazy_re.py
 create mode 100644 src/sonic-config-engine/redis_bcc.py

diff --git a/src/sonic-config-engine/lazy_re.py b/src/sonic-config-engine/lazy_re.py
new file mode 100644
index 0000000000..b51c385c1c
--- /dev/null
+++ b/src/sonic-config-engine/lazy_re.py
@@ -0,0 +1,22 @@
+# monkey patch re.compile to improve import time of some packages
+
+import re
+
+_orig_re_compile = re.compile
+
+
+def __re_compile(*args, **kwargs):
+    class __LazyReCompile(object):
+        def __init__(self, *args, **kwargs):
+            self.args = args
+            self.kwargs = kwargs
+            self.pattern_obj = None
+
+        def __getattr__(self, name):
+            if self.pattern_obj is None:
+                self.pattern_obj = _orig_re_compile(*self.args, **self.kwargs)
+            return getattr(self.pattern_obj, name)
+    return __LazyReCompile(*args, **kwargs)
+
+re.compile = __re_compile
+
diff --git a/src/sonic-config-engine/redis_bcc.py b/src/sonic-config-engine/redis_bcc.py
new file mode 100644
index 0000000000..5ab59b6a69
--- /dev/null
+++ b/src/sonic-config-engine/redis_bcc.py
@@ -0,0 +1,26 @@
+import jinja2
+
+class RedisBytecodeCache(jinja2.BytecodeCache):
+    """ A bytecode cache for jinja2 template that stores bytecode in Redis """
+
+    REDIS_HASH = 'JINJA2_CACHE'
+
+    def __init__(self, client):
+        self._client = client
+        try:
+            self._client.connect(self._client.STATE_DB, retry_on=False)
+        except Exception:
+            self._client = None
+
+    def load_bytecode(self, bucket):
+        if self._client is None:
+            return
+        code = self._client.get(self._client.STATE_DB, self.REDIS_HASH, bucket.key)
+        if code is not None:
+            bucket.bytecode_from_string(code)
+
+    def dump_bytecode(self, bucket):
+        if self._client is None:
+            return
+        self._client.set(self._client.STATE_DB, self.REDIS_HASH, bucket.key, bucket.bytecode_to_string())
+
diff --git a/src/sonic-config-engine/setup.py b/src/sonic-config-engine/setup.py
index 7ca810ce6a..c75a5b5a03 100755
--- a/src/sonic-config-engine/setup.py
+++ b/src/sonic-config-engine/setup.py
@@ -16,7 +16,7 @@ setup(name='sonic-config-engine',
       author='Taoyu Li',
       author_email='taoyl@microsoft.com',
       url='https://github.com/Azure/sonic-buildimage',
-      py_modules=['portconfig', 'minigraph', 'openconfig_acl', 'sonic_device_util', 'config_samples'],
+      py_modules=['portconfig', 'minigraph', 'openconfig_acl', 'sonic_device_util', 'config_samples', 'redis_bcc', 'lazy_re'],
       scripts=['sonic-cfggen'],
       install_requires=['lxml', 'jinja2>=2.10', 'netaddr', 'ipaddr', 'pyyaml', 'pyangbind==0.6.0'],
       test_suite='setup.get_test_suite',
diff --git a/src/sonic-config-engine/sonic-cfggen b/src/sonic-config-engine/sonic-cfggen
index ca550b49e4..7cff6c9fb3 100755
--- a/src/sonic-config-engine/sonic-cfggen
+++ b/src/sonic-config-engine/sonic-cfggen
@@ -16,6 +16,16 @@ See usage string for detail description for arguments.
 """
 
 from __future__ import print_function
+
+# monkey patch re.compile to do lazy regular expression compilation.
+# This is done to improve import time of jinja2, yaml, natsort modules, because they
+# do many regexp compilation at import time, so it will speed up sonic-cfggen invocations
+# that do not require template generation or yaml loading. sonic-cfggen is used in so many places
+# during system boot up that importing jinja2, yaml, natsort every time
+# without lazy regular expression compilation affect boot up time.
+# FIXME: remove this once sonic-cfggen and templates dependencies are replaced with a faster approach
+import lazy_re
+
 import sys
 import os.path
 import argparse
@@ -33,7 +43,8 @@ from sonic_device_util import get_platform_info
 from sonic_device_util import get_system_mac
 from config_samples import generate_sample_config
 from config_samples import get_available_config
-from swsssdk import ConfigDBConnector
+from swsssdk import SonicV2Connector, ConfigDBConnector
+from redis_bcc import RedisBytecodeCache
 from collections import OrderedDict
 from natsort import natsorted
 
@@ -267,7 +278,8 @@ def main():
         paths  = ['/', '/usr/share/sonic/templates', os.path.dirname(template_file)]
         loader = jinja2.FileSystemLoader(paths)
 
-        env = jinja2.Environment(loader=loader, trim_blocks=True)
+        redis_bcc = RedisBytecodeCache(SonicV2Connector(host='127.0.0.1'))
+        env = jinja2.Environment(loader=loader, trim_blocks=True, bytecode_cache=redis_bcc)
         env.filters['sort_by_port_index'] = sort_by_port_index
         env.filters['ipv4'] = is_ipv4
         env.filters['ipv6'] = is_ipv6