diff --git a/thrift/lib/cpp2/server/metrics/StreamMetricCallback.h b/thrift/lib/cpp2/server/metrics/StreamMetricCallback.h index 49af713246b..39a4afac9e3 100644 --- a/thrift/lib/cpp2/server/metrics/StreamMetricCallback.h +++ b/thrift/lib/cpp2/server/metrics/StreamMetricCallback.h @@ -31,6 +31,8 @@ class StreamMetricCallback { virtual void onStreamCancel(std::string_view /* methodName */) = 0; virtual void onStreamRequestN( std::string_view /* methodName */, uint32_t /* credits */) = 0; + virtual void recordCreditsAvailable( + std::string_view /* methodName */, uint32_t /* creditsAvailable */) = 0; }; class NoopStreamMetricCallback final : public StreamMetricCallback { @@ -41,6 +43,7 @@ class NoopStreamMetricCallback final : public StreamMetricCallback { void onStreamComplete(std::string_view) override {} void onStreamCancel(std::string_view) override {} void onStreamRequestN(std::string_view, uint32_t) override {} + void recordCreditsAvailable(std::string_view, uint32_t) override {} }; } // namespace apache::thrift diff --git a/thrift/lib/cpp2/transport/rocket/server/RocketStreamClientCallback.cpp b/thrift/lib/cpp2/transport/rocket/server/RocketStreamClientCallback.cpp index 47263c726ed..b1d08ce0ec8 100644 --- a/thrift/lib/cpp2/transport/rocket/server/RocketStreamClientCallback.cpp +++ b/thrift/lib/cpp2/transport/rocket/server/RocketStreamClientCallback.cpp @@ -226,8 +226,9 @@ bool RocketStreamClientCallback::request(uint32_t tokens) { } cancelTimeout(); - tokens_ += tokens; + streamMetricCallback_.recordCreditsAvailable(rpcMethodName_, tokens_); streamMetricCallback_.onStreamRequestN(rpcMethodName_, tokens); + tokens_ += tokens; return serverCallback()->onStreamRequestN(tokens); } @@ -292,6 +293,7 @@ void RocketStreamClientCallback::scheduleTimeout() { timeoutCallback_ = std::make_unique(*this); } connection_.scheduleStreamTimeout(timeoutCallback_.get()); + streamMetricCallback_.recordCreditsAvailable(rpcMethodName_, tokens_); } void RocketStreamClientCallback::cancelTimeout() {