Implement undrain functionality in gNPSI service.

PiperOrigin-RevId: 636871959
openconfig · May 24, 2024 · b2a7765 · b2a7765
1 parent b27f152
commit b2a7765
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/server/gnpsi_service_impl.cc b/server/gnpsi_service_impl.cc
@@ -145,6 +145,11 @@ void GnpsiServiceImpl::DrainConnections() {
   }
 }
 
+void GnpsiServiceImpl::UndrainConnections() {
+  absl::MutexLock lock(&mu_);
+  service_drained_ = false;
+}
+
 void GnpsiServiceImpl::SendSamplePacket(
     const std::string& sample_packet, ::gnpsi::SFlowMetadata::Version version) {
   absl::MutexLock l(&mu_);

diff --git a/server/gnpsi_service_impl.h b/server/gnpsi_service_impl.h
@@ -47,6 +47,7 @@ class GnpsiSenderInterface {
       const std::string& sample_packet,
       SFlowMetadata::Version version = SFlowMetadata::V5) = 0;
   virtual void DrainConnections() = 0;
+  virtual void UndrainConnections() = 0;
   virtual std::vector<GnpsiStats> GetStats() = 0;
 };
 
@@ -131,6 +132,8 @@ class GnpsiServiceImpl : public ::gnpsi::gNPSI::Service,
   // Closes all current conections and blocks any new incoming connections.
   void DrainConnections() ABSL_LOCKS_EXCLUDED(mu_) override;
 
+  void UndrainConnections() ABSL_LOCKS_EXCLUDED(mu_) override;
+
   // Returns stats per connection collected by the server.
   std::vector<GnpsiStats> GetStats() ABSL_LOCKS_EXCLUDED(mu_) override;