Migrating QueueTaskHeartbeat (#1777)

* Migrating QueueTaskHeartbeat

* changing the name of the input queue

* rename type alias Tasks to Async

* Fix property casing

* fixing types

* Removing IStorageProvider

* fix function name

* address PR comments

src/ApiService/ApiService/OneFuzzTypes/Enums.cs

@@ -44,3 +44,66 @@ public enum WebhookMessageState
     Succeeded,
     Failed
 }
+
+public enum TaskState
+{
+    Init,
+    Waiting,
+    Scheduled,
+    Setting_up,
+    Running,
+    Stopping,
+    Stopped,
+    WaitJob
+}
+
+public enum TaskType
+{
+    Coverage,
+    LibfuzzerFuzz,
+    LibfuzzerCoverage,
+    LibfuzzerCrashReport,
+    LibfuzzerMerge,
+    LibfuzzerRegression,
+    GenericAnalysis,
+    GenericSupervisor,
+    GenericMerge,
+    GenericGenerator,
+    GenericCrashReport,
+    GenericRegression
+}
+
+public enum Os
+{
+    Windows,
+    Linux
+}
+
+public enum ContainerType
+{
+    Analysis,
+    Coverage,
+    Crashes,
+    Inputs,
+    NoRepro,
+    ReadonlyInputs,
+    Reports,
+    Setup,
+    Tools,
+    UniqueInputs,
+    UniqueReports,
+    RegressionReports,
+    Logs
+}
+
+
+public enum StatsFormat
+{
+    AFL
+}
+
+public enum TaskDebugFlag
+{
+    KeepNodeOnFailure,
+    KeepNodeOnCompletion,
+}

src/ApiService/ApiService/OneFuzzTypes/Events.cs

@@ -104,11 +104,11 @@ namespace Microsoft.OneFuzz.Service
     //    ) : BaseEvent();
 
 
-    //record EventTaskHeartbeat(
-    //    JobId: Guid,
-    //    TaskId: Guid,
-    //    Config: TaskConfig
-    //): BaseEvent();
+    record EventTaskHeartbeat(
+       Guid JobId,
+       Guid TaskId,
+       TaskConfig Config
+    ) : BaseEvent();
 
 
     //record EventPing(

src/ApiService/ApiService/OneFuzzTypes/Model.cs

@@ -2,6 +2,8 @@ using Microsoft.OneFuzz.Service.OneFuzzLib.Orm;
 using System;
 using System.Collections.Generic;
 using PoolName = System.String;
+using Region = System.String;
+using Container = System.String;
 
 namespace Microsoft.OneFuzz.Service;
 
@@ -29,9 +31,15 @@ public enum HeartbeatType
     TaskAlive,
 }
 
-public record HeartbeatData(HeartbeatType type);
+public record HeartbeatData(HeartbeatType Type);
 
-public record NodeHeartbeatEntry(Guid NodeId, HeartbeatData[] data);
+public record TaskHeartbeatEntry(
+    Guid TaskId,
+    Guid? JobId,
+    Guid MachineId,
+    HeartbeatData[] Data
+    );
+public record NodeHeartbeatEntry(Guid NodeId, HeartbeatData[] Data);
 
 public record NodeCommandStopIfFree();
 
@@ -79,7 +87,7 @@ public enum NodeState
 
 public record ProxyHeartbeat
 (
-    string Region,
+    Region Region,
     Guid ProxyId,
     List<ProxyForward> Forwards,
     DateTimeOffset TimeStamp
@@ -103,7 +111,7 @@ public partial record Node
 
 public partial record ProxyForward
 (
-    [PartitionKey] string Region,
+    [PartitionKey] Region Region,
     [RowKey] int DstPort,
     int SrcPort,
     string DstIp
@@ -113,7 +121,7 @@ public partial record ProxyConfig
 (
     Uri Url,
     string Notification,
-    string Region,
+    Region Region,
     Guid? ProxyId,
     List<ProxyForward> Forwards,
     string InstanceTelemetryKey,
@@ -123,7 +131,7 @@ public partial record ProxyConfig
 
 public partial record Proxy
 (
-    [PartitionKey] string Region,
+    [PartitionKey] Region Region,
     [RowKey] Guid ProxyId,
     DateTimeOffset? CreatedTimestamp,
     VmState State,
@@ -148,23 +156,102 @@ public record EventMessage(
 ) : EntityBase();
 
 
-//record AnyHttpUrl(AnyUrl):
-//    allowed_schemes = {'http', 'https
-//
+public record TaskDetails(
+
+    TaskType Type,
+    int Duration,
+    string? TargetExe,
+    Dictionary<string, string>? TargetEnv,
+    List<string>? TargetOptions,
+    int? TargetWorkers,
+    bool? TargetOptionsMerge,
+    bool? CheckAsanLog,
+    bool? CheckDebugger,
+    int? CheckRetryCount,
+    bool? CheckFuzzerHelp,
+    bool? ExpectCrashOnFailure,
+    bool? RenameOutput,
+    string? SupervisorExe,
+    Dictionary<string, string>? SupervisorEnv,
+    List<string>? SupervisorOptions,
+    string? SupervisorInputMarker,
+    string? GeneratorExe,
+    Dictionary<string, string>? GeneratorEnv,
+    List<string>? GeneratorOptions,
+    string? AnalyzerExe,
+    Dictionary<string, string>? AnalyzerEnv,
+    List<string> AnalyzerOptions,
+    ContainerType? WaitForFiles,
+    string? StatsFile,
+    StatsFormat? StatsFormat,
+    bool? RebootAfterSetup,
+    int? TargetTimeout,
+    int? EnsembleSyncDelay,
+    bool? PreserveExistingOutputs,
+    List<string>? ReportList,
+    int? MinimizedStackDepth,
+    string? CoverageFilter
+);
+
+public record TaskVm(
+    Region Region,
+    string Sku,
+    string Image,
+    int Count,
+    bool SpotInstance,
+    bool? RebootAfterSetup
+);
+
+public record TaskPool(
+    int Count,
+    PoolName PoolName
+);
+
+public record TaskContainers(
+    ContainerType Type,
+    Container Name
+);
+public record TaskConfig(
+   Guid JobId,
+   List<Guid>? PrereqTasks,
+   TaskDetails Task,
+   TaskVm? Vm,
+   TaskPool? Pool,
+   List<TaskContainers>? Containers,
+   Dictionary<string, string>? Tags,
+   List<TaskDebugFlag>? Debug,
+   bool? Colocate
+   );
 
 
+public record TaskEventSummary(
+    DateTimeOffset? Timestamp,
+    string EventData,
+    string EventType
+    );
 
 
+public record NodeAssignment(
+    Guid NodeId,
+    Guid? ScalesetId,
+    NodeTaskState State
+    );
 
-//public record TaskConfig(
-//    Guid jobId,
-//    List<Guid> PrereqTasks,
-//    TaskDetails Task,
-//    TaskVm? vm,
-//    TaskPool pool: Optional[]
-//    containers: List[TaskContainers]
-//    tags: Dict[str, str]
-//    debug: Optional[List[TaskDebugFlag]]
-//    colocate: Optional[bool]
-//    ): EntityBase();
 
+public record Task(
+    // Timestamp: Optional[datetime] = Field(alias="Timestamp")
+    [PartitionKey] Guid JobId,
+    [RowKey] Guid TaskId,
+    TaskState State,
+    Os Os,
+    TaskConfig Config,
+    Error? Error,
+    Authentication? Auth,
+    DateTimeOffset? Heartbeat,
+    DateTimeOffset? EndTime,
+    UserInfo? UserInfo) : EntityBase()
+{
+    List<TaskEventSummary> Events { get; set; } = new List<TaskEventSummary>();
+    List<NodeAssignment> Nodes { get; set; } = new List<NodeAssignment>();
+
+}

src/ApiService/ApiService/Program.cs

@@ -1,10 +1,14 @@
+// to avoid collision with Task in model.cs
+global using Async = System.Threading.Tasks;
+
 using System;
 using System.Collections.Generic;
 using Microsoft.Extensions.Hosting;
 using Microsoft.Extensions.DependencyInjection;
-using Microsoft.OneFuzz.Service.OneFuzzLib.Orm;
 using ApiService.OneFuzzLib;
 
+
+
 namespace Microsoft.OneFuzz.Service;
 
 public class Program
@@ -34,11 +38,11 @@ public class Program
         .ConfigureServices((context, services) =>
             services
             .AddSingleton<ILogTracerFactory>(_ => new LogTracerFactory(GetLoggers()))
-            .AddSingleton<IStorageProvider>(_ => new StorageProvider(EnvironmentVariables.OneFuzz.FuncStorage ?? throw new InvalidOperationException("Missing account id")))
             .AddSingleton<INodeOperations, NodeOperations>()
             .AddSingleton<IEvents, Events>()
             .AddSingleton<IWebhookOperations, WebhookOperations>()
             .AddSingleton<IWebhookMessageLogOperations, WebhookMessageLogOperations>()
+            .AddSingleton<ITaskOperations, TaskOperations>()
             .AddSingleton<IQueue, Queue>()
             .AddSingleton<ICreds>(_ => new Creds())
             .AddSingleton<IStorage, Storage>()

src/ApiService/ApiService/QueueFileChanges.cs

@@ -2,7 +2,6 @@ using System;
 using Microsoft.Azure.Functions.Worker;
 using System.Collections.Generic;
 using System.Text.Json;
-using System.Threading.Tasks;
 using Microsoft.OneFuzz.Service.OneFuzzLib.Orm;
 using System.Linq;
 
@@ -15,19 +14,17 @@ public class QueueFileChanges
     const int MAX_DEQUEUE_COUNT = 5;
 
     private readonly ILogTracerFactory _loggerFactory;
-    private readonly IStorageProvider _storageProvider;
 
     private readonly IStorage _storage;
 
-    public QueueFileChanges(ILogTracerFactory loggerFactory, IStorageProvider storageProvider, IStorage storage)
+    public QueueFileChanges(ILogTracerFactory loggerFactory, IStorage storage)
     {
         _loggerFactory = loggerFactory;
-        _storageProvider = storageProvider;
         _storage = storage;
     }
 
     [Function("QueueFileChanges")]
-    public Task Run(
+    public Async.Task Run(
         [QueueTrigger("file-changes-refactored", Connection = "AzureWebJobsStorage")] string msg,
         int dequeueCount)
     {
@@ -42,18 +39,18 @@ public class QueueFileChanges
         if (!fileChangeEvent.ContainsKey(eventType)
             || fileChangeEvent[eventType] != "Microsoft.Storage.BlobCreated")
         {
-            return Task.CompletedTask;
+            return Async.Task.CompletedTask;
         }
 
         const string topic = "topic";
         if (!fileChangeEvent.ContainsKey(topic)
             || !_storage.CorpusAccounts(log).Contains(fileChangeEvent[topic]))
         {
-            return Task.CompletedTask;
+            return Async.Task.CompletedTask;
         }
 
         file_added(log, fileChangeEvent, lastTry);
-        return Task.CompletedTask;
+        return Async.Task.CompletedTask;
     }
 
     private void file_added(ILogTracer log, Dictionary<string, string> fileChangeEvent, bool failTaskOnTransientError)

src/ApiService/ApiService/QueueNodeHearbeat.cs

@@ -1,7 +1,6 @@
 using System;
 using Microsoft.Azure.Functions.Worker;
 using System.Text.Json;
-using System.Threading.Tasks;
 using Microsoft.OneFuzz.Service.OneFuzzLib.Orm;
 
 namespace Microsoft.OneFuzz.Service;
@@ -22,7 +21,7 @@ public class QueueNodeHearbeat
     }
 
     [Function("QueueNodeHearbeat")]
-    public async Task Run([QueueTrigger("myqueue-items", Connection = "AzureWebJobsStorage")] string msg)
+    public async Async.Task Run([QueueTrigger("myqueue-items", Connection = "AzureWebJobsStorage")] string msg)
     {
         var log = _loggerFactory.MakeLogTracer(Guid.NewGuid());
         log.Info($"heartbeat: {msg}");

src/ApiService/ApiService/QueueProxyHeartbeat.cs

@@ -1,7 +1,6 @@
 using System;
 using Microsoft.Azure.Functions.Worker;
 using System.Text.Json;
-using System.Threading.Tasks;
 using Microsoft.OneFuzz.Service.OneFuzzLib.Orm;
 
 namespace Microsoft.OneFuzz.Service;
@@ -19,7 +18,7 @@ public class QueueProxyHearbeat
     }
 
     [Function("QueueProxyHearbeat")]
-    public async Task Run([QueueTrigger("myqueue-items", Connection = "AzureWebJobsStorage")] string msg)
+    public async Async.Task Run([QueueTrigger("myqueue-items", Connection = "AzureWebJobsStorage")] string msg)
     {
         var log = _loggerFactory.MakeLogTracer(Guid.NewGuid());
 

src/ApiService/ApiService/QueueTaskHearbeat.cs

@@ -0,0 +1,43 @@
+using System;
+using Microsoft.Azure.Functions.Worker;
+using Microsoft.Extensions.Logging;
+using System.Text.Json;
+using Microsoft.OneFuzz.Service.OneFuzzLib.Orm;
+
+namespace Microsoft.OneFuzz.Service;
+
+
+public class QueueTaskHearbeat
+{
+    private readonly ILogger _logger;
+
+    private readonly IEvents _events;
+    private readonly ITaskOperations _tasks;
+
+    public QueueTaskHearbeat(ILoggerFactory loggerFactory, ITaskOperations tasks, IEvents events)
+    {
+        _logger = loggerFactory.CreateLogger<QueueTaskHearbeat>();
+        _tasks = tasks;
+        _events = events;
+    }
+
+    [Function("QueueTaskHearbeat")]
+    public async Async.Task Run([QueueTrigger("myqueue-items2", Connection = "AzureWebJobsStorage")] string msg)
+    {
+        _logger.LogInformation($"heartbeat: {msg}");
+
+        var hb = JsonSerializer.Deserialize<TaskHeartbeatEntry>(msg, EntityConverter.GetJsonSerializerOptions()).EnsureNotNull($"wrong data {msg}");
+
+        var task = await _tasks.GetByTaskId(hb.TaskId);
+
+        if (task == null)
+        {
+            _logger.LogWarning($"invalid task id: {hb.TaskId}");
+            return;
+        }
+
+        var newTask = task with { Heartbeat = DateTimeOffset.UtcNow };
+        await _tasks.Replace(newTask);
+        await _events.SendEvent(new EventTaskHeartbeat(newTask.JobId, newTask.TaskId, newTask.Config));
+    }
+}

src/ApiService/ApiService/UserCredentials.cs

@@ -5,6 +5,7 @@ using System.Threading.Tasks;
 using Microsoft.Azure.Functions.Worker.Http;
 using Microsoft.IdentityModel.Tokens;
 
+
 namespace Microsoft.OneFuzz.Service;
 
 public class UserCredentials
@@ -53,7 +54,7 @@ public class UserCredentials
 
     static Task<OneFuzzResult<string[]>> GetAllowedTenants()
     {
-        return Task.FromResult(OneFuzzResult<string[]>.Ok(Array.Empty<string>()));
+        return Async.Task.FromResult(OneFuzzResult<string[]>.Ok(Array.Empty<string>()));
     }
 
     /*

src/ApiService/ApiService/onefuzzlib/Events.cs

@@ -5,7 +5,6 @@ using System.Collections.Generic;
 using System.Text;
 using System.Text.Json;
 using System.Text.Json.Serialization;
-using System.Threading.Tasks;
 
 namespace Microsoft.OneFuzz.Service
 {
@@ -20,9 +19,9 @@ namespace Microsoft.OneFuzz.Service
 
     public interface IEvents
     {
-        public Task SendEvent(BaseEvent anEvent);
+        public Async.Task SendEvent(BaseEvent anEvent);
 
-        public Task QueueSignalrEvent(EventMessage message);
+        public Async.Task QueueSignalrEvent(EventMessage message);
     }
 
     public class Events : IEvents
@@ -38,14 +37,14 @@ namespace Microsoft.OneFuzz.Service
             _webhook = webhook;
         }
 
-        public async Task QueueSignalrEvent(EventMessage eventMessage)
+        public async Async.Task QueueSignalrEvent(EventMessage eventMessage)
         {
             var message = new SignalREvent("events", new List<EventMessage>() { eventMessage });
             var encodedMessage = Encoding.UTF8.GetBytes(JsonSerializer.Serialize(message));
             await _queue.SendMessage("signalr-events", encodedMessage, StorageType.Config);
         }
 
-        public async Task SendEvent(BaseEvent anEvent)
+        public async Async.Task SendEvent(BaseEvent anEvent)
         {
             var log = _loggerFactory.MakeLogTracer(Guid.NewGuid());
             var eventType = anEvent.GetEventType();

src/ApiService/ApiService/onefuzzlib/Queue.cs

@@ -8,8 +8,8 @@ using System.Threading.Tasks;
 namespace Microsoft.OneFuzz.Service;
 public interface IQueue
 {
-    Task SendMessage(string name, byte[] message, StorageType storageType, TimeSpan? visibilityTimeout = null, TimeSpan? timeToLive = null);
-    Task<bool> QueueObject<T>(string name, T obj, StorageType storageType, TimeSpan? visibilityTimeout);
+    Async.Task SendMessage(string name, byte[] message, StorageType storageType, TimeSpan? visibilityTimeout = null, TimeSpan? timeToLive = null);
+    Async.Task<bool> QueueObject<T>(string name, T obj, StorageType storageType, TimeSpan? visibilityTimeout);
 }
 
 
@@ -25,7 +25,7 @@ public class Queue : IQueue
     }
 
 
-    public async Task SendMessage(string name, byte[] message, StorageType storageType, TimeSpan? visibilityTimeout = null, TimeSpan? timeToLive = null)
+    public async Async.Task SendMessage(string name, byte[] message, StorageType storageType, TimeSpan? visibilityTimeout = null, TimeSpan? timeToLive = null)
     {
         var queue = GetQueue(name, storageType);
         if (queue != null)

src/ApiService/ApiService/onefuzzlib/TaskOperations.cs

@@ -0,0 +1,28 @@
+using ApiService.OneFuzzLib.Orm;
+using System;
+using System.Linq;
+
+namespace Microsoft.OneFuzz.Service;
+
+public interface ITaskOperations : IOrm<Task>
+{
+    Async.Task<Task?> GetByTaskId(Guid taskId);
+}
+
+public class TaskOperations : Orm<Task>, ITaskOperations
+{
+
+    public TaskOperations(IStorage storage)
+        : base(storage)
+    {
+
+    }
+
+    public async Async.Task<Task?> GetByTaskId(Guid taskId)
+    {
+        var data = QueryAsync(filter: $"RowKey eq '{taskId}'");
+
+        return await data.FirstOrDefaultAsync();
+    }
+
+}

src/ApiService/ApiService/onefuzzlib/WebhookOperations.cs

@@ -2,7 +2,6 @@
 using Microsoft.OneFuzz.Service;
 using System;
 using System.Collections.Generic;
-using System.Threading.Tasks;
 
 namespace ApiService.OneFuzzLib;
 
@@ -29,7 +28,7 @@ public class WebhookMessageLogOperations : Orm<WebhookMessageLog>, IWebhookMessa
     }
 
 
-    public async Task QueueWebhook(WebhookMessageLog webhookLog)
+    public async Async.Task QueueWebhook(WebhookMessageLog webhookLog)
     {
         var log = _loggerFactory.MakeLogTracer(Guid.NewGuid());
         var obj = new WebhookMessageQueueObj(webhookLog.WebhookId, webhookLog.EventId);
@@ -65,7 +64,7 @@ public class WebhookMessageLogOperations : Orm<WebhookMessageLog>, IWebhookMessa
 
 public interface IWebhookOperations
 {
-    Task SendEvent(EventMessage eventMessage);
+    Async.Task SendEvent(EventMessage eventMessage);
 }
 
 public class WebhookOperations : Orm<Webhook>, IWebhookOperations
@@ -77,7 +76,7 @@ public class WebhookOperations : Orm<Webhook>, IWebhookOperations
         _webhookMessageLogOperations = webhookMessageLogOperations;
     }
 
-    async public Task SendEvent(EventMessage eventMessage)
+    async public Async.Task SendEvent(EventMessage eventMessage)
     {
         await foreach (var webhook in GetWebhooksCached())
         {
@@ -89,7 +88,7 @@ public class WebhookOperations : Orm<Webhook>, IWebhookOperations
         }
     }
 
-    async private Task AddEvent(Webhook webhook, EventMessage eventMessage)
+    async private Async.Task AddEvent(Webhook webhook, EventMessage eventMessage)
     {
         var message = new WebhookMessageLog(
              EventId: eventMessage.EventId,

src/ApiService/ApiService/onefuzzlib/orm/StorageProvider.cs

@@ -1,53 +0,0 @@
-using Azure.Data.Tables;
-using System;
-using System.Linq;
-using System.Threading.Tasks;
-using Azure.Core;
-using Azure.ResourceManager.Storage;
-using Azure.ResourceManager;
-using Azure.Identity;
-
-namespace Microsoft.OneFuzz.Service.OneFuzzLib.Orm;
-
-
-public interface IStorageProvider
-{
-    Task<TableClient> GetTableClient(string table);
-    //IAsyncEnumerable<T> QueryAsync<T>(string filter) where T : EntityBase;
-    //Task<bool> Replace<T>(T entity) where T : EntityBase;
-
-}
-
-public class StorageProvider : IStorageProvider
-{
-    private readonly string _accountId;
-    private readonly EntityConverter _entityConverter;
-    private readonly ArmClient _armClient;
-
-    public StorageProvider(string accountId)
-    {
-        _accountId = accountId;
-        _entityConverter = new EntityConverter();
-        _armClient = new ArmClient(new DefaultAzureCredential());
-    }
-
-    public async Task<TableClient> GetTableClient(string table)
-    {
-        var (name, key) = GetStorageAccountNameAndKey(_accountId);
-        var identifier = new ResourceIdentifier(_accountId);
-        var tableClient = new TableServiceClient(new Uri($"https://{identifier.Name}.table.core.windows.net"), new TableSharedKeyCredential(name, key));
-        await tableClient.CreateTableIfNotExistsAsync(table);
-        return tableClient.GetTableClient(table);
-    }
-
-
-    public (string?, string?) GetStorageAccountNameAndKey(string accountId)
-    {
-        var resourceId = new ResourceIdentifier(accountId);
-        var storageAccount = _armClient.GetStorageAccountResource(resourceId);
-        var key = storageAccount.GetKeys().Value.Keys.FirstOrDefault();
-        return (resourceId.Name, key?.Value);
-    }
-
-
-}