scaleset vm instance scale-in protection (#1946)

* scaleset vm instance scale-in protection

* added tests and used to validate the behavior

Co-authored-by: stas <statis@microsoft.com>

src/ApiService/ApiService/OneFuzzTypes/ReturnTypes.cs

@@ -71,6 +71,8 @@ public struct OneFuzzResultVoid {
 
         public Error ErrorV => error;
 
+        public OneFuzzResultVoid() => (error, isOk) = (NoError, true);
+
         private OneFuzzResultVoid(ErrorCode errorCode, string[] errors) => (error, isOk) = (new Error(errorCode, errors), false);
 
         private OneFuzzResultVoid(Error err) => (error, isOk) = (err, false);

src/ApiService/ApiService/TestHooks/VmssTestHooks.cs

@@ -0,0 +1,68 @@
+using System.Net;
+using System.Text.Json;
+using System.Threading.Tasks;
+using Microsoft.Azure.Functions.Worker;
+using Microsoft.Azure.Functions.Worker.Http;
+using Microsoft.OneFuzz.Service;
+using Microsoft.OneFuzz.Service.OneFuzzLib.Orm;
+
+#if DEBUG
+namespace ApiService.TestHooks {
+    public class VmssTestHooks {
+
+        private readonly ILogTracer _log;
+        private readonly IConfigOperations _configOps;
+        private readonly IVmssOperations _vmssOps;
+
+        public VmssTestHooks(ILogTracer log, IConfigOperations configOps, IVmssOperations vmssOps) {
+            _log = log.WithTag("TestHooks", nameof(VmssTestHooks));
+            _configOps = configOps; ;
+            _vmssOps = vmssOps; ;
+        }
+
+
+        [Function("ListInstanceIdsTesHook")]
+        public async Task<HttpResponseData> ListInstanceIds([HttpTrigger(AuthorizationLevel.Anonymous, "get", Route = "testhooks/vmssOperations/listInstanceIds")] HttpRequestData req) {
+            _log.Info("list instance ids");
+            var query = UriExtension.GetQueryComponents(req.Url);
+            var name = UriExtension.GetGuid("name", query) ?? throw new Exception("name must be set");
+            var ids = await _vmssOps.ListInstanceIds(name);
+
+            var json = JsonSerializer.Serialize(ids, EntityConverter.GetJsonSerializerOptions());
+            var resp = req.CreateResponse(HttpStatusCode.OK);
+            await resp.WriteStringAsync(json);
+            return resp;
+        }
+
+        [Function("GetInstanceIdsTesHook")]
+        public async Task<HttpResponseData> GetInstanceId([HttpTrigger(AuthorizationLevel.Anonymous, "get", Route = "testhooks/vmssOperations/getInstanceId")] HttpRequestData req) {
+            _log.Info("list instance ids");
+            var query = UriExtension.GetQueryComponents(req.Url);
+            var name = UriExtension.GetGuid("name", query) ?? throw new Exception("name must be set");
+            var vmId = UriExtension.GetGuid("vmId", query) ?? throw new Exception("vmId must be set");
+            var id = await _vmssOps.GetInstanceId(name, vmId);
+
+            var json = JsonSerializer.Serialize(id, EntityConverter.GetJsonSerializerOptions());
+            var resp = req.CreateResponse(HttpStatusCode.OK);
+            await resp.WriteStringAsync(json);
+            return resp;
+        }
+
+        [Function("UpdateScaleInProtectionTestHook")]
+        public async Task<HttpResponseData> UpdateScaleInProtection([HttpTrigger(AuthorizationLevel.Anonymous, "put", Route = "testhooks/vmssOperations/updateScaleInProtection")] HttpRequestData req) {
+            _log.Info("list instance ids");
+            var query = UriExtension.GetQueryComponents(req.Url);
+            var name = UriExtension.GetGuid("name", query) ?? throw new Exception("name must be set");
+            var vmId = UriExtension.GetGuid("vmId", query) ?? throw new Exception("vmId must be set");
+            var protectFromScaleIn = UriExtension.GetBool("protectFromScaleIn", query);
+            var id = await _vmssOps.UpdateScaleInProtection(name, vmId, protectFromScaleIn);
+
+            var json = JsonSerializer.Serialize(id, EntityConverter.GetJsonSerializerOptions());
+            var resp = req.CreateResponse(HttpStatusCode.OK);
+            await resp.WriteStringAsync(json);
+            return resp;
+        }
+    }
+}
+
+#endif

src/ApiService/ApiService/onefuzzlib/VmssOperations.cs

@@ -19,6 +19,8 @@ public interface IVmssOperations {
 
 public class VmssOperations : IVmssOperations {
 
+    string INSTANCE_NOT_FOUND = " is not an active Virtual Machine Scale Set VM instanceId.";
+
     ILogTracer _log;
     ICreds _creds;
 
@@ -105,7 +107,7 @@ public async Async.Task<IDictionary<Guid, string>> ListInstanceIds(Guid name) {
                         }
                     }
                 }
-            } catch (CloudException ex) {
+            } catch (Exception ex) when (ex is RequestFailedException || ex is CloudException) {
                 _log.Exception(ex, $"vm does not exist {name}");
             }
         }
@@ -126,7 +128,7 @@ public async Async.Task<OneFuzzResult<VirtualMachineScaleSetVmResource>> GetInst
                     return OneFuzzResult<VirtualMachineScaleSetVmResource>.Ok(response);
                 }
             }
-        } catch (CloudException ex) {
+        } catch (Exception ex) when (ex is RequestFailedException || ex is CloudException) {
             _log.Exception(ex, $"unable to find vm instance: {name}:{vmId}");
             return OneFuzzResult<VirtualMachineScaleSetVmResource>.Error(ErrorCode.UNABLE_TO_FIND, $"unable to find vm instance: {name}:{vmId}");
         }
@@ -159,14 +161,27 @@ public async Async.Task<OneFuzzResultVoid> UpdateScaleInProtection(Guid name, Gu
             instanceVm.Data.ProtectionPolicy = newProtectionPolicy;
 
             var scaleSet = GetVmssResource(name);
-
-            VirtualMachineScaleSetVmInstanceRequiredIds ids = new VirtualMachineScaleSetVmInstanceRequiredIds(new[] { instanceVm.Data.InstanceId });
-            var updateRes = await scaleSet.UpdateInstancesAsync(WaitUntil.Started, ids);
-
-            //TODO: finish this after UpdateInstance method is fixed
-            //https://github.com/Azure/azure-sdk-for-net/issues/28491
-
-            throw new NotImplementedException("Update instance does not work as expected. See https://github.com/Azure/azure-sdk-for-net/issues/28491");
+            var vmCollection = scaleSet.GetVirtualMachineScaleSetVms();
+            try {
+                var r = await vmCollection.CreateOrUpdateAsync(WaitUntil.Started, instanceVm.Data.InstanceId, instanceVm.Data);
+                if (r.GetRawResponse().IsError) {
+                    var msg = $"failed to update scale in protection on vm {vmId} for scaleset {name}";
+                    _log.WithHttpStatus((r.GetRawResponse().Status, r.GetRawResponse().ReasonPhrase)).Error(msg);
+                    return OneFuzzResultVoid.Error(ErrorCode.UNABLE_TO_UPDATE, msg);
+                } else {
+                    return OneFuzzResultVoid.Ok();
+                }
+            } catch (Exception ex) when (ex is RequestFailedException || ex is CloudException) {
+
+                if (ex.Message.Contains(INSTANCE_NOT_FOUND) && protectFromScaleIn == false) {
+                    _log.Info($"Tried to remove scale in protection on node {name} {vmId} but instance no longer exists");
+                    return OneFuzzResultVoid.Ok();
+                } else {
+                    var msg = $"failed to update scale in protection on vm {vmId} for scaleset {name}";
+                    _log.Exception(ex, msg);
+                    return OneFuzzResultVoid.Error(ErrorCode.UNABLE_TO_UPDATE, ex.Message);
+                }
+            }
         }
     }