From 336ef68d0d5eae781b688197690e54bbf8988515 Mon Sep 17 00:00:00 2001
From: Sonic Build Admin <sonicbld@microsoft.com>
Date: Tue, 11 Mar 2025 23:44:57 +0000
Subject: [PATCH] Use pzstd to compress the Docker in SWI slim images

<!--
     Please make sure you've read and understood our contributing guidelines:
     https://github.com/Azure/SONiC/blob/gh-pages/CONTRIBUTING.md

     ** Make sure all your commits include a signature generated with `git commit -s` **

     If this is a bug fix, make sure your description includes "fixes #xxxx", or
     "closes #xxxx" or "resolves #xxxx"

     Please provide the following information:
-->

#### Why I did it
Same with https://github.com/sonic-net/sonic-buildimage/pull/21824

Enabled the docker inram feature for slim image.
It would extract the docker image to ram during the boot, so it would take extra times during boot.

##### Work item tracking
- Microsoft ADO **(number only)**:
31323281

#### How I did it
Use pzstd which is more efficient tool to compress and decompress the docker file to reduce the boot time.
Currently, we do not modify the "FILESYSTEM_DOCKERFS=dockerfs.tar.gz" in onie-image.conf, so for slim image, we still use dockerfs.tar.gz as file name but actually with zstd compressed.
We tried to find out a way to adjust the file name in onie-image.conf, but it seems not easy to do that.
So we use the file cmd to determine the compressing type in union-mount.j2, then use the related cmd to extract the docker file.
Plan to support zstd for all types of images in the future to unify the docker file image .

#### How to verify it
Boot swi slim image and normal image, boot mellanox bin image, all boot successfully.

<!--
If PR needs to be backported, then the PR must be tested against the base branch and the earliest backport release branch and provide tested image version on these two branches. For example, if the PR is requested for master, 202211 and 202012, then the requester needs to provide test results on master and 202012.
-->

#### Which release branch to backport (provide reason below if selected)

<!--
- Note we only backport fixes to a release branch, *not* features!
- Please also provide a reason for the backporting below.
- e.g.
- [x] 202006
-->

- [ ] 201811
- [ ] 201911
- [ ] 202006
- [ ] 202012
- [ ] 202106
- [ ] 202111
- [ ] 202205
- [ ] 202211
- [ ] 202305

#### Tested branch (Please provide the tested image version)

<!--
- Please provide tested image version
- e.g.
- [x] 20201231.100
-->

- [ ] <!-- image version 1 -->
- [ ] <!-- image version 2 -->

#### Description for the changelog
<!--
Write a short (one line) summary that describes the changes in this
pull request for inclusion in the changelog:
-->

<!--
 Ensure to add label/tag for the feature raised. example - PR#2174 under sonic-utilities repo. where, Generic Config and Update feature has been labelled as GCU.
-->

#### Link to config_db schema for YANG module changes
<!--
Provide a link to config_db schema for the table for which YANG model
is defined
Link should point to correct section on https://github.com/Azure/sonic-buildimage/blob/master/src/sonic-yang-models/doc/Configuration.md
-->

#### A picture of a cute animal (not mandatory but encouraged)
---
 build_debian.sh                      | 12 +++++++++++-
 files/dsc/install_debian.j2          |  4 ++++
 files/initramfs-tools/file           | 18 ++++++++++++++++++
 files/initramfs-tools/pzstd          | 16 ++++++++++++++++
 files/initramfs-tools/union-mount.j2 | 12 +++++++++---
 5 files changed, 58 insertions(+), 4 deletions(-)
 create mode 100644 files/initramfs-tools/file
 create mode 100644 files/initramfs-tools/pzstd

diff --git a/build_debian.sh b/build_debian.sh
index 7fe04902fb2..764fef95520 100755
--- a/build_debian.sh
+++ b/build_debian.sh
@@ -391,6 +391,12 @@ sudo LANG=C DEBIAN_FRONTEND=noninteractive chroot $FILESYSTEM_ROOT apt-get -y in
     zstd                    \
     nvme-cli
 
+sudo cp files/initramfs-tools/pzstd $FILESYSTEM_ROOT/etc/initramfs-tools/hooks/pzstd
+sudo chmod +x $FILESYSTEM_ROOT/etc/initramfs-tools/hooks/pzstd
+
+sudo cp files/initramfs-tools/file $FILESYSTEM_ROOT/etc/initramfs-tools/hooks/file
+sudo chmod +x $FILESYSTEM_ROOT/etc/initramfs-tools/hooks/file
+
 # Have systemd create the auditd log directory
 sudo mkdir -p ${FILESYSTEM_ROOT}/etc/systemd/system/auditd.service.d
 sudo tee ${FILESYSTEM_ROOT}/etc/systemd/system/auditd.service.d/log-directory.conf >/dev/null <<EOF
@@ -851,7 +857,11 @@ if [[ $MULTIARCH_QEMU_ENVIRON == y || $CROSS_BUILD_ENVIRON == y ]]; then
 fi
 
 ## Compress docker files
-pushd $FILESYSTEM_ROOT && sudo tar -I pigz -cf $OLDPWD/$FILESYSTEM_DOCKERFS -C ${DOCKERFS_PATH}var/lib/docker .; popd
+if [ "$BUILD_REDUCE_IMAGE_SIZE" = "y" ]; then
+    pushd $FILESYSTEM_ROOT && sudo tar -I pzstd -cf $OLDPWD/$FILESYSTEM_DOCKERFS -C ${DOCKERFS_PATH}var/lib/docker .; popd
+else
+    pushd $FILESYSTEM_ROOT && sudo tar -I pigz -cf $OLDPWD/$FILESYSTEM_DOCKERFS -C ${DOCKERFS_PATH}var/lib/docker .; popd
+fi
 
 ## Compress together with /boot, /var/lib/docker and $PLATFORM_DIR as an installer payload zip file
 pushd $FILESYSTEM_ROOT && sudo tar -I pigz -cf platform.tar.gz -C $PLATFORM_DIR . && sudo zip -n .gz $OLDPWD/$INSTALLER_PAYLOAD -r boot/ platform.tar.gz; popd
diff --git a/files/dsc/install_debian.j2 b/files/dsc/install_debian.j2
index def17524976..a8397c974c9 100755
--- a/files/dsc/install_debian.j2
+++ b/files/dsc/install_debian.j2
@@ -15,7 +15,11 @@ image_dir=image-$image_version
 
 INSTALLER_PAYLOAD=fs.zip
 DOCKERFS_DIR=docker
+{% if BUILD_REDUCE_IMAGE_SIZE == "y" -%}
+FILESYSTEM_DOCKERFS=dockerfs.tar.zstd
+{%- else -%}
 FILESYSTEM_DOCKERFS=dockerfs.tar.gz
+{%- endif %}
 BL_CONF=boot.conf
 
 DATA_PARTUUID=6ED62003-DD8D-44B8-9538-0A2B7C7E628F
diff --git a/files/initramfs-tools/file b/files/initramfs-tools/file
new file mode 100644
index 00000000000..cff6acd18fd
--- /dev/null
+++ b/files/initramfs-tools/file
@@ -0,0 +1,18 @@
+#!/bin/sh
+set -e
+PREREQ=""
+prereqs() {
+    echo "$PREREQ"
+}
+case "$1" in
+    prereqs)
+        prereqs
+        exit 0
+        ;;
+esac
+. /usr/share/initramfs-tools/hook-functions
+# Include file binary
+copy_exec /usr/bin/file /usr/bin
+# Include magic database
+copy_exec /usr/lib/file/magic.mgc /etc
+exit 0
diff --git a/files/initramfs-tools/pzstd b/files/initramfs-tools/pzstd
new file mode 100644
index 00000000000..8775d480b62
--- /dev/null
+++ b/files/initramfs-tools/pzstd
@@ -0,0 +1,16 @@
+#!/bin/sh
+set -e
+PREREQ=""
+prereqs() {
+    echo "$PREREQ"
+}
+case "$1" in
+    prereqs)
+        prereqs
+        exit 0
+        ;;
+esac
+. /usr/share/initramfs-tools/hook-functions
+# Include pzstd binary
+copy_exec /usr/bin/pzstd /usr/bin
+exit 0
diff --git a/files/initramfs-tools/union-mount.j2 b/files/initramfs-tools/union-mount.j2
index 8f8abb8f6af..de2cb5c7596 100644
--- a/files/initramfs-tools/union-mount.j2
+++ b/files/initramfs-tools/union-mount.j2
@@ -130,9 +130,15 @@ extract_dockerfs()
 {
     echo "Extracting {{ FILESYSTEM_DOCKERFS }}"
     if [ -f "${rootmnt}/host/$image_dir/{{ FILESYSTEM_DOCKERFS }}" ] && [ "$secureboot" = false ]; then
-        # Extract dockerfs.tar.gz into /var/lib/docker unless the system booted with secureboot
-        # In secureboot dockerfs.tar.gz cannot be trusted as it does not have a signature
-        tar xz --numeric-owner -f ${rootmnt}/host/$image_dir/{{ FILESYSTEM_DOCKERFS }} -C ${rootmnt}/var/lib/docker
+        # Check if the file is zstd compressed
+        file_type=$(file -b --mime-type "${rootmnt}/host/$image_dir/{{ FILESYSTEM_DOCKERFS }}")
+        if [ "$file_type" = "application/zstd" ]; then
+            echo "Detected zstd compression, extracting with pzstd..."
+            pzstd -d -q ${rootmnt}/host/$image_dir/{{ FILESYSTEM_DOCKERFS }} -c | tar x --numeric-owner -C ${rootmnt}/var/lib/docker
+        else
+            echo "Using default extraction method (gzip assumed)..."
+            tar xz --numeric-owner -f "${rootmnt}/host/$image_dir/{{ FILESYSTEM_DOCKERFS }}" -C "${rootmnt}/var/lib/docker"
+        fi
     elif [ "$bootloader" = "aboot" ] && unzip -l "$swi_path" | grep -q {{ FILESYSTEM_DOCKERFS }}; then
         # Aboot swi images also support extracting dockerfs.tar.gz directly from them
         unzip -qp "$swi_path" {{ FILESYSTEM_DOCKERFS }} | tar xz --numeric-owner -C ${rootmnt}/var/lib/docker