diff --git a/Source/TurboLinkGrpc/Private/TurboLinkGrpcContext.cpp b/Source/TurboLinkGrpc/Private/TurboLinkGrpcContext.cpp index 1dc9bad..1ff927c 100644 --- a/Source/TurboLinkGrpc/Private/TurboLinkGrpcContext.cpp +++ b/Source/TurboLinkGrpc/Private/TurboLinkGrpcContext.cpp @@ -28,12 +28,14 @@ void GrpcContext::UpdateState(EGrpcContextState NewState) InitialTag = Service->TurboLinkManager->GetNextTag(AsShared()); WriteTag = Service->TurboLinkManager->GetNextTag(AsShared()); ReadTag = Service->TurboLinkManager->GetNextTag(AsShared()); + FinishTag = Service->TurboLinkManager->GetNextTag(AsShared()); } else if (ContextState == EGrpcContextState::Done) { Service->TurboLinkManager->RemoveTag(InitialTag); Service->TurboLinkManager->RemoveTag(WriteTag); Service->TurboLinkManager->RemoveTag(ReadTag); + Service->TurboLinkManager->RemoveTag(FinishTag); } if (Client->OnContextStateChange.IsBound()) diff --git a/Source/TurboLinkGrpc/Private/TurboLinkGrpcContext.h b/Source/TurboLinkGrpc/Private/TurboLinkGrpcContext.h index 010fa8a..b2b7304 100644 --- a/Source/TurboLinkGrpc/Private/TurboLinkGrpcContext.h +++ b/Source/TurboLinkGrpc/Private/TurboLinkGrpcContext.h @@ -46,6 +46,7 @@ class GrpcContext : public TSharedFromThis void* InitialTag = nullptr; void* WriteTag = nullptr; void* ReadTag = nullptr; + void* FinishTag = nullptr; protected: void UpdateState(EGrpcContextState NewState); @@ -139,14 +140,22 @@ class GrpcContext_Ping_Stream : public TGrpcContext protected: void OnRpcEventInternal(bool Ok, const void* EventTag, typename Super::FRpcCallbackFunc RpcCallbackFunc) { + if (!Ok) { - Super::RpcReaderWriter->Finish(&(Super::RpcStatus), Super::ReadTag); - Super::UpdateState(EGrpcContextState::Done); + Super::RpcReaderWriter->Finish(&(Super::RpcStatus), Super::FinishTag); return; } FGrpcResult result = GrpcContext::MakeGrpcResult(Super::RpcStatus); + if (EventTag == Super::FinishTag) { + if (RpcCallbackFunc) + { + RpcCallbackFunc(result, nullptr); + } + Super::UpdateState(EGrpcContextState::Done); + return; + } if (Super::RpcStatus.ok()) { if (Super::GetState() == EGrpcContextState::Initialing) @@ -293,11 +302,20 @@ class GrpcContext_Stream_Stream : public TGrpcContext { if (!Ok) { - Super::RpcReaderWriter->Finish(&(Super::RpcStatus), Super::ReadTag); + Super::RpcReaderWriter->Finish(&(Super::RpcStatus), Super::FinishTag); return; } FGrpcResult result = GrpcContext::MakeGrpcResult(Super::RpcStatus); + if (EventTag == Super::FinishTag) + { + if (RpcCallbackFunc) + { + RpcCallbackFunc(result, &(Super::RpcResponse)); + } + Super::UpdateState(EGrpcContextState::Done); + return; + } if (Super::RpcStatus.ok()) { if (Super::GetState() == EGrpcContextState::Initialing)