Fix the lifetime of InferPayload

src/infer_payload.cc

@@ -33,28 +33,28 @@ InferPayload::InferPayload(
     std::function<void(std::unique_ptr<InferResponse>)> callback)
     : is_decoupled_(is_decoupled), is_promise_set_(false), callback_(callback)
 {
-  prev_promise_.reset(new std::promise<std::unique_ptr<InferResponse>>());
-}
-
-InferPayload::~InferPayload()
-{
-  prev_promise_.reset();
+  promise_.reset(new std::promise<std::unique_ptr<InferResponse>>());
 }
 
 void
-InferPayload::SetValueForPrevPromise(
-    std::unique_ptr<InferResponse> infer_response)
+InferPayload::SetValue(std::unique_ptr<InferResponse> infer_response)
 {
-  prev_promise_->set_value(std::move(infer_response));
-  prev_promise_.reset();
-  is_promise_set_ = true;
+  // Only set value to the promise with the first response. Call the callback
+  // function to send decoupled response to the stub.
+  std::lock_guard<std::mutex> lock(mutex_);
+  if (!is_promise_set_) {
+    is_promise_set_ = true;
+    promise_->set_value(std::move(infer_response));
+  } else {
+    Callback(std::move(infer_response));
+  }
 }
 
 void
 InferPayload::SetFuture(
     std::future<std::unique_ptr<InferResponse>>& response_future)
 {
-  response_future = prev_promise_->get_future();
+  response_future = promise_->get_future();
 }
 
 bool

src/infer_payload.h

@@ -48,9 +48,8 @@ class InferPayload {
   InferPayload(
       const bool is_decouple,
       std::function<void(std::unique_ptr<InferResponse>)> callback);
-  ~InferPayload();
 
-  void SetValueForPrevPromise(std::unique_ptr<InferResponse> infer_response);
+  void SetValue(std::unique_ptr<InferResponse> infer_response);
   void SetFuture(std::future<std::unique_ptr<InferResponse>>& response_future);
   bool IsDecoupled();
   bool IsPromiseSet();
@@ -60,8 +59,9 @@ class InferPayload {
   std::shared_ptr<ResponseAllocatorUserp> ResponseAllocUserp();
 
  private:
-  std::unique_ptr<std::promise<std::unique_ptr<InferResponse>>> prev_promise_;
+  std::unique_ptr<std::promise<std::unique_ptr<InferResponse>>> promise_;
   bool is_decoupled_;
+  std::mutex mutex_;
   bool is_promise_set_;
   std::function<void(std::unique_ptr<InferResponse>)> callback_;
   std::shared_ptr<ResponseAllocatorUserp> response_alloc_userp_;

src/request_executor.cc

@@ -177,13 +177,7 @@ InferResponseComplete(
         output_tensors, pb_error, true /* is_last_response */, userp /* id */);
   }
 
-  // Only set value to the promise with the first response. Call the callback
-  // function to send decoupled response to the stub.
-  if (p->IsPromiseSet()) {
-    p->Callback(std::move(infer_response));
-  } else {
-    p->SetValueForPrevPromise(std::move(infer_response));
-  }
+  p->SetValue(std::move(infer_response));
 }
 
 TRITONSERVER_Error*