Enhance smartswitch environment variables parsing (#21209)

Fix #20284

In 202405 and above, two extra steps are added before the start of every container which checks NUM_DPU and IS_DPU_DEVICE by parsing the platform.json file using the jq tool. This is only relevant for Smartswitch. However, this is adding some delay during the reconciliation phase of WR/FR resulting

How I did it
Set the environment variables for systemd by systemd-sonic-generator.

Signed-off-by: Ze Gan <ganze718@gmail.com>

Author: Pterosaur

files/build_templates/docker_image_ctl.j2

@@ -372,18 +372,6 @@ start() {
         source $PLATFORM_ENV_CONF
     fi
 
-    # Parse the platform.json file to get the platform specific information
-    PLATFORM_JSON=/usr/share/sonic/device/$PLATFORM/platform.json
-    if [ -f "$PLATFORM_JSON" ]; then
-        NUM_DPU=$(jq -r '.DPUS | length' $PLATFORM_JSON 2>/dev/null)
-        jq -e '.DPU' $PLATFORM_JSON >/dev/null
-        if [[ $? -eq 0 ]]; then
-            IS_DPU_DEVICE="true"
-        else
-            IS_DPU_DEVICE="false"
-        fi
-    fi
-
     {%- if sonic_asic_platform == "broadcom" %}
     {%- if docker_container_name == "syncd" %}
     # Set the SYNCD_SHM_SIZE if this variable not defined

src/systemd-sonic-generator/ssg-test.cc

@@ -231,6 +231,7 @@ class SsgMainTest : public SsgFunctionTest {
             while (getline(file, line) && !found) {
                 if (str == line) {
                     found = true;
+                    break;
                 }
             }
         return found;
@@ -364,6 +365,32 @@ class SsgMainTest : public SsgFunctionTest {
             "system", !cfg.is_smart_switch_npu && !cfg.is_smart_switch_dpu, cfg.num_dpus, false);
     }
 
+    void validate_environment_variable(const SsgMainConfig &cfg) {
+        std::unordered_map<std::string, std::string> env_vars;
+        env_vars["IS_DPU_DEVICE"] = (cfg.is_smart_switch_dpu ? "true" : "false");
+        env_vars["NUM_DPU"] = std::to_string(cfg.num_dpus);
+
+        std::vector<std::string> checked_service_list;
+
+        checked_service_list.insert(checked_service_list.end(), common_service_list.begin(), common_service_list.end());
+        if (cfg.num_dpus > 0) {
+            checked_service_list.insert(checked_service_list.end(), npu_service_list.begin(), npu_service_list.end());
+        }
+        if (cfg.num_asics > 1) {
+            checked_service_list.insert(checked_service_list.end(), multi_asic_service_list.begin(), multi_asic_service_list.end());
+        }
+
+        for (const auto &target: checked_service_list) {
+            if (find_string_in_file("[Service]", target)) {
+                for (const auto& item : env_vars) {
+                    std::string str = "Environment=\"" + item.first + "=" + item.second + "\"";
+                    EXPECT_EQ(find_string_in_file(str, target), true)
+                        << "Error validating " + str + " in " + target;
+                }
+            }
+        }
+    }
+
     /* ssg_main test routine.
      * input: num_asics    number of asics
      */
@@ -431,6 +458,9 @@ class SsgMainTest : public SsgFunctionTest {
 
         /* Validate Test Unit file for dependency creation. */
         validate_depedency_in_unit_file(cfg);
+
+        /* Validate environment variables */
+        validate_environment_variable(cfg);
     }
 
     /* Save global variables before running tests */
@@ -719,7 +749,6 @@ TEST_F(SsgMainTest, ssg_main_argv) {
         std::vector<std::string> arguments = {
                     "ssg_main",
                 };
-
         /* Create argv list for ssg_main. */
         for (const auto& arg : arguments) {
             argv_.push_back((char*)arg.data());

src/systemd-sonic-generator/systemd-sonic-generator.cpp

@@ -11,6 +11,9 @@
 #include <string>
 #include <sstream>
 #include <unordered_set>
+#include <fstream>
+#include <unordered_map>
+#include <regex>
 
 #define MAX_NUM_TARGETS 48
 #define MAX_NUM_INSTALL_LINES 48
@@ -385,6 +388,56 @@ static void replace_multi_inst_dep(const char *src) {
     rename(tmp_file_path, src);
 }
 
+static void update_environment(const std::string &src)
+{
+    std::ifstream src_file(src);
+    std::ofstream tmp_file(src + ".tmp");
+    bool has_service_section = false;
+    std::string line;
+
+    // locate the [Service] section
+    while (std::getline(src_file, line)) {
+        tmp_file << line << std::endl;
+        if (line.find("[Service]") != std::string::npos) {
+            has_service_section = true;
+            break;
+        }
+    }
+
+    if (!has_service_section) {
+        tmp_file << "[Service]" << std::endl;
+    }
+
+    std::unordered_map<std::string, std::string> env_vars;
+    env_vars["IS_DPU_DEVICE"] = (smart_switch_dpu ? "true" : "false");
+    env_vars["NUM_DPU"] = std::to_string(num_dpus);
+
+    for (const auto& [key, value] : env_vars) {
+        tmp_file << "Environment=\"" << key << "=" << value << "\"" << std::endl;
+    }
+
+    // copy the rest of the file
+    if (has_service_section) {
+        const static std::regex env_var_regex("Environment=\"?(\\w+)");
+        while (std::getline(src_file, line)) {
+            // skip the existing environment variables
+            std::smatch match;
+            if (std::regex_search(line, match, env_var_regex)) {
+                if (env_vars.find(match[1]) != env_vars.end()) {
+                    continue;
+                }
+            }
+            tmp_file << line << std::endl;
+        }
+    }
+
+    src_file.close();
+    tmp_file.close();
+    // remove the original file and rename the temporary file
+    remove(src.c_str());
+    rename((src + ".tmp").c_str(), src.c_str());
+}
+
 int get_install_targets(std::string unit_file, char* targets[]) {
     /***
     Returns install targets for a unit file
@@ -1155,6 +1208,8 @@ int ssg_main(int argc, char **argv) {
             free(targets[j]);
         }
 
+        update_environment(get_unit_file_prefix() + unit_instance);
+
         free(unit_files[i]);
     }