From 77bd2347c710845e906e430e4baee3203ca8419b Mon Sep 17 00:00:00 2001
From: gNPSI Team <copybara-servicebot@google.com>
Date: Thu, 29 Jan 2026 23:58:01 -0800
Subject: [PATCH] Support addition of CongestionTelemetry to exported data

PiperOrigin-RevId: 863093443
---
 proto/gnpsi/gnpsi.proto           | 2 +-
 server/gnpsi_relay_server_test.cc | 6 +++---
 server/gnpsi_service_impl.cc      | 7 ++++++-
 server/gnpsi_service_impl.h       | 9 ++++++---
 server/mock_gnpsi_service_impl.h  | 3 ++-
 5 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/proto/gnpsi/gnpsi.proto b/proto/gnpsi/gnpsi.proto
index 3933994..bc26bcd 100644
--- a/proto/gnpsi/gnpsi.proto
+++ b/proto/gnpsi/gnpsi.proto
@@ -90,7 +90,7 @@ message CongestionTelemetry {
   int32 egress_port = 103;
 
   // Sequence number of the sample.
-  int32 sample_sequence_number = 104;
+  int64 sample_sequence_number = 104;
 }
 
 // gNPSI sample that can contain SFlow/NetFlow/IPFIX data.
diff --git a/server/gnpsi_relay_server_test.cc b/server/gnpsi_relay_server_test.cc
index bb9c3d2..2d62152 100644
--- a/server/gnpsi_relay_server_test.cc
+++ b/server/gnpsi_relay_server_test.cc
@@ -58,7 +58,7 @@ TEST_F(GnpsiRelayServerTest, RelaySuccess) {
         return -1;
       }));
   // Assert a send call made in case of a successful read
-  EXPECT_CALL(gnpsi_service_impl_, SendSamplePacket(_, _)).Times(1);
+  EXPECT_CALL(gnpsi_service_impl_, SendSamplePacket(_, _, _)).Times(1);
   relay_server_.StartRelayAndWait(gnpsi_service_impl_);
 }
 
@@ -77,7 +77,7 @@ TEST_F(GnpsiRelayServerTest, NoDeathOnNonCriticalReadError) {
         return -1;
       }));
   // Assert no send call is made in case non fatal error is encountered
-  EXPECT_CALL(gnpsi_service_impl_, SendSamplePacket(_, _)).Times(0);
+  EXPECT_CALL(gnpsi_service_impl_, SendSamplePacket(_, _, _)).Times(0);
   relay_server_.StartRelayAndWait(gnpsi_service_impl_);
 }
 
@@ -108,7 +108,7 @@ TEST_F(GnpsiRelayServerTest, DeathOnCriticalReadError) {
     return -1;
   }));
   // Assert no send call is made in case fatal error is encountered
-  EXPECT_CALL(gnpsi_service_impl_, SendSamplePacket(_, _)).Times(0);
+  EXPECT_CALL(gnpsi_service_impl_, SendSamplePacket(_, _, _)).Times(0);
   relay_server_.StartRelayAndWait(gnpsi_service_impl_);
 }
 }  // namespace gnpsi
diff --git a/server/gnpsi_service_impl.cc b/server/gnpsi_service_impl.cc
index a862cae..c60111e 100644
--- a/server/gnpsi_service_impl.cc
+++ b/server/gnpsi_service_impl.cc
@@ -151,12 +151,17 @@ void GnpsiServiceImpl::UndrainConnections() {
 }
 
 void GnpsiServiceImpl::SendSamplePacket(
-    const std::string& sample_packet, ::gnpsi::SFlowMetadata::Version version) {
+    const std::string& sample_packet, ::gnpsi::SFlowMetadata::Version version,
+    std::vector<::gnpsi::CongestionTelemetry> congestion_telemetry) {
   absl::MutexLock l(&mu_);
   Sample response;
   response.set_packet(sample_packet);
   response.set_timestamp(absl::ToUnixNanos(absl::Now()));
   response.mutable_sflow_metadata()->set_version(version);
+  if (!congestion_telemetry.empty()) {
+    response.mutable_congestion_telemetry()->Add(congestion_telemetry.begin(),
+                                                 congestion_telemetry.end());
+  }
   for (auto it = gnpsi_connections_.begin(), end = gnpsi_connections_.end();
        it != end; it++) {
     auto connection = *it;
diff --git a/server/gnpsi_service_impl.h b/server/gnpsi_service_impl.h
index 147dd70..a88f5ce 100644
--- a/server/gnpsi_service_impl.h
+++ b/server/gnpsi_service_impl.h
@@ -45,7 +45,8 @@ class GnpsiSenderInterface {
   virtual ~GnpsiSenderInterface() = default;
   virtual void SendSamplePacket(
       const std::string& sample_packet,
-      SFlowMetadata::Version version = SFlowMetadata::V5) = 0;
+      SFlowMetadata::Version version = SFlowMetadata::V5,
+      std::vector<CongestionTelemetry> congestion_telemetry = {}) = 0;
   virtual void DrainConnections() = 0;
   virtual void UndrainConnections() = 0;
   virtual std::vector<GnpsiStats> GetStats() = 0;
@@ -125,8 +126,10 @@ class GnpsiServiceImpl : public ::gnpsi::gNPSI::Service,
       ABSL_LOCKS_EXCLUDED(mu_) override;
 
   // Sends Sample response to each client.
-  void SendSamplePacket(const std::string& sample_packet,
-                        SFlowMetadata::Version version = SFlowMetadata::V5)
+  void SendSamplePacket(
+      const std::string& sample_packet,
+      SFlowMetadata::Version version = SFlowMetadata::V5,
+      std::vector<CongestionTelemetry> congestion_telemetry = {})
       ABSL_LOCKS_EXCLUDED(mu_) override;
 
   // Closes all current conections and blocks any new incoming connections.
diff --git a/server/mock_gnpsi_service_impl.h b/server/mock_gnpsi_service_impl.h
index 9f829e5..a7192e2 100644
--- a/server/mock_gnpsi_service_impl.h
+++ b/server/mock_gnpsi_service_impl.h
@@ -11,7 +11,8 @@ class MockGnpsiServiceImpl : public GnpsiSenderInterface {
   virtual ~MockGnpsiServiceImpl() = default;
   MOCK_METHOD(void, SendSamplePacket,
               (const std::string& sample_packet,
-               SFlowMetadata::Version version),
+               SFlowMetadata::Version version,
+               std::vector<CongestionTelemetry> congestion_telemetry),
               (override));
   MOCK_METHOD(void, DrainConnections, (), (override));
   MOCK_METHOD(void, UndrainConnections, (), (override));