Fix bug for scale-in protection (#3144)

* Fix bug for scale-in protection * Update metric name
microsoft · May 31, 2023 · 0f6073b · 0f6073b
1 parent 98007be
commit 0f6073b
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 1 deletion.
diff --git a/src/ApiService/ApiService/Functions/AgentCanSchedule.cs b/src/ApiService/ApiService/Functions/AgentCanSchedule.cs
@@ -62,6 +62,14 @@ private async Async.Task<HttpResponseData> Post(HttpRequestData req) {
         }
         _ = scp.OkV; // node could be updated but we don't use it after this
         allowed = scp.IsOk;
+
+        if (allowed) {
+            _log.Metric($"TaskAllowedToSchedule", 1, new Dictionary<string, string> {
+                {"MachineId", node.MachineId.ToString()},
+                {"TaskId", task is not null ? task.TaskId.ToString() : string.Empty}
+            });
+        }
+
         return await RequestHandling.Ok(req, new CanSchedule(Allowed: allowed, WorkStopped: workStopped, Reason: reason));
     }
 }
diff --git a/src/ApiService/ApiService/onefuzzlib/NodeOperations.cs b/src/ApiService/ApiService/onefuzzlib/NodeOperations.cs
@@ -91,6 +91,9 @@ public async Task<OneFuzzResult<Node>> AcquireScaleInProtection(Node node) {
         if (node.ScalesetId is ScalesetId scalesetId &&
             await TryGetNodeInfo(node) is NodeInfo nodeInfo) {
 
+            var metricDimensions = new Dictionary<string, string> {
+                {"MachineId", node.MachineId.ToString()}
+            };
             _logTracer.Info($"Setting scale-in protection on node {node.MachineId:Tag:MachineId}");
 
             var instanceId = node.InstanceId;
@@ -110,17 +113,25 @@ await TryGetNodeInfo(node) is NodeInfo nodeInfo) {
             var r = await _context.VmssOperations.UpdateScaleInProtection(nodeInfo.Scaleset, instanceId, protectFromScaleIn: true);
             if (!r.IsOk) {
                 _logTracer.Error(r.ErrorV);
+                _logTracer.Metric($"FailedAcquiringScaleInProtection", 1, metricDimensions);
+                return r.ErrorV;
             }
+
+            _logTracer.Metric($"AcquiredScaleInProtection", 1, metricDimensions);
+            return OneFuzzResult.Ok(node);
         }
 
-        return OneFuzzResult.Ok(node);
+        return Error.Create(ErrorCode.INVALID_NODE, "Failed getting NodeInfo. Cannot acquire scale-in protection");
     }
 
     public async Task<OneFuzzResultVoid> ReleaseScaleInProtection(Node node) {
         if (!node.DebugKeepNode &&
             node.ScalesetId is ScalesetId scalesetId &&
             await TryGetNodeInfo(node) is NodeInfo nodeInfo) {
 
+            var metricDimensions = new Dictionary<string, string> {
+                {"MachineId", node.MachineId.ToString()}
+            };
             _logTracer.Info($"Removing scale-in protection on node {node.MachineId:Tag:MachineId}");
 
             var instanceId = node.InstanceId;
@@ -136,7 +147,11 @@ await TryGetNodeInfo(node) is NodeInfo nodeInfo) {
             var r = await _context.VmssOperations.UpdateScaleInProtection(nodeInfo.Scaleset, instanceId, protectFromScaleIn: false);
             if (!r.IsOk) {
                 _logTracer.Error(r.ErrorV);
+                _logTracer.Metric($"FailedReleasingScaleInProtection", 1, metricDimensions);
+                return r;
             }
+
+            _logTracer.Metric($"ReleasedScaleInProection", 1, metricDimensions);
             return r;
         }