From e90a0e89f7c515387c70df1de4b4c0769330f20b Mon Sep 17 00:00:00 2001
From: Feng Zhou <55177366+fengzhou-msft@users.noreply.github.com>
Date: Mon, 27 Jul 2020 07:50:41 +0000
Subject: [PATCH] {Telemetry} Add command init and invoke time in telemetry
 (#14273)

---
 .flake8                                        |  3 +++
 src/azure-cli-core/azure/cli/core/telemetry.py | 14 ++++++++++++++
 src/azure-cli/azure/cli/__main__.py            | 15 +++++++++------
 3 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/.flake8 b/.flake8
index 2c402923c..2f18f7296 100644
--- a/.flake8
+++ b/.flake8
@@ -8,6 +8,9 @@ ignore =
     F811,  # redefinition of unused, to be removed in the future
     C901   # code flow is too complex, too many violations, to be removed in the future
     W504   # line break after binary operator effect on readability is subjective
+per-file-ignores =
+    # module level import not at top of file, elevate timer to measure import time
+    src/azure-cli/azure/cli/__main__.py:E402
 exclude = 
     azure_bdist_wheel.py
     build
diff --git a/src/azure-cli-core/azure/cli/core/telemetry.py b/src/azure-cli-core/azure/cli/core/telemetry.py
index b36f8a6f8..cc10885e8 100644
--- a/src/azure-cli-core/azure/cli/core/telemetry.py
+++ b/src/azure-cli-core/azure/cli/core/telemetry.py
@@ -45,6 +45,8 @@ class TelemetrySession:  # pylint: disable=too-many-instance-attributes
         self.extension_management_detail = None
         self.raw_command = None
         self.mode = 'default'
+        self.init_time_elapsed = None
+        self.invoke_time_elapsed = None
         # A dictionary with the application insight instrumentation key
         # as the key and an array of telemetry events as value
         self.events = defaultdict(list)
@@ -164,6 +166,8 @@ class TelemetrySession:  # pylint: disable=too-many-instance-attributes
                               lambda: '{},{}'.format(locale.getdefaultlocale()[0], locale.getdefaultlocale()[1]))
         set_custom_properties(result, 'StartTime', str(self.start_time))
         set_custom_properties(result, 'EndTime', str(self.end_time))
+        set_custom_properties(result, 'InitTimeElapsed', str(self.init_time_elapsed))
+        set_custom_properties(result, 'InvokeTimeElapsed', str(self.invoke_time_elapsed))
         set_custom_properties(result, 'OutputType', self.output_type)
         set_custom_properties(result, 'RawCommand', self.raw_command)
         set_custom_properties(result, 'Params', ','.join(self.parameters or []))
@@ -227,6 +231,16 @@ def start(mode=None):
     _session.start_time = datetime.datetime.utcnow()
 
 
+@decorators.suppress_all_exceptions()
+def set_init_time_elapsed(init_time_elapsed):
+    _session.init_time_elapsed = init_time_elapsed
+
+
+@decorators.suppress_all_exceptions()
+def set_invoke_time_elapsed(invoke_time_elapsed):
+    _session.invoke_time_elapsed = invoke_time_elapsed
+
+
 @_user_agrees_to_telemetry
 @decorators.suppress_all_exceptions()
 def flush():
diff --git a/src/azure-cli/azure/cli/__main__.py b/src/azure-cli/azure/cli/__main__.py
index 1558a8266..06eef18cc 100644
--- a/src/azure-cli/azure/cli/__main__.py
+++ b/src/azure-cli/azure/cli/__main__.py
@@ -2,9 +2,13 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
+# pylint: disable=wrong-import-position
+
+import timeit
+# Log the start time
+start_time = timeit.default_timer()
 
 import sys
-import timeit
 import uuid
 
 import azure.cli.core.telemetry as telemetry
@@ -12,9 +16,6 @@ from azure.cli.core import get_default_cli
 from knack.completion import ARGCOMPLETE_ENV_NAME
 from knack.log import get_logger
 
-# Log the start time
-# TODO: Disable E402 "Module level import not at top of file" so that imports can also be timed
-start_time = timeit.default_timer()
 
 # A workaround for https://bugs.python.org/issue32502 (https://github.com/Azure/azure-cli/issues/5184)
 # If uuid1 raises ValueError, use uuid4 instead.
@@ -60,8 +61,6 @@ except SystemExit as ex:  # some code directly call sys.exit, this is to make su
     raise ex
 
 finally:
-    telemetry.conclude()
-
     try:
         # Log the invoke finish time
         invoke_finish_time = timeit.default_timer()
@@ -71,3 +70,7 @@ finally:
                     invoke_finish_time - init_finish_time)
     except NameError:
         pass
+
+    telemetry.set_init_time_elapsed("{:.6f}".format(init_finish_time - start_time))
+    telemetry.set_invoke_time_elapsed("{:.6f}".format(invoke_finish_time - init_finish_time))
+    telemetry.conclude()