Merge pull request #31 from mrc-ide/mrc-4785

Enforce maximum object size

R/task.R

@@ -574,6 +574,8 @@ task_variables <- function(names, envir, environment, root, call = NULL) {
 
     locals <- rlang::env_get_list(envir, nms_locals, inherit = TRUE,
                                   last = topenv())
+    check_locals_size(locals, call = call)
+
     validate_globals <- getOption("hipercow.validate_globals", FALSE)
     if (validate_globals && length(nms_globals) > 0) {
       globals <- rlang::env_get_list(envir, nms_globals, inherit = TRUE,
@@ -622,3 +624,29 @@ show_progress <- function(progress, call = NULL) {
     progress
   }
 }
+
+
+check_locals_size <- function(locals, call = NULL) {
+  if (length(locals) == 0) {
+    return()
+  }
+  max_size <- getOption("hipercow.max_size_local", 1e6)
+  if (!is.finite(max_size)) {
+    return()
+  }
+  size <- vnapply(locals, object.size)
+  err <- names(locals)[size > max_size]
+  if (length(err) > 0) {
+    max_size_bytes <- format_bytes(max_size)
+    cli::cli_abort(
+      c("Object{?s} too large to save with task: {squote(err)}",
+        x = "Objects saved with a hipercow task can only be {max_size_bytes}",
+        i = paste("You can increase the limit by increasing the value of",
+                  "the option 'hipercow.max_size_local', even using 'Inf' to",
+                  "disable this check entirely"),
+        i = paste("Better again, create large objects from your 'sources'",
+                  "argument to your environment, and then advertise this",
+                  "using the 'globals' argument")),
+      call = call)
+  }
+}

R/util.R

@@ -53,6 +53,11 @@ squote <- function(x) {
 }
 
 
+vnapply <- function(...) {
+  vapply(..., FUN.VALUE = 1)
+}
+
+
 vcapply <- function(...) {
   vapply(..., FUN.VALUE = "")
 }
@@ -70,3 +75,14 @@ saverds_if_different <- function(object, path) {
   }
   !skip
 }
+
+
+format_bytes <- function(x) {
+  if (x >= 1e6) {
+    sprintf("%s MB", round(x / 1e6, 3))
+  } else if (x >= 1e3) {
+    sprintf("%s kB", round(x / 1e3, 3))
+  } else {
+    sprintf("%s bytes", x)
+  }
+}

tests/testthat/test-task.R

@@ -268,3 +268,36 @@ test_that("can map task status to logical for task_wait", {
   expect_error(final_status_to_logical("created"),
                "Unhandled status 'created'")
 })
+
+
+test_that("check that locals are not too big", {
+  expect_no_error(check_locals_size(list()))
+  expect_no_error(
+    withr::with_options(list(hipercow.max_size_local = -Inf),
+                        check_locals_size(list(a = 1))))
+  expect_no_error(
+    withr::with_options(list(hipercow.max_size_local = 10000),
+                        check_locals_size(list(a = 1))))
+  expect_error(
+    withr::with_options(list(hipercow.max_size_local = 1),
+                        check_locals_size(list(a = 1))),
+    "Object too large to save with task: 'a'")
+  expect_error(
+    withr::with_options(list(hipercow.max_size_local = 1),
+                        check_locals_size(list(a = 1, b = 2))),
+    "Objects too large to save with task: 'a'.+ 'b'")
+})
+
+
+test_that("prevent large objects being saved by default", {
+  withr::local_options(hipercow.max_size_local = NULL)
+
+  path <- withr::local_tempdir()
+  init_quietly(path)
+  a <- runif(1e6)
+  err <- expect_error(
+    withr::with_dir(path, task_create_expr(mean(a))),
+    "Object too large to save with task: 'a'")
+  expect_match(err$body[[1]],
+               "Objects saved with a hipercow task can only be 1 MB")
+})

tests/testthat/test-util.R

@@ -108,3 +108,13 @@ test_that("saverds_if_different does not update file when not different", {
   expect_false(saverds_if_different(list(1, "two"), path))
   mockery::expect_called(mock_saverds, 0)
 })
+
+
+test_that("can format bytes", {
+  expect_equal(format_bytes(100), "100 bytes")
+  expect_equal(format_bytes(999), "999 bytes")
+  expect_equal(format_bytes(1000), "1 kB")
+  expect_equal(format_bytes(999999), "999.999 kB")
+  expect_equal(format_bytes(1000000), "1 MB")
+  expect_equal(format_bytes(1000000000), "1000 MB")
+})