From b2a7765b272315dd912fb4567daed2ccd533bbf2 Mon Sep 17 00:00:00 2001
From: gNPSI Team <copybara-servicebot@google.com>
Date: Fri, 24 May 2024 04:39:45 -0700
Subject: [PATCH] Implement undrain functionality in gNPSI service.

PiperOrigin-RevId: 636871959
---
 server/gnpsi_service_impl.cc | 5 +++++
 server/gnpsi_service_impl.h  | 3 +++
 2 files changed, 8 insertions(+)

diff --git a/server/gnpsi_service_impl.cc b/server/gnpsi_service_impl.cc
index 6cf1bf5..a862cae 100644
--- 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
index 0f200ed..a70b0a7 100644
--- 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;