[Thinkit] Create a function that rewrites table entries to only rely on ports from a given gnmi_config.

tests/lib/BUILD.bazel

@@ -58,18 +58,55 @@ cc_library(
     srcs = ["switch_test_setup_helpers.cc"],
     hdrs = ["switch_test_setup_helpers.h"],
     deps = [
+        "//gutil:collections",
+        "//gutil:proto",
         "//lib/gnmi:gnmi_helper",
         "//lib/gnmi:openconfig_cc_proto",
         "//lib/validator:validator_lib",
+        "//p4_pdpi:ir_cc_proto",
+        "//p4_pdpi:ir_tools",
         "//p4_pdpi:p4_runtime_session",
         "//thinkit:mirror_testbed",
         "//thinkit:switch",
         "@com_github_p4lang_p4runtime//:p4info_cc_proto",
         "@com_github_p4lang_p4runtime//:p4runtime_cc_proto",
+        "@com_github_p4lang_p4runtime//:p4types_cc_proto",
+        "@com_google_absl//absl/container:btree",
+        "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/time",
         "@com_google_absl//absl/types:optional",
+        "@com_google_absl//absl/types:span",
+    ],
+)
+
+# go/golden-test-with-coverage
+cc_test(
+    name = "switch_test_setup_helpers_golden_test_runner",
+    srcs = ["switch_test_setup_helpers_golden_test_runner.cc"],
+    data = ["switch_test_setup_helpers.expected"],
+    linkstatic = True,
+    deps = [
+        ":switch_test_setup_helpers",
+        "//gutil:status_matchers",
+        "//gutil:testing",
+        "//p4_pdpi:ir_cc_proto",
+        "//sai_p4/instantiations/google:instantiations",
+        "//sai_p4/instantiations/google:sai_p4info_cc",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/types:span",
+        "@com_google_protobuf//:protobuf",
+    ],
+)
+
+cmd_diff_test(
+    name = "switch_test_setup_helpers_golden_test",
+    actual_cmd = "$(execpath :switch_test_setup_helpers_golden_test_runner)",
+    expected = "//tests/lib:switch_test_setup_helpers.expected",
+    tools = [
+        ":switch_test_setup_helpers_golden_test_runner",
     ],
 )

tests/lib/switch_test_setup_helpers.cc

@@ -1,24 +1,39 @@
 #include "tests/lib/switch_test_setup_helpers.h"
 
+#include <deque>
+#include <functional>
+#include <future>  // NOLINT: third_party library.
 #include <memory>
 #include <optional>
 #include <string>
 #include <utility>
+#include <vector>
 
+#include "absl/container/btree_set.h"
+#include "absl/container/flat_hash_map.h"
 #include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
 #include "absl/time/time.h"
 #include "absl/types/optional.h"
+#include "absl/types/span.h"
+#include "glog/logging.h"
+#include "gutil/collections.h"
+#include "gutil/proto.h"
 #include "gutil/status.h"
 #include "lib/gnmi/gnmi_helper.h"
 #include "lib/gnmi/openconfig.pb.h"
 #include "lib/validator/validator_lib.h"
+#include "p4/config/v1/p4types.pb.h"
 #include "p4/v1/p4runtime.pb.h"
+#include "p4_pdpi/ir_tools.h"
 #include "p4_pdpi/p4_runtime_session.h"
 
 namespace pins_test {
 namespace {
 
 constexpr absl::Duration kGnmiTimeoutDefault = absl::Minutes(3);
+constexpr char kPortNamedType[] = "port_id_t";
 
 absl::StatusOr<std::unique_ptr<pdpi::P4RuntimeSession>>
 CreateP4RuntimeSessionAndClearExistingState(
@@ -50,6 +65,34 @@ absl::Status SetForwardingPipelineConfigAndAssureClearedTables(
   return absl::OkStatus();
 }
 
+// Uses the `port_map` to remap any P4runtime ports in `entries`.
+absl::Status RewritePortsInTableEntries(
+    const pdpi::IrP4Info& info, std::vector<pdpi::IrTableEntry>& entries,
+    const absl::flat_hash_map<std::string, std::string>& port_map) {
+  p4::config::v1::P4NamedType port_type;
+  port_type.set_name(kPortNamedType);
+  RETURN_IF_ERROR(pdpi::TransformValuesOfType(
+      info, port_type, entries, [&](absl::string_view old_port) {
+        return gutil::FindOrStatus(port_map, std::string(old_port));
+      }));
+
+  // Watch ports do not have a named type, but we still consider them ports so
+  // we have to deal with them specifically rather than using the generic
+  // rewriting function above.
+  for (pdpi::IrTableEntry& entry : entries) {
+    if (entry.has_action_set()) {
+      for (auto& action : *entry.mutable_action_set()->mutable_actions()) {
+        if (!action.watch_port().empty()) {
+          ASSIGN_OR_RETURN(*action.mutable_watch_port(),
+                           gutil::FindOrStatus(port_map, action.watch_port()));
+        }
+      }
+    }
+  }
+
+  return absl::OkStatus();
+}
+
 }  // namespace
 
 absl::StatusOr<std::unique_ptr<pdpi::P4RuntimeSession>>
@@ -181,4 +224,90 @@ absl::Status WaitForEnabledInterfacesToBeUp(
                                      /*with_healthz=*/false);
 }
 
+// Gets every P4runtime port used in `entries`.
+absl::StatusOr<absl::btree_set<std::string>> GetPortsUsed(
+    const pdpi::IrP4Info& info, std::vector<pdpi::IrTableEntry> entries) {
+  absl::btree_set<std::string> ports;
+  p4::config::v1::P4NamedType port_type;
+  port_type.set_name(kPortNamedType);
+  RETURN_IF_ERROR(pdpi::TransformValuesOfType(info, port_type, entries,
+                                              [&](absl::string_view port) {
+                                                ports.insert(std::string(port));
+                                                return std::string(port);
+                                              }));
+
+  // Watch ports do not have a named type, but we still consider them ports so
+  // we have to deal with them specifically rather than using the generic
+  // function above.
+  for (pdpi::IrTableEntry& entry : entries) {
+    if (entry.has_action_set()) {
+      for (const auto& action : entry.action_set().actions()) {
+        if (!action.watch_port().empty()) {
+          ports.insert(action.watch_port());
+        }
+      }
+    }
+  }
+
+  return ports;
+}
+
+// Remaps ports in a round-robin fashion, but starts by fixing those that are
+// both used in `entries` and in `ports`.
+absl::Status RewritePortsInTableEntries(
+    const pdpi::IrP4Info& info, absl::Span<const std::string> new_ports,
+    std::vector<pdpi::IrTableEntry>& entries) {
+  if (new_ports.empty()) {
+    return absl::InvalidArgumentError("`new_ports` may not be empty");
+  }
+
+  // We get the ports used in the original entries so we can preserve any that
+  // exist in both, and then remap them in a round-robin fashion.
+  ASSIGN_OR_RETURN(absl::btree_set<std::string> ports_used_originally,
+                   GetPortsUsed(info, entries));
+
+  absl::flat_hash_map<std::string, std::string> old_to_new_port_id;
+  // Queue of which new port to map an old port to next.
+  std::deque<std::string> new_port_queue;
+
+  for (const auto& new_port : new_ports) {
+    if (ports_used_originally.contains(new_port)) {
+      // Make sure that existing ports are preserved and add them to the back of
+      // the queue for balancing.
+      old_to_new_port_id[new_port] = new_port;
+      new_port_queue.push_back(new_port);
+    } else {
+      new_port_queue.push_front(new_port);
+    }
+  }
+
+  // Map all old ports to new ports.
+  for (const auto& old_port : ports_used_originally) {
+    // If a port is already mapped, we should not remap it.
+    if (old_to_new_port_id.contains(old_port)) continue;
+    const std::string new_port = new_port_queue.front();
+    old_to_new_port_id[old_port] = new_port;
+    new_port_queue.pop_front();
+    new_port_queue.push_back(new_port);
+  }
+
+  return RewritePortsInTableEntries(info, entries, old_to_new_port_id);
+}
+
+absl::Status RewritePortsInTableEntries(
+    const pdpi::IrP4Info& info, absl::string_view gnmi_config,
+    std::vector<pdpi::IrTableEntry>& entries) {
+  ASSIGN_OR_RETURN(absl::btree_set<std::string> valid_port_ids_set,
+                   GetAllPortIds(gnmi_config));
+
+  if (valid_port_ids_set.empty()) {
+    return absl::InvalidArgumentError(
+        "given gnmi_config had no valid port ids");
+  }
+
+  std::vector<std::string> new_ports(valid_port_ids_set.begin(),
+                                     valid_port_ids_set.end());
+  return RewritePortsInTableEntries(info, new_ports, entries);
+}
+
 }  // namespace pins_test

tests/lib/switch_test_setup_helpers.expected

@@ -246,7 +246,7 @@ matches {
   }
 }
 action {
-  name: "set_nexthop"
+  name: "set_ip_nexthop"
   params {
     name: "router_interface_id"
     value {
@@ -726,4 +726,3 @@ modified: action_set.actions[3].watch_port: "14" -> "1"
 modified: action_set.actions[4].watch_port: "15" -> "a_port"
 modified: action_set.actions[5].watch_port: "16" -> "2"
 modified: action_set.actions[6].watch_port: "17" -> "1"
-

tests/lib/switch_test_setup_helpers.h

@@ -5,12 +5,16 @@
 #include <optional>
 #include <string>
 #include <utility>
+#include <vector>
 
+#include "absl/container/btree_set.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
+#include "absl/types/span.h"
 #include "absl/time/time.h"
 #include "lib/gnmi/openconfig.pb.h"
 #include "p4/config/v1/p4info.pb.h"
+#include "p4_pdpi/ir.pb.h"
 #include "p4_pdpi/p4_runtime_session.h"
 #include "thinkit/mirror_testbed.h"
 #include "thinkit/switch.h"
@@ -64,6 +68,32 @@ absl::Status WaitForEnabledInterfacesToBeUp(
         std::function<void(const openconfig::Interfaces& enabled_interfaces)>>
         on_failure = std::nullopt);
 
+// Gets a count of which P4runtime ports are used in `entries` and how
+// frequently.
+absl::StatusOr<absl::btree_set<std::string>> GetPortsUsed(
+    const pdpi::IrP4Info& info, std::vector<pdpi::IrTableEntry> entries);
+
+// Rewrites the ports of the given table `entries` to only use the given
+// non-empty `ports`. Uses `info` to determine which values refer to ports.
+// Specifically, for each port `x` in the set of table entries, this
+// function remaps it to f(x) such that:
+// - if `x \in ports` then `f(x) = x`
+// - `x` is remapped to an `f(x) \in ports` (chosen deterministically)
+//   such that `\forall. p1,p2 \in ports. |{f(x) = p1 | x \in
+//   all_ports(entries)}| <= |{f(x) = p2 | x \in all_ports(entries)}| + 1`.
+//
+// This function makes it possible to use a list of pre-existing table entries
+// with any set of ports desired (or configured on the switch).
+absl::Status RewritePortsInTableEntries(
+    const pdpi::IrP4Info& info, absl::Span<const std::string> new_ports,
+    std::vector<pdpi::IrTableEntry>& entries);
+
+// Extracts the available ports from the given `gnmi_config` and rewrites
+// entries as per above.
+absl::Status RewritePortsInTableEntries(
+    const pdpi::IrP4Info& info, absl::string_view gnmi_config,
+    std::vector<pdpi::IrTableEntry>& entries);
+
 }  // namespace pins_test
 
 #endif  // PINS_TESTS_LIB_SWITCH_TEST_SETUP_HELPERS_H_

tests/lib/switch_test_setup_helpers_golden_test_runner.cc

@@ -307,7 +307,7 @@ int main(int argc, char** argv) {
               exact { str: "nexthop-1" }
             }
             action {
-              name: "set_nexthop"
+              name: "set_ip_nexthop"
               params {
                 name: "router_interface_id"
                 value { str: "router-interface-1" }