Replace queue_task_hearbeat (#1899)

* Replace queue_task_hearbeat

* dont rename statsFormat

* using hashsets for the helpers

src/ApiService/ApiService/AgentCanSchedule.cs

@@ -54,7 +54,7 @@ public class AgentCanSchedule {
         }
 
         var task = await _taskOperations.GetByTaskId(canScheduleRequest.TaskId);
-        workStopped = task == null || TaskStateHelper.ShuttingDown().Contains(task.State);
+        workStopped = task == null || TaskStateHelper.ShuttingDown.Contains(task.State);
 
         if (allowed) {
             allowed = (await _nodeOperations.AcquireScaleInProtection(node)).IsOk;

src/ApiService/ApiService/OneFuzzTypes/Enums.cs

@@ -1,5 +1,4 @@
-using System.Collections.Concurrent;
-using Microsoft.OneFuzz.Service.OneFuzzLib.Orm;
+using Microsoft.OneFuzz.Service.OneFuzzLib.Orm;
 
 namespace Microsoft.OneFuzz.Service;
 
@@ -96,6 +95,7 @@ public enum ContainerType {
 }
 
 
+[SkipRename]
 public enum StatsFormat {
     AFL
 }
@@ -136,133 +136,55 @@ public static class JobStateHelper {
 
 
 public static class ScalesetStateHelper {
-    static ConcurrentDictionary<string, ScalesetState[]> _states = new ConcurrentDictionary<string, ScalesetState[]>();
-
-    /// set of states that indicate the scaleset can be updated
-    public static ScalesetState[] CanUpdate() {
-        return
-        _states.GetOrAdd(nameof(CanUpdate), k => new[]{
-            ScalesetState.Running,
-            ScalesetState.Resize
-        });
-    }
-
-    /// set of states that indicate work is needed during eventing
-    public static ScalesetState[] NeedsWork() {
-        return
-        _states.GetOrAdd(nameof(NeedsWork), k => new[]{
+    private static readonly IReadOnlySet<ScalesetState> _canUpdate = new HashSet<ScalesetState> { ScalesetState.Init, ScalesetState.Resize };
+    private static readonly IReadOnlySet<ScalesetState> _needsWork =
+        new HashSet<ScalesetState>{
             ScalesetState.Init,
             ScalesetState.Setup,
             ScalesetState.Resize,
             ScalesetState.Shutdown,
-            ScalesetState.Halt,
-        });
-    }
+            ScalesetState.Halt
+        };
+    private static readonly IReadOnlySet<ScalesetState> _available = new HashSet<ScalesetState> { ScalesetState.Resize, ScalesetState.Running };
+    private static readonly IReadOnlySet<ScalesetState> _resizing = new HashSet<ScalesetState> { ScalesetState.Halt, ScalesetState.Init, ScalesetState.Setup };
+
+    /// set of states that indicate the scaleset can be updated
+    public static IReadOnlySet<ScalesetState> CanUpdate => _canUpdate;
+
+    /// set of states that indicate work is needed during eventing
+    public static IReadOnlySet<ScalesetState> NeedsWork => _needsWork;
 
     /// set of states that indicate if it's available for work
-    public static ScalesetState[] Available() {
-        return
-        _states.GetOrAdd(nameof(Available), k => {
-            return
-                new[]{
-                ScalesetState.Resize,
-                ScalesetState.Running,
-            };
-        });
-    }
+    public static IReadOnlySet<ScalesetState> Available => _available;
 
     /// set of states that indicate scaleset is resizing
-    public static ScalesetState[] Resizing() {
-        return
-        _states.GetOrAdd(nameof(Resizing), k => {
-            return
-                new[]{
-                ScalesetState.Halt,
-                ScalesetState.Init,
-                ScalesetState.Setup,
-            };
-        });
-    }
+    public static IReadOnlySet<ScalesetState> Resizing => _resizing;
 }
 
 
 public static class VmStateHelper {
 
-    static ConcurrentDictionary<string, VmState[]> _states = new ConcurrentDictionary<string, VmState[]>();
-    public static VmState[] NeedsWork() {
-        return
-        _states.GetOrAdd(nameof(VmStateHelper.NeedsWork), k => {
-            return
-                new[]{
-                VmState.Init,
-                VmState.ExtensionsLaunch,
-                VmState.Stopping
-            };
-        });
-    }
+    private static readonly IReadOnlySet<VmState> _needsWork = new HashSet<VmState> { VmState.Init, VmState.Init, VmState.ExtensionsLaunch, VmState.Stopping };
+    private static readonly IReadOnlySet<VmState> _available = new HashSet<VmState> { VmState.Init, VmState.ExtensionsLaunch, VmState.ExtensionsFailed, VmState.VmAllocationFailed, VmState.Running, };
 
-    public static VmState[] Available() {
-        return
-        _states.GetOrAdd(nameof(VmStateHelper.Available), k => {
-            return
-                new[]{
-                VmState.Init,
-                VmState.ExtensionsLaunch,
-                VmState.ExtensionsFailed,
-                VmState.VmAllocationFailed,
-                VmState.Running,
-            };
-        });
-    }
+    public static IReadOnlySet<VmState> NeedsWork => _needsWork;
+    public static IReadOnlySet<VmState> Available => _available;
 }
 
 public static class TaskStateHelper {
-    static ConcurrentDictionary<string, TaskState[]> _states = new ConcurrentDictionary<string, TaskState[]>();
-    public static TaskState[] Available() {
-        return
-        _states.GetOrAdd(nameof(Available), k => {
-            return
-                 new[]{
-                    TaskState.Waiting,
-                    TaskState.Scheduled,
-                    TaskState.SettingUp,
-                    TaskState.Running,
-                    TaskState.WaitJob
-                 };
-        });
-    }
 
-    internal static TaskState[] NeedsWork() {
-        return
-        _states.GetOrAdd(nameof(TaskStateHelper.HasStarted), k =>
-            new[]{
-                TaskState.Init,
-                TaskState.Stopping
-            }
-        );
-    }
+    private static readonly IReadOnlySet<TaskState> _available = new HashSet<TaskState> { TaskState.Waiting, TaskState.Scheduled, TaskState.SettingUp, TaskState.Running, TaskState.WaitJob };
+    private static readonly IReadOnlySet<TaskState> _needsWork = new HashSet<TaskState> { TaskState.Init, TaskState.Stopping };
+    private static readonly IReadOnlySet<TaskState> _shuttingDown = new HashSet<TaskState> { TaskState.Stopping, TaskState.Stopped };
+    private static readonly IReadOnlySet<TaskState> _hasStarted = new HashSet<TaskState> { TaskState.Running, TaskState.Stopping, TaskState.Stopped };
 
+    public static IReadOnlySet<TaskState> Available => _available;
 
-    public static TaskState[] ShuttingDown() {
-        return
-        _states.GetOrAdd(nameof(TaskStateHelper.ShuttingDown), k =>
-            new[]{
-                TaskState.Stopping,
-                TaskState.Stopping,
-            }
-        );
-    }
+    public static IReadOnlySet<TaskState> NeedsWork => _needsWork;
 
-    internal static TaskState[] HasStarted() {
-        return
-        _states.GetOrAdd(nameof(TaskStateHelper.HasStarted), k =>
-            new[]{
-                TaskState.Running,
-                TaskState.Stopping,
-                TaskState.Stopped
-            }
-        );
-    }
+    public static IReadOnlySet<TaskState> ShuttingDown => _shuttingDown;
+
+    public static IReadOnlySet<TaskState> HasStarted => _hasStarted;
 
 }
 public enum PoolState {
@@ -273,26 +195,11 @@ public enum PoolState {
 }
 
 public static class PoolStateHelper {
-    static ConcurrentDictionary<string, PoolState[]> _states = new ConcurrentDictionary<string, PoolState[]>();
-    public static PoolState[] NeedsWork() {
-        return
-        _states.GetOrAdd("NeedsWork", k =>
-            new[]{
-                PoolState.Init,
-                PoolState.Shutdown,
-                PoolState.Halt
-            }
-        );
-    }
+    private static readonly IReadOnlySet<PoolState> _needsWork = new HashSet<PoolState> { PoolState.Init, PoolState.Shutdown, PoolState.Halt };
+    private static readonly IReadOnlySet<PoolState> _available = new HashSet<PoolState> { PoolState.Running };
 
-    public static PoolState[] Available() {
-        return
-        _states.GetOrAdd("Available", k =>
-                new[]{
-                PoolState.Running
-                }
-        );
-    }
+    public static IReadOnlySet<PoolState> NeedsWork => _needsWork;
+    public static IReadOnlySet<PoolState> Available => _available;
 }
 
 public enum Architecture {
@@ -345,8 +252,6 @@ public enum ContainerPermission {
 }
 
 
-
-
 public enum Compare {
     Equal,
     AtLeast,

src/ApiService/ApiService/QueueTaskHearbeat.cs

@@ -18,8 +18,8 @@ public class QueueTaskHearbeat {
         _events = events;
     }
 
-    //[Function("QueueTaskHearbeat")]
-    public async Async.Task Run([QueueTrigger("myqueue-items2", Connection = "AzureWebJobsStorage")] string msg) {
+    [Function("QueueTaskHearbeat")]
+    public async Async.Task Run([QueueTrigger("task-heartbeat", Connection = "AzureWebJobsStorage")] string msg) {
         _logger.LogInformation($"heartbeat: {msg}");
 
         var hb = JsonSerializer.Deserialize<TaskHeartbeatEntry>(msg, EntityConverter.GetJsonSerializerOptions()).EnsureNotNull($"wrong data {msg}");

src/ApiService/ApiService/TimerProxy.cs

@@ -33,7 +33,7 @@ public partial class TimerProxy {
         var proxies = await _proxYOperations.QueryAsync().ToListAsync();
 
         foreach (var proxy in proxies) {
-            if (VmStateHelper.Available().Contains(proxy.State)) {
+            if (VmStateHelper.Available.Contains(proxy.State)) {
                 // Note, outdated checked at the start, but set at the end of this loop.
                 // As this function is called via a timer, this works around a user
                 // requesting to use the proxy while this function is checking if it's
@@ -49,7 +49,7 @@ public partial class TimerProxy {
                 }
             }
 
-            if (VmStateHelper.NeedsWork().Contains(proxy.State)) {
+            if (VmStateHelper.NeedsWork.Contains(proxy.State)) {
                 _logger.Error($"scaleset-proxy: update state. proxy:{proxy.Region} state:{proxy.State}");
                 await _proxYOperations.ProcessStateUpdate(proxy);
             }

src/ApiService/ApiService/TimerRepro.cs

@@ -25,7 +25,7 @@ public class TimerRepro {
 
         var expiredVmIds = expired.Select(repro => repro?.VmId);
 
-        await foreach (var repro in _reproOperations.SearchStates(VmStateHelper.NeedsWork())) {
+        await foreach (var repro in _reproOperations.SearchStates(VmStateHelper.NeedsWork)) {
             if (await expiredVmIds.ContainsAsync(repro.VmId)) {
                 // this VM already got processed during the expired phase
                 continue;

src/ApiService/ApiService/TimerTasks.cs

@@ -44,7 +44,7 @@ public class TimerTasks {
             await _jobOperations.ProcessStateUpdates(job);
         }
 
-        var tasks = _taskOperations.SearchStates(states: TaskStateHelper.NeedsWork());
+        var tasks = _taskOperations.SearchStates(states: TaskStateHelper.NeedsWork);
         await foreach (var task in tasks) {
             _logger.Info($"update task: {task.TaskId}");
             await _taskOperations.ProcessStateUpdate(task);

src/ApiService/ApiService/onefuzzlib/NodeOperations.cs

@@ -134,7 +134,7 @@ public class NodeOperations : StatefulOrm<Node, NodeState>, INodeOperations {
             }
             var scaleset = scalesetResult.OkV!;
 
-            if (!ScalesetStateHelper.Available().Contains(scaleset.State)) {
+            if (!ScalesetStateHelper.Available.Contains(scaleset.State)) {
                 _logTracer.Info($"can_process_new_work scaleset not available for work. scaleset_id:{node.ScalesetId} machine_id:{node.MachineId}");
                 return false;
             }
@@ -147,7 +147,7 @@ public class NodeOperations : StatefulOrm<Node, NodeState>, INodeOperations {
         }
 
         var pool = poolResult.OkV!;
-        if (!PoolStateHelper.Available().Contains(pool.State)) {
+        if (!PoolStateHelper.Available.Contains(pool.State)) {
             _logTracer.Info($"can_schedule - pool is not available for work. pool_name:{node.PoolName} machine_id:{node.MachineId}");
             return false;
         }

src/ApiService/ApiService/onefuzzlib/NotificationOperations.cs

@@ -78,7 +78,7 @@ public class NotificationOperations : Orm<Notification>, INotificationOperations
 
     public IAsyncEnumerable<(Task, IEnumerable<string>)> GetQueueTasks() {
         // Nullability mismatch: We filter tuples where the containers are null
-        return _context.TaskOperations.SearchStates(states: TaskStateHelper.Available())
+        return _context.TaskOperations.SearchStates(states: TaskStateHelper.Available)
             .Select(task => (task, _context.TaskOperations.GetInputContainerQueues(task.Config)))
             .Where(taskTuple => taskTuple.Item2 != null)!;
     }

src/ApiService/ApiService/onefuzzlib/ProxyOperations.cs

@@ -40,7 +40,7 @@ public class ProxyOperations : StatefulOrm<Proxy, VmState>, IProxyOperations {
                 continue;
             }
 
-            if (!VmStateHelper.Available().Contains(proxy.State)) {
+            if (!VmStateHelper.Available.Contains(proxy.State)) {
                 continue;
             }
             return proxy;
@@ -71,7 +71,7 @@ public class ProxyOperations : StatefulOrm<Proxy, VmState>, IProxyOperations {
     }
 
     public bool IsOutdated(Proxy proxy) {
-        if (!VmStateHelper.Available().Contains(proxy.State)) {
+        if (!VmStateHelper.Available.Contains(proxy.State)) {
             return false;
         }
 

src/ApiService/ApiService/onefuzzlib/Scheduler.cs

@@ -58,7 +58,7 @@ public class Scheduler : IScheduler {
     }
 
     private async Async.Task<bool> ScheduleWorkset(WorkSet workSet, Pool pool, int count) {
-        if (!PoolStateHelper.Available().Contains(pool.State)) {
+        if (!PoolStateHelper.Available.Contains(pool.State)) {
             _logTracer.Info($"pool not available for work: {pool.Name} state: {pool.State}");
             return false;
         }

src/ApiService/ApiService/onefuzzlib/TaskOperations.cs

@@ -71,19 +71,19 @@ public class TaskOperations : StatefulOrm<Task, TaskState>, ITaskOperations {
     }
 
     public async Async.Task MarkStopping(Task task) {
-        if (TaskStateHelper.ShuttingDown().Contains(task.State)) {
+        if (TaskStateHelper.ShuttingDown.Contains(task.State)) {
             _logTracer.Verbose($"ignoring post - task stop calls to stop {task.JobId}:{task.TaskId}");
             return;
         }
 
-        if (TaskStateHelper.HasStarted().Contains(task.State)) {
+        if (TaskStateHelper.HasStarted.Contains(task.State)) {
             await MarkFailed(task, new Error(Code: ErrorCode.TASK_FAILED, Errors: new[] { "task never started" }));
 
         }
     }
 
     public async Async.Task MarkFailed(Task task, Error error, List<Task>? taskInJob = null) {
-        if (TaskStateHelper.ShuttingDown().Contains(task.State)) {
+        if (TaskStateHelper.ShuttingDown.Contains(task.State)) {
             _logTracer.Verbose(
                 $"ignoring post-task stop failures for {task.JobId}:{task.TaskId}"
             );
@@ -210,7 +210,7 @@ public class TaskOperations : StatefulOrm<Task, TaskState>, ITaskOperations {
                     return false;
                 }
 
-                if (!TaskStateHelper.HasStarted().Contains(t.State)) {
+                if (!TaskStateHelper.HasStarted.Contains(t.State)) {
                     return false;
                 }
             }