build_image.py recognize BOARD_SYSTEMIMAGE_PARTITION_RESERVED_SIZE

- Copy "use_logical_partitions" to *_image_info.txt before sending
  it to build_image.py, so that the script can use this variable.

- build_image.py emits an additional properties file to inform
  the build system about the system image size.

Test: `make systemimage`

Test: `make systemimage` with the following:
    - install a large file to system image
    fails as expected (because _PARTITION_SIZE is exceeded)

Test: `make systemimage` with the following:
    - set PRODUCT_USE_LOGICAL_PARTITIONS to true
    - set BOARD_SYSTEMIMAGE_PARTITION_RESERVED_SIZE
    fails as expected (BOARD_SYSTEMIMAGE_PARTITION_SIZE needs
    to be undefined)

Test: `make systemimage` with the following:
    - install a large file to system image
    - set PRODUCT_USE_LOGICAL_PARTIIONS to true
    - add a small BOARD_SYSTEMIMAGE_PARTITION_RESERVED_SIZE
    - remove BOARD_SYSTEMIMAGE_PARTITION_SIZE
    build succeeds.

Test: same for systemotherimage

Bug: 79106666

Change-Id: I574062882acd1ecd633ac38c5a8c5351b90a32d8
diff --git a/core/Makefile b/core/Makefile
index 38ee46c..27252e2 100644
--- a/core/Makefile
+++ b/core/Makefile
@@ -1086,6 +1086,17 @@
 INTERNAL_USERIMAGES_DEPS += $(MKE2FS_CONF)
 endif
 
+ifeq (true,$(USE_LOGICAL_PARTITIONS))
+
+ifeq ($(PRODUCTS.$(INTERNAL_PRODUCT).PRODUCT_SUPPORTS_VERITY),true)
+  $(error vboot 1.0 doesn't support logical partition)
+endif
+
+# TODO(b/80195851): Should not define BOARD_AVB_SYSTEM_KEY_PATH without
+# BOARD_AVB_SYSTEM_DETACHED_VBMETA.
+
+endif # USE_LOGICAL_PARTITIONS
+
 # $(1): the path of the output dictionary file
 # $(2): a subset of "system vendor cache userdata product oem"
 # $(3): additional "key=value" pairs to append to the dictionary file.
@@ -1104,6 +1115,7 @@
     $(if $(BOARD_SYSTEMIMAGE_SQUASHFS_DISABLE_4K_ALIGN),$(hide) echo "system_squashfs_disable_4k_align=$(BOARD_SYSTEMIMAGE_SQUASHFS_DISABLE_4K_ALIGN)" >> $(1))
     $(if $(PRODUCTS.$(INTERNAL_PRODUCT).PRODUCT_SYSTEM_BASE_FS_PATH),$(hide) echo "system_base_fs_file=$(PRODUCTS.$(INTERNAL_PRODUCT).PRODUCT_SYSTEM_BASE_FS_PATH)" >> $(1))
     $(if $(PRODUCTS.$(INTERNAL_PRODUCT).PRODUCT_SYSTEM_HEADROOM),$(hide) echo "system_headroom=$(PRODUCTS.$(INTERNAL_PRODUCT).PRODUCT_SYSTEM_HEADROOM)" >> $(1))
+    $(if $(BOARD_SYSTEMIMAGE_PARTITION_RESERVED_SIZE),$(hide) echo "system_reserved_size=$(BOARD_SYSTEMIMAGE_PARTITION_RESERVED_SIZE)" >> $(1))
 )
 $(if $(BOARD_EXT4_SHARE_DUP_BLOCKS),$(hide) echo "ext4_share_dup_blocks=$(BOARD_EXT4_SHARE_DUP_BLOCKS)" >> $(1))
 $(if $(filter $(2),userdata),\
@@ -1190,6 +1202,7 @@
 $(if $(filter true,$(BOARD_BUILD_SYSTEM_ROOT_IMAGE)),\
     $(hide) echo "system_root_image=true" >> $(1)
     $(hide) echo "ramdisk_dir=$(TARGET_ROOT_OUT)" >> $(1))
+$(if $(USE_LOGICAL_PARTITIONS),$(hide) echo "use_logical_partitions=true" >> $(1))
 $(if $(3),$(hide) $(foreach kv,$(3),echo "$(kv)" >> $(1);))
 endef
 
@@ -1199,6 +1212,12 @@
 $(call generate-image-prop-dictionary,$(1),system vendor cache userdata product oem,$(2))
 endef
 
+# $(1): the path of the input dictionary file, where each line has the format key=value
+# $(2): the key to look up
+define read-image-prop-dictionary
+$$(grep '$(2)=' $(1) | cut -f2- -d'=')
+endef
+
 # $(1): modules list
 # $(2): output dir
 # $(3): mount point
@@ -1595,10 +1614,8 @@
   $(hide) PATH=$(foreach p,$(INTERNAL_USERIMAGES_BINARY_PATHS),$(p):)$$PATH \
       build/make/tools/releasetools/build_image.py \
       $(TARGET_OUT) $(systemimage_intermediates)/system_image_info.txt $(1) $(TARGET_OUT) \
-      || ( echo "Out of space? the tree size of $(TARGET_OUT) is (MB): " 1>&2 ;\
-           du -sm $(TARGET_OUT) 1>&2;\
-           echo "The max is $$(( $(BOARD_SYSTEMIMAGE_PARTITION_SIZE) / 1048576 )) MB." 1>&2 ;\
-           mkdir -p $(DIST_DIR); cp $(INSTALLED_FILES_FILE) $(DIST_DIR)/installed-files-rescued.txt; \
+      $(systemimage_intermediates)/generated_system_image_info.txt \
+      || ( mkdir -p $(DIST_DIR); cp $(INSTALLED_FILES_FILE) $(DIST_DIR)/installed-files-rescued.txt; \
            exit 1 )
 endef
 
@@ -1640,7 +1657,9 @@
 $(INSTALLED_SYSTEMIMAGE): $(BUILT_SYSTEMIMAGE) $(RECOVERY_FROM_BOOT_PATCH)
 	@echo "Install system fs image: $@"
 	$(copy-file-to-target)
-	$(hide) $(call assert-max-image-size,$@ $(RECOVERY_FROM_BOOT_PATCH),$(BOARD_SYSTEMIMAGE_PARTITION_SIZE))
+	$(hide) $(call assert-max-image-size,$@ $(RECOVERY_FROM_BOOT_PATCH),\
+		$(call read-image-prop-dictionary,\
+			$(systemimage_intermediates)/generated_system_image_info.txt,system_size))
 
 systemimage: $(INSTALLED_SYSTEMIMAGE)
 
@@ -1649,7 +1668,9 @@
 	            | $(INTERNAL_USERIMAGES_DEPS)
 	@echo "make $@: ignoring dependencies"
 	$(call build-systemimage-target,$(INSTALLED_SYSTEMIMAGE))
-	$(hide) $(call assert-max-image-size,$(INSTALLED_SYSTEMIMAGE),$(BOARD_SYSTEMIMAGE_PARTITION_SIZE))
+	$(hide) $(call assert-max-image-size,$(INSTALLED_SYSTEMIMAGE),\
+		$(call read-image-prop-dictionary,\
+			$(systemimage_intermediates)/generated_system_image_info.txt,system_size))
 
 ifneq (,$(filter systemimage-nodeps snod, $(MAKECMDGOALS)))
 ifeq (true,$(WITH_DEXPREOPT))
@@ -1991,8 +2012,11 @@
   $(call generate-image-prop-dictionary, $(systemotherimage_intermediates)/system_other_image_info.txt,system,skip_fsck=true)
   $(hide) PATH=$(foreach p,$(INTERNAL_USERIMAGES_BINARY_PATHS),$(p):)$$PATH \
       build/make/tools/releasetools/build_image.py \
-      $(TARGET_OUT_SYSTEM_OTHER) $(systemotherimage_intermediates)/system_other_image_info.txt $(INSTALLED_SYSTEMOTHERIMAGE_TARGET) $(TARGET_OUT)
-  $(hide) $(call assert-max-image-size,$(INSTALLED_SYSTEMOTHERIMAGE_TARGET),$(BOARD_SYSTEMIMAGE_PARTITION_SIZE))
+      $(TARGET_OUT_SYSTEM_OTHER) $(systemotherimage_intermediates)/system_other_image_info.txt $(INSTALLED_SYSTEMOTHERIMAGE_TARGET) $(TARGET_OUT)\
+      $(systemotherimage_intermediates)/generated_system_other_image_info.txt
+  $(hide) $(call assert-max-image-size,$(INSTALLED_SYSTEMOTHERIMAGE_TARGET),\
+    $(call read-image-prop-dictionary,\
+      $(systemotherimage_intermediates)/generated_system_other_image_info.txt,system_size))
 endef
 
 # We just build this directly to the install location.
@@ -2292,10 +2316,6 @@
   $(error BOARD_BOOTIMAGE_PARTITION_SIZE must be set for BOARD_AVB_ENABLE)
 endif
 
-ifndef BOARD_SYSTEMIMAGE_PARTITION_SIZE
-  $(error BOARD_SYSTEMIMAGE_PARTITION_SIZE must be set for BOARD_AVB_ENABLE)
-endif
-
 # $(1): the directory to extract public keys to
 define extract-avb-chain-public-keys
   $(if $(BOARD_AVB_BOOT_KEY_PATH),\
diff --git a/core/config.mk b/core/config.mk
index 7bbb3e5..f92375c 100644
--- a/core/config.mk
+++ b/core/config.mk
@@ -928,7 +928,15 @@
 
 ifeq ($(USE_LOGICAL_PARTITIONS),true)
   BOARD_KERNEL_CMDLINE += androidboot.logical_partitions=1
+
+ifneq ($(BOARD_SYSTEMIMAGE_PARTITION_SIZE),)
+ifneq ($(BOARD_SYSTEMIMAGE_PARTITION_RESERVED_SIZE),)
+$(error Should not define BOARD_SYSTEMIMAGE_PARTITION_SIZE and \
+    BOARD_SYSTEMIMAGE_PARTITION_RESERVED_SIZE together)
 endif
+endif
+
+endif # USE_LOGICAL_PARTITIONS
 
 # ###############################################################
 # Set up final options.
diff --git a/core/product.mk b/core/product.mk
index f22a3e5..6a5add3 100644
--- a/core/product.mk
+++ b/core/product.mk
@@ -376,6 +376,10 @@
 	WITH_DEXPREOPT \
 	WITH_DEXPREOPT_BOOT_IMG_AND_SYSTEM_SERVER_ONLY
 
+# Logical partitions related variables.
+_product_stash_var_list += \
+	BOARD_SYSTEMIMAGE_PARTITION_RESERVED_SIZE \
+
 #
 # Mark the variables in _product_stash_var_list as readonly
 #
diff --git a/tools/releasetools/build_image.py b/tools/releasetools/build_image.py
index 2406f4a..4e690fe 100755
--- a/tools/releasetools/build_image.py
+++ b/tools/releasetools/build_image.py
@@ -18,8 +18,10 @@
 Builds output_image from the given input_directory, properties_file,
 and writes the image to target_output_directory.
 
+If argument generated_prop_file exists, write additional properties to the file.
+
 Usage:  build_image.py input_directory properties_file output_image \\
-            target_output_directory
+            target_output_directory [generated_prop_file]
 """
 
 from __future__ import print_function
@@ -40,22 +42,29 @@
 
 FIXED_SALT = "aee087a5be3b982978c923f566a94613496b417f2af592639bc80d141e34dfe7"
 BLOCK_SIZE = 4096
+BYTES_IN_MB = 1024 * 1024
 
 
-def RunCommand(cmd, verbose=None):
+def RunCommand(cmd, verbose=None, env=None):
   """Echo and run the given command.
 
   Args:
     cmd: the command represented as a list of strings.
     verbose: show commands being executed.
+    env: a dictionary of additional environment variables.
   Returns:
     A tuple of the output and the exit code.
   """
+  env_copy = None
+  if env is not None:
+    env_copy = os.environ.copy()
+    env_copy.update(env)
   if verbose is None:
     verbose = OPTIONS.verbose
   if verbose:
     print("Running: " + " ".join(cmd))
-  p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+  p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
+                       env=env_copy)
   output, _ = p.communicate()
 
   if verbose:
@@ -103,6 +112,24 @@
   return verity_size
 
 
+def GetDiskUsage(path):
+  """Return number of bytes that "path" occupies on host.
+
+  Args:
+    path: The directory or file to calculate size on
+  Returns:
+    True and the number of bytes if successful,
+    False and 0 otherwise.
+  """
+  env = {"POSIXLY_CORRECT": "1"}
+  cmd = ["du", "-s", path]
+  output, exit_code = RunCommand(cmd, verbose=False, env=env)
+  if exit_code != 0:
+    return False, 0
+  # POSIX du returns number of blocks with block size 512
+  return True, int(output.split()[0]) * 512
+
+
 def GetSimgSize(image_file):
   simg = sparse_img.SparseImage(image_file, build_map=False)
   return simg.blocksize * simg.total_blocks
@@ -442,6 +469,8 @@
 
 def BuildImage(in_dir, prop_dict, out_file, target_out=None):
   """Build an image to out_file from in_dir with property prop_dict.
+  After the function call, values in prop_dict is updated with
+  computed values.
 
   Args:
     in_dir: path of input directory.
@@ -486,6 +515,21 @@
   verity_supported = prop_dict.get("verity") == "true"
   verity_fec_supported = prop_dict.get("verity_fec") == "true"
 
+  if (prop_dict.get("use_logical_partitions") == "true" and
+      "partition_size" not in prop_dict):
+    # if partition_size is not defined, use output of `du' + reserved_size
+    success, size = GetDiskUsage(origin_in)
+    if not success:
+      return False
+    if OPTIONS.verbose:
+      print("The tree size of %s is %d MB." % (origin_in, size // BYTES_IN_MB))
+    size += int(prop_dict.get("partition_reserved_size", 0))
+    # Round this up to a multiple of 4K so that avbtool works
+    size = common.RoundUpTo4K(size)
+    prop_dict["partition_size"] = str(size)
+    if OPTIONS.verbose:
+      print("Allocating %d MB for %s." % (size // BYTES_IN_MB, out_file))
+
   # Adjust the partition size to make room for the hashes if this is to be
   # verified.
   if verity_supported and is_verity_partition:
@@ -613,6 +657,17 @@
   if exit_code != 0:
     print("Error: '%s' failed with exit code %d:\n%s" % (
         build_command, exit_code, mkfs_output))
+    success, du = GetDiskUsage(origin_in)
+    du_str = ("%d bytes (%d MB)" % (du, du // BYTES_IN_MB)
+             ) if success else "unknown"
+    print("Out of space? The tree size of %s is %s.\n" % (
+        origin_in, du_str))
+    print("The max is %d bytes (%d MB).\n" % (
+        int(prop_dict["partition_size"]),
+        int(prop_dict["partition_size"]) // BYTES_IN_MB))
+    print("Reserved space is %d bytes (%d MB).\n" % (
+        int(prop_dict.get("partition_reserved_size", 0)),
+        int(prop_dict.get("partition_reserved_size", 0)) // BYTES_IN_MB))
     return False
 
   # Check if there's enough headroom space available for ext4 image.
@@ -713,6 +768,7 @@
       "avb_enable",
       "avb_avbtool",
       "avb_salt",
+      "use_logical_partitions",
   )
   for p in common_props:
     copy_prop(p, p)
@@ -745,6 +801,7 @@
     copy_prop("system_extfs_inode_count", "extfs_inode_count")
     if not copy_prop("system_extfs_rsv_pct", "extfs_rsv_pct"):
       d["extfs_rsv_pct"] = "0"
+    copy_prop("system_reserved_size", "partition_reserved_size")
   elif mount_point == "system_other":
     # We inherit the selinux policies of /system since we contain some of its
     # files.
@@ -767,6 +824,7 @@
     copy_prop("system_extfs_inode_count", "extfs_inode_count")
     if not copy_prop("system_extfs_rsv_pct", "extfs_rsv_pct"):
       d["extfs_rsv_pct"] = "0"
+    copy_prop("system_reserved_size", "partition_reserved_size")
   elif mount_point == "data":
     # Copy the generic fs type first, override with specific one if available.
     copy_prop("fs_type", "fs_type")
@@ -842,8 +900,27 @@
   return d
 
 
+def GlobalDictFromImageProp(image_prop, mount_point):
+  d = {}
+  def copy_prop(src_p, dest_p):
+    if src_p in image_prop:
+      d[dest_p] = image_prop[src_p]
+      return True
+    return False
+  if mount_point == "system":
+    copy_prop("partition_size", "system_size")
+  elif mount_point == "system_other":
+    copy_prop("partition_size", "system_size")
+  return d
+
+
+def SaveGlobalDict(filename, glob_dict):
+  with open(filename, "w") as f:
+    f.writelines(["%s=%s" % (key, value) for (key, value) in glob_dict.items()])
+
+
 def main(argv):
-  if len(argv) != 4:
+  if len(argv) < 4 or len(argv) > 5:
     print(__doc__)
     sys.exit(1)
 
@@ -851,6 +928,7 @@
   glob_dict_file = argv[1]
   out_file = argv[2]
   target_out = argv[3]
+  prop_file_out = argv[4] if len(argv) >= 5 else None
 
   glob_dict = LoadGlobalDict(glob_dict_file)
   if "mount_point" in glob_dict:
@@ -885,6 +963,9 @@
           file=sys.stderr)
     sys.exit(1)
 
+  if prop_file_out:
+    glob_dict_out = GlobalDictFromImageProp(image_properties, mount_point)
+    SaveGlobalDict(prop_file_out, glob_dict_out)
 
 if __name__ == '__main__':
   try: