Skip to content

Commit 9e15a18

Browse files
authored
Cherry-pick: Fix: Add mutex lock for state completion check in gRPC streaming to prevent race condition (#7619)
1 parent 29c7a28 commit 9e15a18

File tree

1 file changed

+12
-9
lines changed

1 file changed

+12
-9
lines changed

src/grpc/stream_infer_handler.cc

Lines changed: 12 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -537,15 +537,18 @@ ModelStreamInferHandler::Process(InferHandler::State* state, bool rpc_ok)
537537
} else if (state->step_ == Steps::WRITEREADY) {
538538
// Finish the state if all the transactions associated with
539539
// the state have completed.
540-
if (state->IsComplete()) {
541-
state->context_->DecrementRequestCounter();
542-
finished = Finish(state);
543-
} else {
544-
LOG_ERROR << "Should not print this! Decoupled should NOT write via "
545-
"WRITEREADY!";
546-
// Remove the state from the completion queue
547-
std::lock_guard<std::recursive_mutex> lock(state->step_mtx_);
548-
state->step_ = Steps::ISSUED;
540+
std::lock_guard<std::recursive_mutex> lk1(state->context_->mu_);
541+
{
542+
if (state->IsComplete()) {
543+
state->context_->DecrementRequestCounter();
544+
finished = Finish(state);
545+
} else {
546+
LOG_ERROR << "Should not print this! Decoupled should NOT write via "
547+
"WRITEREADY!";
548+
// Remove the state from the completion queue
549+
std::lock_guard<std::recursive_mutex> lock(state->step_mtx_);
550+
state->step_ = Steps::ISSUED;
551+
}
549552
}
550553
}
551554
}

0 commit comments

Comments
 (0)