Make min_available_memory_mb configurable (microsoft#3577)

kananb · nharper285 · web-flow · commit d9ac2a7da934 · 2023-10-18T15:09:55.000-07:00
* Add min_available_memory_mb to cli and service configs to be passed to agent

* Add min_available_memory_mb to task create and template parameters

* Add logging to indicate the configured min_available_memory

* Add missing parameter

* Remove unnecessary import

* Update webhook_events.py

* test change

* original version of docs.

* Update webhook_events.py again

* Don't start memory checking task when min_available_memory_mb is zero

* cargo fmt

* Remove comments

* Add min_available_memory_mb to libfuzzer's _create_tasks()

* Add a megabyte converted value for min available memory logs

* Use float instead of int for division

---------

Co-authored-by: Noah Harper &lt;noharper@microsoft.com&gt;
diff --git a/docs/webhook_events.md b/docs/webhook_events.md
@@ -445,6 +445,11 @@ If webhook is set to have Event Grid message format then the payload will look a
                     "title": "Generator Options",
                     "type": "array"
                 },
+                "min_available_memory_mb": {
+                    "minimum": 0,
+                    "title": "Min Available Memory Mb",
+                    "type": "integer"
+                },
                 "minimized_stack_depth": {
                     "title": "Minimized Stack Depth",
                     "type": "integer"
@@ -2332,6 +2337,11 @@ If webhook is set to have Event Grid message format then the payload will look a
                     "title": "Generator Options",
                     "type": "array"
                 },
+                "min_available_memory_mb": {
+                    "minimum": 0,
+                    "title": "Min Available Memory Mb",
+                    "type": "integer"
+                },
                 "minimized_stack_depth": {
                     "title": "Minimized Stack Depth",
                     "type": "integer"
@@ -3054,6 +3064,11 @@ If webhook is set to have Event Grid message format then the payload will look a
                     "title": "Generator Options",
                     "type": "array"
                 },
+                "min_available_memory_mb": {
+                    "minimum": 0,
+                    "title": "Min Available Memory Mb",
+                    "type": "integer"
+                },
                 "minimized_stack_depth": {
                     "title": "Minimized Stack Depth",
                     "type": "integer"
@@ -3572,6 +3587,11 @@ If webhook is set to have Event Grid message format then the payload will look a
                     "title": "Generator Options",
                     "type": "array"
                 },
+                "min_available_memory_mb": {
+                    "minimum": 0,
+                    "title": "Min Available Memory Mb",
+                    "type": "integer"
+                },
                 "minimized_stack_depth": {
                     "title": "Minimized Stack Depth",
                     "type": "integer"
@@ -4056,6 +4076,11 @@ If webhook is set to have Event Grid message format then the payload will look a
                     "title": "Generator Options",
                     "type": "array"
                 },
+                "min_available_memory_mb": {
+                    "minimum": 0,
+                    "title": "Min Available Memory Mb",
+                    "type": "integer"
+                },
                 "minimized_stack_depth": {
                     "title": "Minimized Stack Depth",
                     "type": "integer"
@@ -4514,6 +4539,11 @@ If webhook is set to have Event Grid message format then the payload will look a
                     "title": "Generator Options",
                     "type": "array"
                 },
+                "min_available_memory_mb": {
+                    "minimum": 0,
+                    "title": "Min Available Memory Mb",
+                    "type": "integer"
+                },
                 "minimized_stack_depth": {
                     "title": "Minimized Stack Depth",
                     "type": "integer"
@@ -4999,6 +5029,11 @@ If webhook is set to have Event Grid message format then the payload will look a
                     "title": "Generator Options",
                     "type": "array"
                 },
+                "min_available_memory_mb": {
+                    "minimum": 0,
+                    "title": "Min Available Memory Mb",
+                    "type": "integer"
+                },
                 "minimized_stack_depth": {
                     "title": "Minimized Stack Depth",
                     "type": "integer"
@@ -6743,6 +6778,11 @@ If webhook is set to have Event Grid message format then the payload will look a
                     "title": "Generator Options",
                     "type": "array"
                 },
+                "min_available_memory_mb": {
+                    "minimum": 0,
+                    "title": "Min Available Memory Mb",
+                    "type": "integer"
+                },
                 "minimized_stack_depth": {
                     "title": "Minimized Stack Depth",
                     "type": "integer"
diff --git a/src/ApiService/ApiService/OneFuzzTypes/Model.cs b/src/ApiService/ApiService/OneFuzzTypes/Model.cs
@@ -239,6 +239,7 @@ public record TaskDetails(
     List<string>? ReportList = null,
     long? MinimizedStackDepth = null,
     Dictionary<string, string>? TaskEnv = null,
+    ulong? MinAvailableMemoryMb = null,
 
     // Deprecated. Retained for processing old table data.
     string? CoverageFilter = null,
@@ -1164,6 +1165,7 @@ Dictionary<string, string> Tags
     public IContainerDef? RegressionReports { get; set; }
     public IContainerDef? ExtraSetup { get; set; }
     public IContainerDef? ExtraOutput { get; set; }
+    public ulong? MinAvailableMemoryMb { get; set; }
 }
 
 public record NodeCommandEnvelope(
diff --git a/src/ApiService/ApiService/onefuzzlib/Config.cs b/src/ApiService/ApiService/onefuzzlib/Config.cs
@@ -73,7 +73,10 @@ private static BlobContainerSasPermissions ConvertPermissions(ContainerPermissio
             HeartbeatQueue: await _queue.GetQueueSas("task-heartbeat", StorageType.Config, QueueSasPermissions.Add) ?? throw new Exception("unable to get heartbeat queue sas"),
             JobResultQueue: await _queue.GetQueueSas("job-result", StorageType.Config, QueueSasPermissions.Add) ?? throw new Exception("unable to get heartbeat queue sas"),
             Tags: task.Config.Tags ?? new Dictionary<string, string>()
-        );
+        ) {
+            // It's okay if this is null because the agent will use a default value if so.
+            MinAvailableMemoryMb = task.Config.Task.MinAvailableMemoryMb,
+        };
 
         if (definition.MonitorQueue != null) {
             config.inputQueue = await _queue.GetQueueSas(task.TaskId.ToString(), StorageType.Corpus, QueueSasPermissions.All);
diff --git a/src/ApiService/ApiService/onefuzzlib/notifications/JinjaTemplateAdapter.cs b/src/ApiService/ApiService/onefuzzlib/notifications/JinjaTemplateAdapter.cs
@@ -166,6 +166,7 @@ public static async Async.Task<TemplateValidationResponse> ValidateScribanTempla
                         targetOptions,
                         1,
                         new Dictionary<string, string>(),
+                        null,
                         "coverage filter",
                         "module allow list",
                         "source allow list",
diff --git a/src/agent/onefuzz-task/src/managed/cmd.rs b/src/agent/onefuzz-task/src/managed/cmd.rs
@@ -2,7 +2,7 @@
 // Licensed under the MIT License.
 use std::path::PathBuf;
 
-use anyhow::{bail, Result};
+use anyhow::Result;
 use clap::{value_parser, Arg, Command};
 use ipc_channel::ipc::{self, IpcReceiver, IpcSender};
 
@@ -96,21 +96,29 @@ pub async fn run(args: &clap::ArgMatches) -> Result<()> {
 
     let min_available_memory_bytes = 1_000_000 * config.common().min_available_memory_mb;
 
-    // If the memory limit is 0, this will resolve immediately with an error.
-    let check_oom = out_of_memory(min_available_memory_bytes);
-
-    let result = tokio::select! {
-        result = config.run() => result,
-
-        // Ignore this task if it returns due to a querying error.
-        Ok(oom) = check_oom => {
-            // Convert the OOM notification to an error, so we can log it below.
-            let err = anyhow::format_err!("out of memory: {} bytes available, {} required", oom.available_bytes, oom.min_bytes);
-            Err(err)
-        },
-
-        _shutdown = shutdown_listener => {
-            Ok(())
+    let result = match min_available_memory_bytes {
+        0 => {
+            log::info!("memory watchdog is disabled: this task may fail suddenly if it runs out of memory.");
+            config.run().await
+        }
+        _ => {
+            // If the memory limit is 0, this will never return.
+            let check_oom = out_of_memory(min_available_memory_bytes);
+
+            tokio::select! {
+                result = config.run() => result,
+
+                // Ignore this task if it returns due to a querying error.
+                Ok(oom) = check_oom => {
+                    // Convert the OOM notification to an error, so we can log it below.
+                    let err = anyhow::format_err!("out of memory: {} bytes available, {} required", oom.available_bytes, oom.min_bytes);
+                    Err(err)
+                },
+
+                _shutdown = shutdown_listener => {
+                    Ok(())
+                }
+            }
         }
     };
 
@@ -131,9 +139,7 @@ const MAX_OOM_QUERY_ERRORS: usize = 5;
 //
 // Parameterized to enable future configuration by VMSS.
 async fn out_of_memory(min_bytes: u64) -> Result<OutOfMemory> {
-    if min_bytes == 0 {
-        bail!("available memory minimum is unreachable");
-    }
+    log::info!("memory watchdog is enabled: this task will fail informatively if there are {} bytes ({:.2}MB) or fewer of usable memory left.", min_bytes, min_bytes as f64 / 1_000_000f64);
 
     let mut consecutive_query_errors = 0;
 
diff --git a/src/cli/onefuzz/api.py b/src/cli/onefuzz/api.py
@@ -1061,6 +1061,7 @@ def create(
         module_allowlist: Optional[str] = None,
         source_allowlist: Optional[str] = None,
         task_env: Optional[Dict[str, str]] = None,
+        min_available_memory_mb: Optional[int] = None,
     ) -> models.Task:
         """
         Create a task
@@ -1138,6 +1139,7 @@ def create(
                 module_allowlist=module_allowlist,
                 source_allowlist=source_allowlist,
                 task_env=task_env,
+                min_available_memory_mb=min_available_memory_mb,
             ),
         )
 
diff --git a/src/cli/onefuzz/templates/afl.py b/src/cli/onefuzz/templates/afl.py
@@ -55,6 +55,7 @@ def basic(
         debug: Optional[List[TaskDebugFlag]] = None,
         ensemble_sync_delay: Optional[int] = None,
         extra_setup_container: Optional[Container] = None,
+        min_available_memory_mb: Optional[int] = None,
     ) -> Optional[Job]:
         """
         Basic AFL job
@@ -166,6 +167,7 @@ def basic(
             target_env=target_env,
             debug=debug,
             ensemble_sync_delay=ensemble_sync_delay,
+            min_available_memory_mb=min_available_memory_mb,
         )
 
         report_containers = [
diff --git a/src/cli/onefuzz/templates/libfuzzer.py b/src/cli/onefuzz/templates/libfuzzer.py
@@ -83,6 +83,7 @@ def _create_tasks(
         analyzer_env: Optional[Dict[str, str]] = None,
         tools: Optional[Container] = None,
         task_env: Optional[Dict[str, str]] = None,
+        min_available_memory_mb: Optional[int] = None,
     ) -> None:
         target_options = target_options or []
         regression_containers: List[Tuple[ContainerType, Container]] = [
@@ -127,6 +128,7 @@ def _create_tasks(
             colocate=colocate_all_tasks or colocate_secondary_tasks,
             minimized_stack_depth=minimized_stack_depth,
             task_env=task_env,
+            min_available_memory_mb=min_available_memory_mb,
         )
 
         fuzzer_containers: List[Tuple[ContainerType, Container]] = [
@@ -180,6 +182,7 @@ def _create_tasks(
             check_fuzzer_help=check_fuzzer_help,
             expect_crash_on_failure=expect_crash_on_failure,
             task_env=task_env,
+            min_available_memory_mb=min_available_memory_mb,
         )
 
         prereq_tasks = [fuzzer_task.task_id, regression_task.task_id]
@@ -243,6 +246,7 @@ def _create_tasks(
             module_allowlist=module_allowlist,
             source_allowlist=source_allowlist,
             task_env=task_env,
+            min_available_memory_mb=min_available_memory_mb,
         )
 
         report_containers: List[Tuple[ContainerType, Container]] = [
@@ -280,6 +284,7 @@ def _create_tasks(
             colocate=colocate_all_tasks or colocate_secondary_tasks,
             minimized_stack_depth=minimized_stack_depth,
             task_env=task_env,
+            min_available_memory_mb=min_available_memory_mb,
         )
 
         if analyzer_exe is not None:
@@ -320,6 +325,7 @@ def _create_tasks(
                 debug=debug,
                 target_timeout=target_timeout,
                 task_env=task_env,
+                min_available_memory_mb=min_available_memory_mb,
             )
 
     def basic(
@@ -369,6 +375,7 @@ def basic(
         extra_output_container: Optional[Container] = None,
         crashes: Optional[Container] = None,
         task_env: Optional[Dict[str, str]] = None,
+        min_available_memory_mb: Optional[int] = None,
     ) -> Optional[Job]:
         """
         Basic libfuzzer job
@@ -507,6 +514,7 @@ def basic(
             tools=tools,
             task_env=task_env,
             target_timeout=target_timeout,
+            min_available_memory_mb=min_available_memory_mb,
         )
 
         self.logger.info("done creating tasks")
@@ -543,6 +551,7 @@ def merge(
         no_check_fuzzer_help: bool = False,
         extra_setup_container: Optional[Container] = None,
         task_env: Optional[Dict[str, str]] = None,
+        min_available_memory_mb: Optional[int] = None,
     ) -> Optional[Job]:
         """
         libfuzzer merge task
@@ -657,6 +666,7 @@ def merge(
             preserve_existing_outputs=preserve_existing_outputs,
             check_fuzzer_help=check_fuzzer_help,
             task_env=task_env,
+            min_available_memory_mb=min_available_memory_mb,
         )
 
         self.logger.info("done creating tasks")
@@ -698,6 +708,7 @@ def dotnet(
         extra_setup_container: Optional[Container] = None,
         crashes: Optional[Container] = None,
         task_env: Optional[Dict[str, str]] = None,
+        min_available_memory_mb: Optional[int] = None,
     ) -> Optional[Job]:
         pool = self.onefuzz.pools.get(pool_name)
 
@@ -820,6 +831,7 @@ def dotnet(
             expect_crash_on_failure=expect_crash_on_failure,
             check_fuzzer_help=False,
             task_env=task_env,
+            min_available_memory_mb=min_available_memory_mb,
         )
 
         # Ensure the fuzzing task starts before we schedule the coverage and
@@ -874,6 +886,7 @@ def dotnet(
             debug=debug,
             colocate=colocate_all_tasks or colocate_secondary_tasks,
             task_env=task_env,
+            min_available_memory_mb=min_available_memory_mb,
         )
 
         report_containers = [
@@ -915,6 +928,7 @@ def dotnet(
             debug=debug,
             colocate=colocate_all_tasks or colocate_secondary_tasks,
             task_env=task_env,
+            min_available_memory_mb=min_available_memory_mb,
         )
 
         self.logger.info("done creating tasks")
@@ -956,6 +970,7 @@ def qemu_user(
         crashes: Optional[Container] = None,
         readonly_inputs: Optional[Container] = None,
         task_env: Optional[Dict[str, str]] = None,
+        min_available_memory_mb: Optional[int] = None,
     ) -> Optional[Job]:
         """
         libfuzzer tasks, wrapped via qemu-user (PREVIEW FEATURE)
@@ -1125,6 +1140,7 @@ def qemu_user(
             expect_crash_on_failure=False,
             check_fuzzer_help=check_fuzzer_help,
             task_env=task_env,
+            min_available_memory_mb=min_available_memory_mb,
         )
 
         report_containers = [
@@ -1167,6 +1183,7 @@ def qemu_user(
             expect_crash_on_failure=False,
             check_fuzzer_help=check_fuzzer_help,
             task_env=task_env,
+            min_available_memory_mb=min_available_memory_mb,
         )
 
         self.logger.info("done creating tasks")
diff --git a/src/cli/onefuzz/templates/ossfuzz.py b/src/cli/onefuzz/templates/ossfuzz.py
@@ -120,6 +120,7 @@ def libfuzzer(
         debug: Optional[List[TaskDebugFlag]] = None,
         ensemble_sync_delay: Optional[int] = None,
         extra_setup_container: Optional[Container] = None,
+        min_available_memory_mb: Optional[int] = None,
     ) -> None:
         """
         OssFuzz style libfuzzer jobs
@@ -257,6 +258,7 @@ def libfuzzer(
                 tags=helper.tags,
                 debug=debug,
                 ensemble_sync_delay=ensemble_sync_delay,
+                min_available_memory_mb=min_available_memory_mb,
             )
             helpers.append(helper)
         base_helper.wait()
diff --git a/src/cli/onefuzz/templates/radamsa.py b/src/cli/onefuzz/templates/radamsa.py
diff --git a/src/cli/onefuzz/templates/regression.py b/src/cli/onefuzz/templates/regression.py
diff --git a/src/pytypes/onefuzztypes/models.py b/src/pytypes/onefuzztypes/models.py
