diff --git a/src/carnot/funcs/metadata/metadata_ops.cc b/src/carnot/funcs/metadata/metadata_ops.cc index 3fe4e21692d..0d2dbe7cb1a 100644 --- a/src/carnot/funcs/metadata/metadata_ops.cc +++ b/src/carnot/funcs/metadata/metadata_ops.cc @@ -51,6 +51,7 @@ void RegisterMetadataOpsOrDie(px::carnot::udf::Registry* registry) { registry->RegisterOrDie("pod_id_to_pod_labels"); registry->RegisterOrDie("pod_id_to_namespace"); registry->RegisterOrDie("pod_id_to_node_name"); + registry->RegisterOrDie("pod_id_to_host_ip"); registry->RegisterOrDie("pod_id_to_replicaset_name"); registry->RegisterOrDie("pod_id_to_replicaset_id"); registry->RegisterOrDie("pod_id_to_deployment_name"); @@ -68,6 +69,7 @@ void RegisterMetadataOpsOrDie(px::carnot::udf::Registry* registry) { registry->RegisterOrDie("pod_name_to_deployment_id"); registry->RegisterOrDie("pod_name_to_pod_id"); registry->RegisterOrDie("pod_name_to_pod_ip"); + registry->RegisterOrDie("pod_name_to_host_ip"); registry->RegisterOrDie("pod_name_to_service_name"); registry->RegisterOrDie("pod_name_to_service_id"); registry->RegisterOrDie("pod_name_to_start_time"); diff --git a/src/carnot/funcs/metadata/metadata_ops.h b/src/carnot/funcs/metadata/metadata_ops.h index 241079858a4..33534a7c46f 100644 --- a/src/carnot/funcs/metadata/metadata_ops.h +++ b/src/carnot/funcs/metadata/metadata_ops.h @@ -188,6 +188,29 @@ class PodNameToPodIPUDF : public ScalarUDF { } }; +class PodNameToHostIPUDF : public ScalarUDF { + public: + StringValue Exec(FunctionContext* ctx, StringValue pod_name) { + auto md = GetMetadataState(ctx); + StringValue pod_id = PodNameToPodIDUDF::GetPodID(md, pod_name); + const px::md::PodInfo* pod_info = md->k8s_metadata_state().PodInfoByID(pod_id); + if (pod_info == nullptr) { + return ""; + } + return pod_info->host_ip(); + } + + static udf::ScalarUDFDocBuilder Doc() { + return udf::ScalarUDFDocBuilder("Get the IP address of the node a pod is running on.") + .Details( + "Gets the IP address of the host (node) that the pod (specified by name) is " + "assigned to.") + .Example("df.host_ip = px.pod_name_to_host_ip(df.pod_name)") + .Arg("pod_name", "The name of the pod to get the host IP for.") + .Returns("The host IP for the pod name passed in."); + } +}; + class PodIDToNamespaceUDF : public ScalarUDF { public: StringValue Exec(FunctionContext* ctx, StringValue pod_id) { @@ -1945,6 +1968,29 @@ class PodIDToNodeNameUDF : public ScalarUDF { } }; +class PodIDToHostIPUDF : public ScalarUDF { + public: + StringValue Exec(FunctionContext* ctx, StringValue pod_id) { + auto md = GetMetadataState(ctx); + + const auto* pod_info = md->k8s_metadata_state().PodInfoByID(pod_id); + if (pod_info == nullptr) { + return ""; + } + return pod_info->host_ip(); + } + + static udf::ScalarUDFDocBuilder Doc() { + return udf::ScalarUDFDocBuilder("Get the IP address of the node a pod ID is running on.") + .Details( + "Gets the IP address of the host (node) that the Pod (specified by Pod ID) is " + "assigned to.") + .Example("df.host_ip = px.pod_id_to_host_ip(df.pod_id)") + .Arg("pod_id", "The Pod ID of the Pod to get the host IP for.") + .Returns("The host IP for the Pod ID passed in."); + } +}; + /** * @brief Returns the ReplicaSet name of a pod ID passed in. */ diff --git a/src/carnot/funcs/metadata/metadata_ops_test.cc b/src/carnot/funcs/metadata/metadata_ops_test.cc index b9d225326bb..eda4bd9d14a 100644 --- a/src/carnot/funcs/metadata/metadata_ops_test.cc +++ b/src/carnot/funcs/metadata/metadata_ops_test.cc @@ -311,6 +311,14 @@ TEST_F(MetadataOpsTest, pod_id_to_node_name_test) { udf_tester.ForInput("123_uid").Expect(""); } +TEST_F(MetadataOpsTest, pod_id_to_host_ip_test) { + auto function_ctx = std::make_unique(metadata_state_, nullptr); + auto udf_tester = px::carnot::udf::UDFTester(std::move(function_ctx)); + udf_tester.ForInput("1_uid").Expect("2.2.2.2"); + // This pod is not available, should return empty. + udf_tester.ForInput("123_uid").Expect(""); +} + TEST_F(MetadataOpsTest, pod_id_to_replicaset_name_test) { auto function_ctx = std::make_unique(metadata_state_, nullptr); auto udf_tester = px::carnot::udf::UDFTester(std::move(function_ctx)); @@ -651,6 +659,14 @@ TEST_F(MetadataOpsTest, pod_name_to_pod_ip) { udf_tester.ForInput("pl/nonexistent").Expect(""); } +TEST_F(MetadataOpsTest, pod_name_to_host_ip) { + auto function_ctx = std::make_unique(metadata_state_, nullptr); + auto udf_tester = px::carnot::udf::UDFTester(std::move(function_ctx)); + udf_tester.ForInput("pl/running_pod").Expect("2.2.2.2"); + udf_tester.ForInput("pl/terminating_pod").Expect(""); + udf_tester.ForInput("pl/nonexistent").Expect(""); +} + TEST_F(MetadataOpsTest, container_id_to_container_status) { ContainerIDToContainerStatusUDF status_udf; diff --git a/src/shared/k8s/metadatapb/test_proto.h b/src/shared/k8s/metadatapb/test_proto.h index ab05aa34f9c..cf1d75542df 100644 --- a/src/shared/k8s/metadatapb/test_proto.h +++ b/src/shared/k8s/metadatapb/test_proto.h @@ -56,6 +56,7 @@ qos_class: QOS_CLASS_GUARANTEED node_name: "test_node" hostname: "test_host" pod_ip: "1.1.1.1" +host_ip: "2.2.2.2" phase: RUNNING message: "Running message" reason: "Running reason" diff --git a/src/shared/metadata/k8s_objects.h b/src/shared/metadata/k8s_objects.h index bb7ac10a117..77d8b611706 100644 --- a/src/shared/metadata/k8s_objects.h +++ b/src/shared/metadata/k8s_objects.h @@ -322,11 +322,13 @@ class PodInfo : public K8sMetadataObject { void set_node_name(std::string_view node_name) { node_name_ = node_name; } void set_hostname(std::string_view hostname) { hostname_ = hostname; } void set_pod_ip(std::string_view pod_ip) { pod_ip_ = pod_ip; } + void set_host_ip(std::string_view host_ip) { host_ip_ = host_ip; } void set_pod_labels(std::string labels) { labels_ = labels; } const std::string& node_name() const { return node_name_; } const std::string& hostname() const { return hostname_; } const std::string& pod_ip() const { return pod_ip_; } + const std::string& host_ip() const { return host_ip_; } const std::string& labels() const { return labels_; } const absl::flat_hash_set& containers() const { return containers_; } @@ -367,6 +369,8 @@ class PodInfo : public K8sMetadataObject { std::string node_name_; std::string hostname_; std::string pod_ip_; + // IP address of the host (node) the pod is assigned to. + std::string host_ip_; std::string labels_; }; diff --git a/src/shared/metadata/metadata_state.cc b/src/shared/metadata/metadata_state.cc index 098d95179c5..e9162addad2 100644 --- a/src/shared/metadata/metadata_state.cc +++ b/src/shared/metadata/metadata_state.cc @@ -287,6 +287,7 @@ Status K8sMetadataState::HandlePodUpdate(const PodUpdate& update) { pod_info->set_node_name(update.node_name()); pod_info->set_hostname(update.hostname()); pod_info->set_pod_ip(update.pod_ip()); + pod_info->set_host_ip(update.host_ip()); pod_info->set_phase(ConvertToPodPhase(update.phase())); pod_info->set_conditions(ConvertToPodConditions(update.conditions())); pod_info->set_phase_message(update.message()); diff --git a/src/shared/metadata/metadata_state_test.cc b/src/shared/metadata/metadata_state_test.cc index d463df5cc87..0fb6662d471 100644 --- a/src/shared/metadata/metadata_state_test.cc +++ b/src/shared/metadata/metadata_state_test.cc @@ -385,6 +385,7 @@ TEST(K8sMetadataStateTest, HandlePodUpdate) { EXPECT_EQ("a_node", pod_info->node_name()); EXPECT_EQ("a_host", pod_info->hostname()); EXPECT_EQ("1.2.3.4", pod_info->pod_ip()); + EXPECT_EQ("5.6.7.8", pod_info->host_ip()); EXPECT_THAT(pod_info->containers(), UnorderedElementsAre("container0_uid")); EXPECT_THAT(pod_info->owner_references(), UnorderedElementsAre(OwnerReference{"1234", "rs0", "ReplicaSet"}));