grpc: implement TTL manager to re-buffer pending messages for bidirectional gRPC

source/common/grpc/BUILD

@@ -213,5 +213,15 @@ envoy_cc_library(
     deps = [
         ":typed_async_client_lib",
         "//source/common/protobuf:utility_lib",
+        "@com_google_absl//absl/container:btree",
+    ],
+)
+
+envoy_cc_library(
+    name = "buffered_message_ttl_manager_lib",
+    hdrs = ["buffered_message_ttl_manager.h"],
+    deps = [
+        ":buffered_async_client_lib",
+        "//envoy/event:dispatcher_interface",
     ],
 )

source/common/grpc/buffered_async_client.h

@@ -12,11 +12,33 @@ namespace Grpc {
 
 enum class BufferState { Buffered, PendingFlush };
 
+class BufferedAsyncClientCallbacks {
+public:
+  virtual ~BufferedAsyncClientCallbacks() = default;
+
+  /**
+   * Called if a message related with id has been proceeded successfully.
+   *
+   * @param id the ID related with buffered messages. This ID must be already published by client.
+   * If not, this callback will do nothing.
+   */
+  virtual void onSuccess(uint64_t id) PURE;
+
+  /**
+   * Called if a message related with id has not been proceeded.
+   *
+   * @param id the ID related with buffered messages. This ID must be already published by client.
+   * If not, this callback will do nothing.
+   */
+  virtual void onError(uint64_t id) PURE;
+};
+
 // This class wraps bidirectional gRPC and provides message arrival guarantee.
 // It stores messages to be sent or in the process of being sent in a buffer,
 // and can track the status of the message based on the ID assigned to each message.
 // If a message fails to be sent, it can be re-buffered to guarantee its arrival.
-template <class RequestType, class ResponseType> class BufferedAsyncClient {
+template <class RequestType, class ResponseType>
+class BufferedAsyncClient : public BufferedAsyncClientCallbacks {
 public:
   BufferedAsyncClient(uint32_t max_buffer_bytes, const Protobuf::MethodDescriptor& service_method,
                       Grpc::AsyncStreamCallbacks<ResponseType>& callbacks,
@@ -73,10 +95,13 @@ template <class RequestType, class ResponseType> class BufferedAsyncClient {
     return inflight_message_ids;
   }
 
-  void onSuccess(uint64_t message_id) { erasePendingMessage(message_id); }
+  void onSuccess(uint64_t message_id) override { erasePendingMessage(message_id); }
 
-  void onError(uint64_t message_id) {
-    if (message_buffer_.find(message_id) == message_buffer_.end()) {
+  void onError(uint64_t message_id) override {
+    const auto& message_it = message_buffer_.find(message_id);
+
+    if (message_it == message_buffer_.end() ||
+        message_it->second.first != Grpc::BufferState::PendingFlush) {
       return;
     }
 

source/common/grpc/buffered_message_ttl_manager.h

@@ -0,0 +1,66 @@
+#pragma once
+
+#include <queue>
+
+#include "envoy/event/dispatcher.h"
+
+#include "source/common/grpc/buffered_async_client.h"
+
+namespace Envoy {
+namespace Grpc {
+
+class BufferedMessageTtlManager {
+public:
+  BufferedMessageTtlManager(Event::Dispatcher& dispatcher, BufferedAsyncClientCallbacks& callbacks,
+                            std::chrono::milliseconds message_ack_timeout)
+      : dispatcher_(dispatcher), message_ack_timeout_(message_ack_timeout), callbacks_(callbacks),
+        timer_(dispatcher_.createTimer([this] { checkMessages(); })) {}
+
+  ~BufferedMessageTtlManager() { timer_->disableTimer(); }
+
+  void setDeadline(absl::flat_hash_set<uint64_t>&& ids) {
+    const auto expires_at = dispatcher_.timeSource().monotonicTime() + message_ack_timeout_;
+    deadline_.emplace(expires_at, std::move(ids));
+
+    if (!timer_->enabled()) {
+      timer_->enableTimer(message_ack_timeout_);
+    }
+  }
+
+  const std::queue<std::pair<MonotonicTime, absl::flat_hash_set<uint64_t>>>& deadline() {
+    return deadline_;
+  }
+
+private:
+  void checkMessages() {
+    const auto now = dispatcher_.timeSource().monotonicTime();
+
+    while (!deadline_.empty()) {
+      auto& it = deadline_.front();
+      if (it.first > now) {
+        break;
+      }
+      for (auto&& id : it.second) {
+        // If the retrieved message is a PendingFlush, it means that the message
+        // has timed out. A timeout is treated as an error, and the callback will
+        // re-buffer the message.
+        callbacks_.onError(id);
+      }
+      deadline_.pop();
+    }
+
+    if (!deadline_.empty()) {
+      const auto earliest_timepoint = deadline_.front().first;
+      timer_->enableTimer(
+          std::chrono::duration_cast<std::chrono::milliseconds>(earliest_timepoint - now));
+    }
+  }
+
+  Event::Dispatcher& dispatcher_;
+  std::chrono::milliseconds message_ack_timeout_;
+  BufferedAsyncClientCallbacks& callbacks_;
+  Event::TimerPtr timer_;
+  std::queue<std::pair<MonotonicTime, absl::flat_hash_set<uint64_t>>> deadline_;
+};
+} // namespace Grpc
+} // namespace Envoy

test/common/grpc/BUILD

@@ -202,3 +202,13 @@ envoy_cc_test(
         "@envoy_api//envoy/config/core/v3:pkg_cc_proto",
     ],
 )
+
+envoy_cc_test(
+    name = "buffered_message_ttl_manager_test",
+    srcs = ["buffered_message_ttl_manager_test.cc"],
+    deps = [
+        "//source/common/event:dispatcher_lib",
+        "//source/common/grpc:buffered_message_ttl_manager_lib",
+        "//test/test_common:utility_lib",
+    ],
+)

test/common/grpc/buffered_message_ttl_manager_test.cc

@@ -0,0 +1,67 @@
+#include <chrono>
+#include <memory>
+
+#include "source/common/event/dispatcher_impl.h"
+#include "source/common/grpc/buffered_message_ttl_manager.h"
+
+#include "test/test_common/utility.h"
+
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+namespace Envoy {
+namespace Grpc {
+namespace {
+
+class MockBufferedAsyncClientCallbacks : public BufferedAsyncClientCallbacks {
+public:
+  MOCK_METHOD(void, onSuccess, (uint64_t), ());
+  MOCK_METHOD(void, onError, (uint64_t), ());
+};
+
+using testing::Invoke;
+using testing::NiceMock;
+
+class BufferedMessageTtlManagerTest : public testing::Test {
+public:
+  BufferedMessageTtlManagerTest()
+      : api_(Api::createApiForTest()), dispatcher_(api_->allocateDispatcher("test_thread")) {}
+
+  void SetUp() override {
+    ttl_manager_ = std::make_shared<BufferedMessageTtlManager>(*dispatcher_, callbacks_, msec_);
+  }
+
+  Api::ApiPtr api_;
+  Event::DispatcherPtr dispatcher_;
+  std::shared_ptr<BufferedMessageTtlManager> ttl_manager_;
+  NiceMock<MockBufferedAsyncClientCallbacks> callbacks_;
+  std::chrono::milliseconds msec_{1000};
+};
+
+TEST_F(BufferedMessageTtlManagerTest, Basic) {
+  absl::flat_hash_set<uint64_t> ids{0};
+  EXPECT_CALL(callbacks_, onError(_)).WillOnce(Invoke([this] {
+    EXPECT_EQ(ttl_manager_->deadline().size(), 1);
+    absl::flat_hash_set<uint64_t> ids{1, 2};
+    EXPECT_CALL(callbacks_, onError(_)).Times(2);
+    ttl_manager_->setDeadline(std::move(ids));
+  }));
+  ttl_manager_->setDeadline(std::move(ids));
+
+  dispatcher_->run(Event::Dispatcher::RunType::Block);
+  EXPECT_EQ(ttl_manager_->deadline().size(), 0);
+
+  // Test if deadline queue is empty after queue cleared once.
+  absl::flat_hash_set<uint64_t> ids2{3};
+  EXPECT_CALL(callbacks_, onError(_)).WillOnce(Invoke([this] {
+    EXPECT_EQ(ttl_manager_->deadline().size(), 1);
+  }));
+  ttl_manager_->setDeadline(std::move(ids2));
+
+  dispatcher_->run(Event::Dispatcher::RunType::Block);
+  EXPECT_EQ(ttl_manager_->deadline().size(), 0);
+}
+
+} // namespace
+} // namespace Grpc
+} // namespace Envoy