Repository: microsoft/onefuzz
File: src/agent/onefuzz-task/src/managed/cmd.rs (Rust)

Code before:
    oneshot_receiver.send(receive_from_task)?;

    let shutdown_listener = task::spawn_blocking(move || loop {
        match receive_from_agent.recv() {
            Ok(msg) => info!("Received unexpected message from agent: {:?}", msg),
            Err(ipc::IpcError::Disconnected) => {
                info!("Agent disconnected from the IPC channel. Shutting down");
                break;
            }
            Err(ipc::IpcError::Bincode(e)) => {
                error!("BinCode error receiving message from agent: {:?}", e);
                break;
            }
            Err(ipc::IpcError::Io(e)) => {
                error!("IO error receiving message from agent: {:?}", e);
                break;
            }
        }
    });

    init_telemetry(config.common()).await;

    let min_available_memory_bytes = 1_000_000 * config.common().min_available_memory_mb;

    // If the memory limit is 0, this will never return.
    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(())
        }
    };

    if let Err(err) = &result {
        error!("error running task: {:?}", err);
    }

    onefuzz_telemetry::try_flush_and_close().await;

    result
}


Review comment (suggestion):
let's disable the check altogether if the memory limit is not >0 instead of launching a task that will loop forever

Code after:
    oneshot_receiver.send(receive_from_task)?;

    let shutdown_listener = task::spawn_blocking(move || loop {
        match receive_from_agent.recv() {
            Ok(msg) => info!("Received unexpected message from agent: {:?}", msg),
            Err(ipc::IpcError::Disconnected) => {
                info!("Agent disconnected from the IPC channel. Shutting down");
                break;
            }
            Err(ipc::IpcError::Bincode(e)) => {
                error!("BinCode error receiving message from agent: {:?}", e);
                break;
            }
            Err(ipc::IpcError::Io(e)) => {
                error!("IO error receiving message from agent: {:?}", e);
                break;
            }
        }
    });

    init_telemetry(config.common()).await;

    let min_available_memory_bytes = 1_000_000 * config.common().min_available_memory_mb;

    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(())
                }
            }
        }
    };

    if let Err(err) = &result {