refactor grpc: refactor ugrpc::impl::AsyncMethodInvocation

commit_hash:3c4d7444deb7d7fc547142eedf37a42f0c0719f4

Author: kpavlov00

grpc/include/userver/ugrpc/client/impl/async_stream_methods.hpp

@@ -18,10 +18,14 @@ namespace ugrpc::client::impl {
 
 ugrpc::impl::AsyncMethodInvocation::WaitStatus WaitAndTryCancelIfNeeded(
     ugrpc::impl::AsyncMethodInvocation& invocation,
-    grpc::ClientContext& context
+    grpc::ClientContext& client_context
 ) noexcept;
 
-void CheckOk(StreamingCallState& state, ugrpc::impl::AsyncMethodInvocation::WaitStatus status, std::string_view stage);
+void CheckOk(
+    StreamingCallState& state,
+    ugrpc::impl::AsyncMethodInvocation::WaitStatus wait_status,
+    std::string_view stage
+);
 
 void CheckFinishStatus(CallState& state);
 
@@ -33,13 +37,13 @@ void ProcessCancelled(CallState& state, std::string_view stage) noexcept;
 
 void ProcessNetworkError(CallState& state, std::string_view stage) noexcept;
 
-void ThrowIfDeadlineIsExceeded(grpc::ClientContext& context, std::string_view call_name);
+void ThrowIfDeadlineIsExceeded(grpc::ClientContext& client_context, std::string_view call_name);
 
 template <typename GrpcStream>
 void StartCall(GrpcStream& stream, StreamingCallState& state) {
-    ugrpc::impl::AsyncMethodInvocation start_call;
-    stream.StartCall(start_call.GetCompletionTag());
-    CheckOk(state, WaitAndTryCancelIfNeeded(start_call, state.GetClientContext()), "StartCall");
+    ugrpc::impl::AsyncMethodInvocation invocation;
+    stream.StartCall(invocation.GetCompletionTag());
+    CheckOk(state, WaitAndTryCancelIfNeeded(invocation, state.GetClientContext()), "StartCall");
 }
 
 template <typename GrpcStream>
@@ -53,14 +57,13 @@ void Finish(
 
     state.SetFinished();
 
-    FinishAsyncMethodInvocation finish;
-    auto& status = state.GetStatus();
-    stream.Finish(&status, finish.GetCompletionTag());
+    FinishAsyncMethodInvocation invocation;
+    stream.Finish(&state.GetStatus(), invocation.GetCompletionTag());
+    const auto wait_status = WaitAndTryCancelIfNeeded(invocation, state.GetClientContext());
 
-    const auto wait_status = WaitAndTryCancelIfNeeded(finish, state.GetClientContext());
     switch (wait_status) {
         case ugrpc::impl::AsyncMethodInvocation::WaitStatus::kOk:
-            state.GetStatsScope().SetFinishTime(finish.GetFinishTime());
+            state.GetStatsScope().SetFinishTime(invocation.GetFinishTime());
             try {
                 ProcessFinish(state, final_response);
             } catch (const std::exception& ex) {
@@ -76,7 +79,7 @@ void Finish(
             break;
 
         case ugrpc::impl::AsyncMethodInvocation::WaitStatus::kError:
-            state.GetStatsScope().SetFinishTime(finish.GetFinishTime());
+            state.GetStatsScope().SetFinishTime(invocation.GetFinishTime());
             ProcessNetworkError(state, "Finish");
             if (throw_on_error) {
                 ThrowIfDeadlineIsExceeded(state.GetClientContext(), state.GetCallName());
@@ -107,24 +110,16 @@ void FinishAbandoned(GrpcStream& stream, StreamingCallState& state) noexcept try
 
     state.GetClientContext().TryCancel();
 
-    FinishAsyncMethodInvocation finish;
-    stream.Finish(&state.GetStatus(), finish.GetCompletionTag());
-
-    const engine::TaskCancellationBlocker cancel_blocker;
-    const auto wait_status = finish.Wait();
+    FinishAsyncMethodInvocation invocation;
+    stream.Finish(&state.GetStatus(), invocation.GetCompletionTag());
+    const auto ok = invocation.WaitNonCancellable();
 
-    state.GetStatsScope().SetFinishTime(finish.GetFinishTime());
+    state.GetStatsScope().SetFinishTime(invocation.GetFinishTime());
 
-    switch (wait_status) {
-        case ugrpc::impl::AsyncMethodInvocation::WaitStatus::kOk:
-            ProcessFinishAbandoned(state);
-            break;
-        case ugrpc::impl::AsyncMethodInvocation::WaitStatus::kError:
-            ProcessNetworkError(state, "Finish");
-            break;
-        case ugrpc::impl::AsyncMethodInvocation::WaitStatus::kCancelled:
-        case ugrpc::impl::AsyncMethodInvocation::WaitStatus::kDeadline:
-            UINVARIANT(false, "unreachable");
+    if (ok) {
+        ProcessFinishAbandoned(state);
+    } else {
+        ProcessNetworkError(state, "Finish");
     }
 } catch (const std::exception& ex) {
     LOG_WARNING() << "There is a caught exception in 'FinishAbandoned': " << ex;
@@ -133,9 +128,9 @@ void FinishAbandoned(GrpcStream& stream, StreamingCallState& state) noexcept try
 template <typename GrpcStream, typename Response>
 [[nodiscard]] bool Read(GrpcStream& stream, Response& response, StreamingCallState& state) {
     UINVARIANT(IsReadAvailable(state), "'impl::Read' called on a finished call");
-    ugrpc::impl::AsyncMethodInvocation read;
-    stream.Read(&response, read.GetCompletionTag());
-    const auto wait_status = WaitAndTryCancelIfNeeded(read, state.GetClientContext());
+    ugrpc::impl::AsyncMethodInvocation invocation;
+    stream.Read(&response, invocation.GetCompletionTag());
+    const auto wait_status = WaitAndTryCancelIfNeeded(invocation, state.GetClientContext());
     if (wait_status == ugrpc::impl::AsyncMethodInvocation::WaitStatus::kCancelled) {
         state.GetStatsScope().OnCancelled();
     }
@@ -146,8 +141,8 @@ template <typename GrpcStream, typename Response>
 void ReadAsync(GrpcStream& stream, Response& response, StreamingCallState& state) {
     UINVARIANT(IsReadAvailable(state), "'impl::Read' called on a finished call");
     state.EmplaceAsyncMethodInvocation();
-    auto& read = state.GetAsyncMethodInvocation();
-    stream.Read(&response, read.GetCompletionTag());
+    auto& invocation = state.GetAsyncMethodInvocation();
+    stream.Read(&response, invocation.GetCompletionTag());
 }
 
 template <typename GrpcStream, typename Request>
@@ -160,16 +155,16 @@ bool Write(GrpcStream& stream, const Request& request, grpc::WriteOptions option
 
     UINVARIANT(IsWriteAvailable(state), "'impl::Write' called on a stream that is closed for writes");
 
-    ugrpc::impl::AsyncMethodInvocation write;
-    stream.Write(request, options, write.GetCompletionTag());
-    const auto result = WaitAndTryCancelIfNeeded(write, state.GetClientContext());
-    if (result == ugrpc::impl::AsyncMethodInvocation::WaitStatus::kCancelled) {
+    ugrpc::impl::AsyncMethodInvocation invocation;
+    stream.Write(request, options, invocation.GetCompletionTag());
+    const auto wait_status = WaitAndTryCancelIfNeeded(invocation, state.GetClientContext());
+    if (wait_status == ugrpc::impl::AsyncMethodInvocation::WaitStatus::kCancelled) {
         state.GetStatsScope().OnCancelled();
     }
-    if (result != ugrpc::impl::AsyncMethodInvocation::WaitStatus::kOk) {
+    if (wait_status != ugrpc::impl::AsyncMethodInvocation::WaitStatus::kOk) {
         state.SetWritesFinished();
     }
-    return result == ugrpc::impl::AsyncMethodInvocation::WaitStatus::kOk;
+    return wait_status == ugrpc::impl::AsyncMethodInvocation::WaitStatus::kOk;
 }
 
 template <typename GrpcStream, typename Request>
@@ -182,9 +177,9 @@ void WriteAndCheck(GrpcStream& stream, const Request& request, grpc::WriteOption
 
     UINVARIANT(IsWriteAndCheckAvailable(state), "'impl::WriteAndCheck' called on a finished or closed stream");
 
-    ugrpc::impl::AsyncMethodInvocation write;
-    stream.Write(request, options, write.GetCompletionTag());
-    CheckOk(state, WaitAndTryCancelIfNeeded(write, state.GetClientContext()), "WriteAndCheck");
+    ugrpc::impl::AsyncMethodInvocation invocation;
+    stream.Write(request, options, invocation.GetCompletionTag());
+    CheckOk(state, WaitAndTryCancelIfNeeded(invocation, state.GetClientContext()), "WriteAndCheck");
 }
 
 template <typename GrpcStream>
@@ -197,9 +192,9 @@ bool WritesDone(GrpcStream& stream, StreamingCallState& state) {
 
     UINVARIANT(IsWriteAvailable(state), "'impl::WritesDone' called on a stream that is closed for writes");
     state.SetWritesFinished();
-    ugrpc::impl::AsyncMethodInvocation writes_done;
-    stream.WritesDone(writes_done.GetCompletionTag());
-    const auto wait_status = WaitAndTryCancelIfNeeded(writes_done, state.GetClientContext());
+    ugrpc::impl::AsyncMethodInvocation invocation;
+    stream.WritesDone(invocation.GetCompletionTag());
+    const auto wait_status = WaitAndTryCancelIfNeeded(invocation, state.GetClientContext());
     if (wait_status == ugrpc::impl::AsyncMethodInvocation::WaitStatus::kCancelled) {
         state.GetStatsScope().OnCancelled();
     }

grpc/include/userver/ugrpc/client/impl/prepare_call.hpp

@@ -67,11 +67,11 @@ class PrepareUnaryCallProxy<grpc::GenericStub, grpc::ByteBuffer, grpc::ByteBuffe
 
     decltype(auto) operator()(
         StubHandle& stub_handle,
-        grpc::ClientContext* context,
+        grpc::ClientContext* client_context,
         const grpc::ByteBuffer& request,
         grpc::CompletionQueue* cq
     ) const {
-        return impl::PrepareCall(prepare_async_method_, stub_handle, context, method_name_, request, cq);
+        return impl::PrepareCall(prepare_async_method_, stub_handle, client_context, method_name_, request, cq);
     }
 
 private:

grpc/include/userver/ugrpc/client/stream_read_future.hpp

@@ -97,11 +97,12 @@ bool StreamReadFuture<RawStream>::Get() {
     UINVARIANT(state_, "'Get' must be called only once");
     const impl::StreamingCallState::AsyncMethodInvocationGuard guard(*state_);
     auto* const state = std::exchange(state_, nullptr);
-    const auto result = impl::WaitAndTryCancelIfNeeded(state->GetAsyncMethodInvocation(), state->GetClientContext());
-    if (result == ugrpc::impl::AsyncMethodInvocation::WaitStatus::kCancelled) {
+    const auto
+        wait_status = impl::WaitAndTryCancelIfNeeded(state->GetAsyncMethodInvocation(), state->GetClientContext());
+    if (wait_status == ugrpc::impl::AsyncMethodInvocation::WaitStatus::kCancelled) {
         state->GetStatsScope().OnCancelled();
         state->GetStatsScope().Flush();
-    } else if (result == ugrpc::impl::AsyncMethodInvocation::WaitStatus::kError) {
+    } else if (wait_status == ugrpc::impl::AsyncMethodInvocation::WaitStatus::kError) {
         // Finish can only be called once all the data is read, otherwise the
         // underlying gRPC driver hangs.
         impl::Finish(*stream_, *state, /*final_response=*/nullptr, /*throw_on_error=*/true);
@@ -110,7 +111,7 @@ bool StreamReadFuture<RawStream>::Get() {
             RunMiddlewarePipeline(*state, impl::MiddlewareHooks::RecvMessageHooks(*recv_message_));
         }
     }
-    return result == ugrpc::impl::AsyncMethodInvocation::WaitStatus::kOk;
+    return wait_status == ugrpc::impl::AsyncMethodInvocation::WaitStatus::kOk;
 }
 
 template <typename RawStream>

grpc/include/userver/ugrpc/impl/async_method_invocation.hpp

@@ -20,8 +20,6 @@ class AsyncMethodInvocation : public EventBase {
     /// @see EventBase::Notify
     void Notify(bool ok) noexcept override;
 
-    bool IsBusy() const noexcept;
-
     enum class WaitStatus {
         kOk,
         kError,
@@ -52,12 +50,10 @@ class AsyncMethodInvocation : public EventBase {
     // For internal use only.
     engine::impl::ContextAccessor* TryGetContextAccessor() noexcept;
     /// @endcond
-protected:
-    void WaitWhileBusy() noexcept;
 
 private:
+    bool enqueued_{false};
     bool ok_{false};
-    bool busy_{false};
     engine::SingleUseEvent event_;
 };
 

grpc/include/userver/ugrpc/server/impl/service_worker_impl.hpp

@@ -107,6 +107,7 @@ class CallData final {
 
         auto& queue = method_data_.service_data.internals.completion_queues.GetQueue(method_data_.queue_id);
 
+        ugrpc::impl::AsyncMethodInvocation request_call_invocation;
         // the request for an incoming RPC must be performed synchronously
         method_data_.service_data.async_service.template RequestCall<CallTraits>(
             method_data_.method_id,
@@ -115,14 +116,14 @@ class CallData final {
             raw_responder_,
             queue,
             queue,
-            request_call_.GetCompletionTag()
+            request_call_invocation.GetCompletionTag()
         );
 
         // Note: we ignore task cancellations here. Even if notify_when_done has
         // already cancelled this RPC, we want to:
         // 1. listen to further RPCs for the same method
         // 2. handle this RPC correctly, including metrics, logs, etc.
-        if (!request_call_.WaitNonCancellable()) {
+        if (!request_call_invocation.WaitNonCancellable()) {
             // the CompletionQueue is shutting down
 
             // Do not wait for notify_when_done. When queue is shutting down, it will
@@ -197,7 +198,6 @@ class CallData final {
     typename CallTraits::RawContext context_{};
     InitialRequest initial_request_{};
     RawResponder raw_responder_{&context_};
-    ugrpc::impl::AsyncMethodInvocation request_call_;
     std::optional<tracing::InPlaceSpan> span_storage_{};
 };
 

grpc/src/ugrpc/client/channels.cpp

@@ -36,9 +36,9 @@ namespace {
             return false;
         }
 
-        ugrpc::impl::AsyncMethodInvocation operation;
-        channel.NotifyOnStateChange(state, deadline, &queue, operation.GetCompletionTag());
-        if (operation.Wait() != ugrpc::impl::AsyncMethodInvocation::WaitStatus::kOk) {
+        ugrpc::impl::AsyncMethodInvocation invocation;
+        channel.NotifyOnStateChange(state, deadline, &queue, invocation.GetCompletionTag());
+        if (invocation.Wait() != ugrpc::impl::AsyncMethodInvocation::WaitStatus::kOk) {
             return false;
         }
     }

grpc/src/ugrpc/client/impl/async_stream_methods.cpp

@@ -26,8 +26,8 @@ void SetErrorAndResetSpan(CallState& state, std::string_view error_message) noex
 
 }  // namespace
 
-void ThrowIfDeadlineIsExceeded(grpc::ClientContext& context, std::string_view call_name) {
-    const auto raw_deadline = context.raw_deadline();
+void ThrowIfDeadlineIsExceeded(grpc::ClientContext& client_context, std::string_view call_name) {
+    const auto raw_deadline = client_context.raw_deadline();
     const auto deadline = ugrpc::TimespecToDeadline(raw_deadline);
     if (deadline.IsReached()) {
         grpc::Status deadline_status(grpc::StatusCode::DEADLINE_EXCEEDED, "Deadline exceeded");
@@ -37,22 +37,26 @@ void ThrowIfDeadlineIsExceeded(grpc::ClientContext& context, std::string_view ca
 
 ugrpc::impl::AsyncMethodInvocation::WaitStatus WaitAndTryCancelIfNeeded(
     ugrpc::impl::AsyncMethodInvocation& invocation,
-    grpc::ClientContext& context
+    grpc::ClientContext& client_context
 ) noexcept {
     const auto wait_status = invocation.Wait();
     if (ugrpc::impl::AsyncMethodInvocation::WaitStatus::kCancelled == wait_status) {
-        context.TryCancel();
+        client_context.TryCancel();
     }
     return wait_status;
 }
 
-void CheckOk(StreamingCallState& state, ugrpc::impl::AsyncMethodInvocation::WaitStatus status, std::string_view stage) {
-    if (status == ugrpc::impl::AsyncMethodInvocation::WaitStatus::kError) {
+void CheckOk(
+    StreamingCallState& state,
+    ugrpc::impl::AsyncMethodInvocation::WaitStatus wait_status,
+    std::string_view stage
+) {
+    if (wait_status == ugrpc::impl::AsyncMethodInvocation::WaitStatus::kError) {
         state.SetFinished();
         ThrowIfDeadlineIsExceeded(state.GetClientContext(), state.GetCallName());
         ProcessNetworkError(state, stage);
         throw RpcInterruptedError(state.GetCallName(), stage);
-    } else if (status == ugrpc::impl::AsyncMethodInvocation::WaitStatus::kCancelled) {
+    } else if (wait_status == ugrpc::impl::AsyncMethodInvocation::WaitStatus::kCancelled) {
         state.SetFinished();
         ProcessCancelled(state, stage);
         throw RpcCancelledError(state.GetCallName(), stage);

grpc/src/ugrpc/impl/async_method_invocation.cpp

@@ -8,18 +8,20 @@ USERVER_NAMESPACE_BEGIN
 
 namespace ugrpc::impl {
 
-AsyncMethodInvocation::~AsyncMethodInvocation() { WaitWhileBusy(); }
+AsyncMethodInvocation::~AsyncMethodInvocation() {
+    if (enqueued_) {
+        event_.WaitNonCancellable();
+    }
+}
 
 void AsyncMethodInvocation::Notify(bool ok) noexcept {
     ok_ = ok;
     event_.Send();
 }
 
-bool AsyncMethodInvocation::IsBusy() const noexcept { return busy_; }
-
 void* AsyncMethodInvocation::GetCompletionTag() noexcept {
-    UASSERT(!busy_);
-    busy_ = true;
+    UASSERT(!enqueued_);
+    enqueued_ = true;
     return static_cast<EventBase*>(this);
 }
 
@@ -40,7 +42,6 @@ AsyncMethodInvocation::WaitStatus AsyncMethodInvocation::WaitUntil(engine::Deadl
             return WaitStatus::kDeadline;
         }
         case engine::FutureStatus::kReady: {
-            busy_ = false;
             return ok_ ? WaitStatus::kOk : WaitStatus::kError;
         }
     }
@@ -59,13 +60,6 @@ engine::impl::ContextAccessor* AsyncMethodInvocation::TryGetContextAccessor() no
 
 bool AsyncMethodInvocation::IsReady() const noexcept { return event_.IsReady(); }
 
-void AsyncMethodInvocation::WaitWhileBusy() noexcept {
-    if (busy_) {
-        event_.WaitNonCancellable();
-    }
-    busy_ = false;
-}
-
 }  // namespace ugrpc::impl
 
 USERVER_NAMESPACE_END

grpc/tests/base_test.cpp

@@ -115,8 +115,8 @@ ugrpc::client::CallOptions PrepareCallOptions() {
     return call_options;
 }
 
-void CheckClientContext(const grpc::ClientContext& context) {
-    const auto& metadata = context.GetServerTrailingMetadata();
+void CheckClientContext(const grpc::ClientContext& client_context) {
+    const auto& metadata = client_context.GetServerTrailingMetadata();
     const auto iter = metadata.find("resp_header");
     ASSERT_NE(iter, metadata.end());
     EXPECT_EQ(iter->second, "value");

grpc/tests/tracing_test.cpp

@@ -77,8 +77,8 @@ class GrpcTracing : public ugrpc::tests::ServiceFixture<UnitTestServiceWithTraci
     logging::DefaultLoggerLevelScope log_level_scope_{logging::Level::kInfo};
 };
 
-void CheckMetadata(const grpc::ClientContext& context) {
-    const auto& metadata = context.GetServerInitialMetadata();
+void CheckMetadata(const grpc::ClientContext& client_context) {
+    const auto& metadata = client_context.GetServerInitialMetadata();
     const auto& span = tracing::Span::CurrentSpan();
 
     // - TraceId should propagate both to sub-spans within a single service,