From 4c065be92c02fddf239ffb5c1fb347593c22c1ab Mon Sep 17 00:00:00 2001
From: yemohyle <73853907+yemohyle@users.noreply.github.com>
Date: Thu, 4 Feb 2021 18:21:00 -0800
Subject: [PATCH] add local parameter map saving to a local file (#14018)

* add local parameter map saving to a local file

* fix local parameter map saving

* revere serviceUri change in AzPredictorSettings
---
 .../ParameterValuePredictor.cs                | 71 ++++++++++++++++---
 1 file changed, 62 insertions(+), 9 deletions(-)

diff --git a/tools/Az.Tools.Predictor/Az.Tools.Predictor/ParameterValuePredictor.cs b/tools/Az.Tools.Predictor/Az.Tools.Predictor/ParameterValuePredictor.cs
index e0d3d48c63d7..9e708692eb67 100644
--- a/tools/Az.Tools.Predictor/Az.Tools.Predictor/ParameterValuePredictor.cs
+++ b/tools/Az.Tools.Predictor/Az.Tools.Predictor/ParameterValuePredictor.cs
@@ -19,10 +19,10 @@
 using System.Collections.Generic;
 using System.IO;
 using System.Management.Automation.Language;
+using System.Threading.Tasks;
+using System.Threading;
 using System.Text.Json;
 
-
-
 namespace Microsoft.Azure.PowerShell.Tools.AzPredictor
 {
     /// <summary>
@@ -32,7 +32,13 @@ sealed class ParameterValuePredictor
     {
         private readonly ConcurrentDictionary<string, string> _localParameterValues = new ConcurrentDictionary<string, string>();
 
-        private readonly Dictionary<string, Dictionary<string, string>> _command_param_to_resource_map;
+        private System.Threading.Mutex _mutex = new System.Threading.Mutex(false, "paramValueHistoryFile_update");
+
+        private readonly Dictionary<string, Dictionary<string, string>> _commandParamToResourceMap;
+
+        private string _paramValueHistoryFilePath = "";
+        private CancellationTokenSource _cancellationTokenSource;
+
 
         private ITelemetryClient _telemetryClient;
 
@@ -49,17 +55,46 @@ public ParameterValuePredictor(ITelemetryClient telemetryClient)
 
             try
             {
-                _command_param_to_resource_map = JsonSerializer.Deserialize<Dictionary<string, Dictionary<string, string>>>(File.ReadAllText(mappingFilePath), JsonUtilities.DefaultSerializerOptions);
+                _commandParamToResourceMap = JsonSerializer.Deserialize<Dictionary<string, Dictionary<string, string>>>(File.ReadAllText(mappingFilePath), JsonUtilities.DefaultSerializerOptions);
             }
             catch (Exception e)
             {
                 // We don't want it to crash the module when the file doesn't exist or when it's mal-formatted.
                 exception = e;
             }
-
             _telemetryClient.OnLoadParameterMap(new ParameterMapTelemetryData(exception));
+
+            String path = Environment.GetFolderPath(Environment.SpecialFolder.ApplicationData);
+            string[] paths = new string[] { path, "Microsoft", "Windows", "PowerShell", "AzPredictor", "paramValueHistory.json" };
+            _paramValueHistoryFilePath = System.IO.Path.Combine(paths);
+            Directory.CreateDirectory(Path.GetDirectoryName(_paramValueHistoryFilePath));
+
+            Task.Run(() =>
+            {
+                if (System.IO.File.Exists(_paramValueHistoryFilePath))
+                {
+                    _mutex.WaitOne();
+                    try
+                    {
+                        var localParameterValues = JsonSerializer.Deserialize<ConcurrentDictionary<string, string>>(File.ReadAllText(_paramValueHistoryFilePath), JsonUtilities.DefaultSerializerOptions);
+                        foreach (var v in localParameterValues)
+                        {
+                            _localParameterValues.AddOrUpdate(v.Key, key => v.Value, (key, oldValue) => oldValue);
+                        }
+                    }
+                    finally
+                    {
+                        _mutex.ReleaseMutex();
+                    }
+                }
+
+                
+            });
+
         }
 
+        
+
         /// <summary>
         /// Process the command from history
         /// </summary>
@@ -89,7 +124,7 @@ public string GetParameterValueFromAzCommand(string commandNoun, string paramete
             var key = parameterName;
             Dictionary<string, string> commandNounMap = null;
 
-            if (_command_param_to_resource_map?.TryGetValue(commandNoun, out commandNounMap) == true)
+            if (_commandParamToResourceMap?.TryGetValue(commandNoun, out commandNounMap) == true)
             {
                 if (commandNounMap.TryGetValue(parameterName, out var parameterNameMappedValue))
                 {
@@ -146,7 +181,7 @@ private void ExtractLocalParameters(CommandAst command)
             }
 
             Dictionary<string, string> commandNounMap = null;
-            _command_param_to_resource_map?.TryGetValue(commandNoun, out commandNounMap);
+            _commandParamToResourceMap?.TryGetValue(commandNoun, out commandNounMap);
 
             for (int i = 1; i < command.CommandElements.Count;)
             {
@@ -189,8 +224,26 @@ private void ExtractLocalParameters(CommandAst command)
                             parameterKey = mappedValue;
                         }
                     }
-
-                    _localParameterValues.AddOrUpdate(parameterKey, parameterValue, (k, v) => parameterValue);
+                    _cancellationTokenSource?.Cancel();
+                    _cancellationTokenSource = new CancellationTokenSource();
+                    Task.Run(() =>
+                    {
+                        this._localParameterValues.AddOrUpdate(parameterKey, parameterValue, (k, v) => parameterValue);
+                        if (_cancellationTokenSource.IsCancellationRequested)
+                        {
+                            throw new OperationCanceledException();
+                        }
+                        String localParameterValuesJson = JsonSerializer.Serialize<ConcurrentDictionary<string, string>>(_localParameterValues, JsonUtilities.DefaultSerializerOptions);
+                        _mutex.WaitOne();
+                        try
+                        {
+                            System.IO.File.WriteAllText(_paramValueHistoryFilePath, localParameterValuesJson);
+                        }
+                        finally
+                        {
+                            _mutex.ReleaseMutex();
+                        }
+                    });
                 }
             }
         }