Fix notify state destruction and inflight states tracking

src/grpc/infer_handler.cc

@@ -973,6 +973,8 @@ ModelInferHandler::InferResponseComplete(
     return;
   }
 
+  state->context_->EraseInflightState(state);
+
 #ifdef TRITON_ENABLE_TRACING
   state->trace_timestamps_.emplace_back(std::make_pair(
       "INFER_RESPONSE_COMPLETE", TraceManager::CaptureTimestamp()));
@@ -987,7 +989,6 @@ ModelInferHandler::InferResponseComplete(
         "deleting GRPC inference response");
 
     state->step_ = Steps::CANCELLED;
-    state->context_->EraseInflightState(state);
 
     LOG_VERBOSE(1) << "ModelInferHandler::InferResponseComplete, "
                    << state->unique_id_

src/grpc/infer_handler.h

@@ -627,22 +627,29 @@ class InferHandlerState {
         ::grpc::ServerCompletionQueue* cq, const uint64_t unique_id = 0)
         : cq_(cq), unique_id_(unique_id), ongoing_requests_(0),
           step_(Steps::START), finish_ok_(true), ongoing_write_(false),
-          received_notification_(false)
+          notify_state_(nullptr), received_notification_(false)
     {
       ctx_.reset(new ::grpc::ServerContext());
       responder_.reset(new ServerResponderType(ctx_.get()));
     }
 
+    ~Context()
+    {
+      if (notify_state_ != nullptr) {
+        delete notify_state_;
+      }
+    }
+
     void SetCompressionLevel(grpc_compression_level compression_level)
     {
       ctx_->set_compression_level(compression_level);
     }
 
     void GrpcContextAsyncNotifyWhenDone(InferHandlerStateType* state)
     {
-      InferHandlerStateType* wrapped_state =
+      notify_state_ =
           new InferHandlerStateType(Steps::WAITING_NOTIFICATION, state);
-      ctx_->AsyncNotifyWhenDone(wrapped_state);
+      ctx_->AsyncNotifyWhenDone(notify_state_);
     }
 
     void SetReceivedNotification(bool value) { received_notification_ = true; }
@@ -666,8 +673,12 @@ class InferHandlerState {
       all_states_.insert(state);
     }
 
-    // Adds the state object created on this context
-    void EraseState(InferHandlerStateType* state) { all_states_.erase(state); }
+    // Erases the state object created on this context
+    void EraseState(InferHandlerStateType* state)
+    {
+      EraseInflightState(state);
+      all_states_.erase(state);
+    }
 
     bool HandleCompletion()
     {
@@ -975,6 +986,10 @@ class InferHandlerState {
     // True if there is an ongoing write to the grpc stream
     std::atomic<bool> ongoing_write_;
 
+    // The state object that is sent to grpc async notification
+    // for tracking the gRPC stream.
+    InferHandlerState* notify_state_;
+
     // Tracks whether the async notification has been delivered by
     // completion queue.
     bool received_notification_;
@@ -1274,7 +1289,6 @@ InferHandler<
         state->context_->SetReceivedNotification(true);
         LOG_VERBOSE(1) << "Received notification for " << Name() << ", "
                        << state->unique_id_;
-        delete state_wrapper;
       }
       LOG_VERBOSE(2) << "Grpc::CQ::Next() "
                      << state->context_->DebugString(state);

src/grpc/stream_infer_handler.cc

@@ -576,6 +576,15 @@ ModelStreamInferHandler::StreamInferResponseComplete(
     }
   }
 
+  // If receiving the final callback then erase the state from the inflight
+  // state data structure to prevent cancellation being called on the request.
+  // Also make sure that if this state was sent to gRPC async notification
+  // mechanism then the state is not removed as it would be needed for handling
+  // the cancellation if detected.
+  if (state->complete_ && (!state->IsAsyncNotifyState())) {
+    state->context_->EraseInflightState(state);
+  }
+
   if (state->IsGrpcContextCancelled()) {
     std::lock_guard<std::recursive_mutex> lock(state->step_mtx_);
     // Clean-up the received response object.
@@ -593,7 +602,6 @@ ModelStreamInferHandler::StreamInferResponseComplete(
     // that state object can be released.
     if (state->complete_) {
       state->step_ = Steps::CANCELLED;
-      state->context_->EraseInflightState(state);
       state->context_->PutTaskBackToQueue(state);
     }
 
@@ -692,7 +700,6 @@ ModelStreamInferHandler::StreamInferResponseComplete(
       // that state object can be released.
       if (state->complete_) {
         state->step_ = Steps::CANCELLED;
-        state->context_->EraseInflightState(state);
         state->context_->PutTaskBackToQueue(state);
       }