Repository: Azure/azure-cli

Issue:
auth: cache adal token cache in memory
The goal is to use a single copy of token cache across the same process, rather dump it after getting a token. Building up a token cache is non trivial involving reading and writing the ~/.azure/accesstoken.json. For some commands which involves 1000+ rest call, like datalake, the perf hit is not trivial.

Discussion:
codecov-io: # [Codecov](https://codecov.io/gh/Azure/azure-cli/pull/3164?src=pr&el=h1) Report
> Merging [#3164](https://codecov.io/gh/Azure/azure-cli/pull/3164?src=pr&el=desc) into [master](https://codecov.io/gh/Azure/azure-cli/commit/0c5e91ca52cf0ed5f9284968dd3272275020ce89?src=pr&el=desc) will **increase** coverage by `<.01%`.
> The diff coverage is `92.59%`.

[![Impacted file tree graph](https://codecov.io/gh/Azure/azure-cli/pull/3164/graphs/tree.svg?width=650&height=150&src=pr&token=2pog0TKvF8)](https://codecov.io/gh/Azure/azure-cli/pull/3164?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master    #3164      +/-   ##
==========================================
+ Coverage   71.11%   71.12%   +<.01%     
==========================================
  Files         362      362              
  Lines       23914    23929      +15     
  Branches     3653     3658       +5     
==========================================
+ Hits        17006    17019      +13     
  Misses       5812     5812              
- Partials     1096     1098       +2
```


| [Impacted Files](https://codecov.io/gh/Azure/azure-cli/pull/3164?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [src/azure-cli-core/azure/cli/core/\_profile.py](https://codecov.io/gh/Azure/azure-cli/pull/3164?src=pr&el=tree#diff-c3JjL2F6dXJlLWNsaS1jb3JlL2F6dXJlL2NsaS9jb3JlL19wcm9maWxlLnB5) | `87.64% <92.59%> (-0.04%)` | :arrow_down: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/Azure/azure-cli/pull/3164?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/Azure/azure-cli/pull/3164?src=pr&el=footer). Last update [0c5e91c...d097884](https://codecov.io/gh/Azure/azure-cli/pull/3164?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Patch:
diff --git a/src/azure-cli-core/HISTORY.rst b/src/azure-cli-core/HISTORY.rst
index 0db20c494..49f585169 100644
--- a/src/azure-cli-core/HISTORY.rst
+++ b/src/azure-cli-core/HISTORY.rst
@@ -6,6 +6,7 @@ Release History
 ^^^^^^^^^^^^^^^^^^
 * core: capture exceptions caused by unregistered provider and auto-register it   
 * login: avoid the bad exception when the user account has no subscription and no tenants
+* perf: persist adal token cache in memory till process exits (#2603)
 
 2.0.4 (2017-04-28)
 ^^^^^^^^^^^^^^^^^^
diff --git a/src/azure-cli-core/azure/cli/core/_profile.py b/src/azure-cli-core/azure/cli/core/_profile.py
index b5a6c77a6..4cae579da 100644
--- a/src/azure-cli-core/azure/cli/core/_profile.py
+++ b/src/azure-cli-core/azure/cli/core/_profile.py
@@ -93,11 +93,20 @@ class CredentialType(Enum):  # pylint: disable=too-few-public-methods
     rbac = CLOUD.endpoints.active_directory_graph_resource_id
 
 
+_GLOBAL_CREDS_CACHE = None
+
+
 class Profile(object):
-    def __init__(self, storage=None, auth_ctx_factory=None):
+    def __init__(self, storage=None, auth_ctx_factory=None, use_global_creds_cache=True):
         self._storage = storage or ACCOUNT
         self.auth_ctx_factory = auth_ctx_factory or _AUTH_CTX_FACTORY
-        self._creds_cache = CredsCache(self.auth_ctx_factory)
+        if use_global_creds_cache:
+            global _GLOBAL_CREDS_CACHE  # pylint: disable=global-statement
+            if _GLOBAL_CREDS_CACHE is None:
+                _GLOBAL_CREDS_CACHE = CredsCache(self.auth_ctx_factory, async_persist=True)
+            self._creds_cache = _GLOBAL_CREDS_CACHE
+        else:
+            self._creds_cache = CredsCache(self.auth_ctx_factory, async_persist=False)
         self._management_resource_uri = CLOUD.endpoints.management
         self._ad_resource_uri = CLOUD.endpoints.active_directory_resource_id
 
@@ -470,29 +479,39 @@ class CredsCache(object):
     also be handled
     '''
 
-    def __init__(self, auth_ctx_factory=None):
+    def __init__(self, auth_ctx_factory=None, async_persist=True):
         # AZURE_ACCESS_TOKEN_FILE is used by Cloud Console and not meant to be user configured
         self._token_file = (os.environ.get('AZURE_ACCESS_TOKEN_FILE', None) or
                             os.path.join(get_config_dir(), 'accessTokens.json'))
         self._service_principal_creds = []
         self._auth_ctx_factory = auth_ctx_factory or _AUTH_CTX_FACTORY
         self._adal_token_cache_attr = None
+        self._should_flush_to_disk = False
+        self._async_persist = async_persist
+        if async_persist:
+            import atexit
+            atexit.register(self.flush_to_disk)
 
     def persist_cached_creds(self):
-        with os.fdopen(os.open(self._token_file, os.O_RDWR | os.O_CREAT | os.O_TRUNC, 0o600),
-                       'w+') as cred_file:
-            items = self.adal_token_cache.read_items()
-            all_creds = [entry for _, entry in items]
+        self._should_flush_to_disk = True
+        if not self._async_persist:
+            self.flush_to_disk()
+        self.adal_token_cache.has_state_changed = False
 
-            # trim away useless fields (needed for cred sharing with xplat)
-            for i in all_creds:
-                for key in TOKEN_FIELDS_EXCLUDED_FROM_PERSISTENCE:
-                    i.pop(key, None)
+    def flush_to_disk(self):
+        if self._should_flush_to_disk:
+            with os.fdopen(os.open(self._token_file, os.O_RDWR | os.O_CREAT | os.O_TRUNC, 0o600),
+                           'w+') as cred_file:
+                items = self.adal_token_cache.read_items()
+                all_creds = [entry for _, entry in items]
 
-            all_creds.extend(self._service_principal_creds)
-            cred_file.write(json.dumps(all_creds))
+                # trim away useless fields (needed for cred sharing with xplat)
+                for i in all_creds:
+                    for key in TOKEN_FIELDS_EXCLUDED_FROM_PERSISTENCE:
+                        i.pop(key, None)
 
-        self.adal_token_cache.has_state_changed = False
+                all_creds.extend(self._service_principal_creds)
+                cred_file.write(json.dumps(all_creds))
 
     def retrieve_token_for_user(self, username, tenant, resource):
         authority = get_authority_url(tenant)


Test patch:
diff --git a/src/azure-cli-core/tests/test_profile.py b/src/azure-cli-core/tests/test_profile.py
index 6f25f61de..ea3648494 100644
--- a/src/azure-cli-core/tests/test_profile.py
+++ b/src/azure-cli-core/tests/test_profile.py
@@ -76,7 +76,7 @@ class Test_Profile(unittest.TestCase):  # pylint: disable=too-many-public-method
 
     def test_update_add_two_different_subscriptions(self):
         storage_mock = {'subscriptions': None}
-        profile = Profile(storage_mock)
+        profile = Profile(storage_mock, use_global_creds_cache=False)
 
         # add the first and verify
         consolidated = Profile._normalize_properties(self.user1,
@@ -127,7 +127,7 @@ class Test_Profile(unittest.TestCase):  # pylint: disable=too-many-public-method
 
     def test_update_with_same_subscription_added_twice(self):
         storage_mock = {'subscriptions': None}
-        profile = Profile(storage_mock)
+        profile = Profile(storage_mock, use_global_creds_cache=False)
 
         # add one twice and verify we will have one but with new token
         consolidated = Profile._normalize_properties(self.user1,
@@ -149,7 +149,7 @@ class Test_Profile(unittest.TestCase):  # pylint: disable=too-many-public-method
 
     def test_set_active_subscription(self):
         storage_mock = {'subscriptions': None}
-        profile = Profile(storage_mock)
+        profile = Profile(storage_mock, use_global_creds_cache=False)
 
         consolidated = Profile._normalize_properties(self.user1,
                                                      [self.subscription1],
@@ -169,7 +169,7 @@ class Test_Profile(unittest.TestCase):  # pylint: disable=too-many-public-method
 
     def test_default_active_subscription_to_non_disabled_one(self):
         storage_mock = {'subscriptions': None}
-        profile = Profile(storage_mock)
+        profile = Profile(storage_mock, use_global_creds_cache=False)
 
         subscriptions = profile._normalize_properties(
             self.user2, [self.subscription2, self.subscription1], False)
@@ -182,7 +182,7 @@ class Test_Profile(unittest.TestCase):  # pylint: disable=too-many-public-method
 
     def test_get_subscription(self):
         storage_mock = {'subscriptions': None}
-        profile = Profile(storage_mock)
+        profile = Profile(storage_mock, use_global_creds_cache=False)
 
         consolidated = Profile._normalize_properties(self.user1,
                                                      [self.subscription1],
@@ -200,7 +200,7 @@ class Test_Profile(unittest.TestCase):  # pylint: disable=too-many-public-method
 
     def test_get_expanded_subscription_info(self):
         storage_mock = {'subscriptions': None}
-        profile = Profile(storage_mock)
+        profile = Profile(storage_mock, use_global_creds_cache=False)
 
         consolidated = Profile._normalize_properties(self.user1,
                                                      [self.subscription1],
@@ -236,7 +236,7 @@ class Test_Profile(unittest.TestCase):  # pylint: disable=too-many-public-method
                                     lambda _: mock_arm_client)
 
         storage_mock = {'subscriptions': []}
-        profile = Profile(storage_mock)
+        profile = Profile(storage_mock, use_global_creds_cache=False)
         profile._management_resource_uri = 'https://management.core.windows.net/'
         profile.find_subscriptions_on_login(False,
                                             '1234',
@@ -264,7 +264,7 @@ class Test_Profile(unittest.TestCase):  # pylint: disable=too-many-public-method
                                     lambda _: mock_arm_client)
 
         storage_mock = {'subscriptions': []}
-        profile = Profile(storage_mock)
+        profile = Profile(storage_mock, use_global_creds_cache=False)
         profile._management_resource_uri = 'https://management.core.windows.net/'
 
         # action
@@ -288,7 +288,7 @@ class Test_Profile(unittest.TestCase):  # pylint: disable=too-many-public-method
         finder = mock.MagicMock()
         finder.find_through_interactive_flow.return_value = []
         storage_mock = {'subscriptions': []}
-        profile = Profile(storage_mock)
+        profile = Profile(storage_mock, use_global_creds_cache=False)
 
         # action
         result = profile.find_subscriptions_on_login(True,
@@ -308,7 +308,7 @@ class Test_Profile(unittest.TestCase):  # pylint: disable=too-many-public-method
         mock_read_cred_file.return_value = [Test_Profile.token_entry1]
 
         storage_mock = {'subscriptions': None}
-        profile = Profile(storage_mock)
+        profile = Profile(storage_mock, use_global_creds_cache=False)
         consolidated = Profile._normalize_properties(self.user1,
                                                      [self.subscription1],
                                                      False)
@@ -322,7 +322,7 @@ class Test_Profile(unittest.TestCase):  # pylint: disable=too-many-public-method
     @mock.patch('azure.cli.core._profile._load_tokens_from_file', return_value=None)
     def test_create_token_cache(self, mock_read_file):
         mock_read_file.return_value = []
-        profile = Profile()
+        profile = Profile(use_global_creds_cache=False)
         cache = profile._creds_cache.adal_token_cache
         self.assertFalse(cache.read_items())
         self.assertTrue(mock_read_file.called)
@@ -330,7 +330,7 @@ class Test_Profile(unittest.TestCase):  # pylint: disable=too-many-public-method
     @mock.patch('azure.cli.core._profile._load_tokens_from_file', autospec=True)
     def test_load_cached_tokens(self, mock_read_file):
         mock_read_file.return_value = [Test_Profile.token_entry1]
-        profile = Profile()
+        profile = Profile(use_global_creds_cache=False)
         cache = profile._creds_cache.adal_token_cache
         matched = cache.find({
             "_authority": "https://login.microsoftonline.com/common",
@@ -348,7 +348,7 @@ class Test_Profile(unittest.TestCase):  # pylint: disable=too-many-public-method
         mock_get_token.return_value = (some_token_type, Test_Profile.raw_token1)
         # setup
         storage_mock = {'subscriptions': None}
-        profile = Profile(storage_mock)
+        profile = Profile(storage_mock, use_global_creds_cache=False)
         consolidated = Profile._normalize_properties(self.user1,
                                                      [self.subscription1],
                                                      False)
@@ -375,7 +375,7 @@ class Test_Profile(unittest.TestCase):  # pylint: disable=too-many-public-method
         mock_get_token.return_value = (some_token_type, Test_Profile.raw_token1)
         # setup
         storage_mock = {'subscriptions': None}
-        profile = Profile(storage_mock)
+        profile = Profile(storage_mock, use_global_creds_cache=False)
         consolidated = Profile._normalize_properties(self.user1, [self.subscription1],
                                                      False)
         profile._set_subscriptions(consolidated)
@@ -395,7 +395,7 @@ class Test_Profile(unittest.TestCase):  # pylint: disable=too-many-public-method
         mock_read_cred_file.return_value = [Test_Profile.token_entry1]
 
         storage_mock = {'subscriptions': None}
-        profile = Profile(storage_mock)
+        profile = Profile(storage_mock, use_global_creds_cache=False)
         consolidated = Profile._normalize_properties(self.user1,
                                                      [self.subscription1],
                                                      False)
@@ -413,7 +413,7 @@ class Test_Profile(unittest.TestCase):  # pylint: disable=too-many-public-method
     def test_logout_all(self, mock_delete_cred_file):
         # setup
         storage_mock = {'subscriptions': None}
-        profile = Profile(storage_mock)
+        profile = Profile(storage_mock, use_global_creds_cache=False)
         consolidated = Profile._normalize_properties(self.user1,
                                                      [self.subscription1],
                                                      False)
@@ -584,7 +584,7 @@ class Test_Profile(unittest.TestCase):  # pylint: disable=too-many-public-method
         mock_read_file.return_value = [self.token_entry1, test_sp]
 
         # action
-        creds_cache = CredsCache()
+        creds_cache = CredsCache(async_persist=False)
 
         # assert
         token_entries = [entry for _, entry in creds_cache.load_adal_token_cache().read_items()]
@@ -601,7 +601,7 @@ class Test_Profile(unittest.TestCase):  # pylint: disable=too-many-public-method
         mock_read_file.return_value = [test_sp]
 
         # action
-        creds_cache = CredsCache()
+        creds_cache = CredsCache(async_persist=False)
         creds_cache.load_adal_token_cache()
 
         # assert
@@ -623,7 +623,7 @@ class Test_Profile(unittest.TestCase):  # pylint: disable=too-many-public-method
         }
         mock_open_for_write.return_value = FileHandleStub()
         mock_read_file.return_value = [self.token_entry1, test_sp]
-        creds_cache = CredsCache()
+        creds_cache = CredsCache(async_persist=False)
 
         # action
         creds_cache.save_service_principal_cred(test_sp2)
@@ -645,7 +645,7 @@ class Test_Profile(unittest.TestCase):  # pylint: disable=too-many-public-method
         }
         mock_open_for_write.return_value = FileHandleStub()
         mock_read_file.return_value = [test_sp]
-        creds_cache = CredsCache()
+        creds_cache = CredsCache(async_persist=False)
 
         # action
         creds_cache.save_service_principal_cred(test_sp)
@@ -664,7 +664,7 @@ class Test_Profile(unittest.TestCase):  # pylint: disable=too-many-public-method
         }
         mock_open_for_write.return_value = FileHandleStub()
         mock_read_file.return_value = [self.token_entry1, test_sp]
-        creds_cache = CredsCache()
+        creds_cache = CredsCache(async_persist=False)
 
         # action #1, logout a user
         creds_cache.remove_cached_creds(self.user1)
@@ -704,7 +704,7 @@ class Test_Profile(unittest.TestCase):  # pylint: disable=too-many-public-method
         mock_adal_auth_context.acquire_token.side_effect = acquire_token_side_effect
         mock_open_for_write.return_value = FileHandleStub()
         mock_read_file.return_value = [self.token_entry1]
-        creds_cache = CredsCache(auth_ctx_factory=get_auth_context)
+        creds_cache = CredsCache(auth_ctx_factory=get_auth_context, async_persist=False)
 
         # action
         mgmt_resource = 'https://management.core.windows.net/'