Refactor code after copilot review

Signed-off-by: Bogdan Pereanu <bogdan.pereanu@intel.com>

Author: pereanub

src/inference/src/dev/iasync_infer_request.cpp

@@ -66,7 +66,7 @@ ov::IAsyncInferRequest::IAsyncInferRequest(const std::shared_ptr<IInferRequest>&
 void ov::IAsyncInferRequest::wait() {
     // Just use the last '_futures' member to wait pipeline completion
     auto future = [this] {
-        std::lock_guard<std::mutex> lock{m_mutex};
+        std::lock_guard lock{m_mutex};
         return m_futures.empty() ? std::shared_future<void>{} : m_futures.back();
     }();
     if (future.valid()) {
@@ -79,7 +79,7 @@ bool ov::IAsyncInferRequest::wait_for(const std::chrono::milliseconds& timeout)
 
     // Just use the last '_futures' member to wait pipeline completion
     auto future = [this] {
-        std::lock_guard<std::mutex> lock{m_mutex};
+        std::lock_guard lock{m_mutex};
         return m_futures.empty() ? std::shared_future<void>{} : m_futures.back();
     }();
 
@@ -98,15 +98,15 @@ bool ov::IAsyncInferRequest::wait_for(const std::chrono::milliseconds& timeout)
 }
 
 void ov::IAsyncInferRequest::cancel() {
-    std::lock_guard<std::mutex> lock{m_mutex};
+    std::lock_guard lock{m_mutex};
     if (m_state == InferState::BUSY) {
         m_state = InferState::CANCELLED;
     }
 }
 
 void ov::IAsyncInferRequest::set_callback(std::function<void(std::exception_ptr)> callback) {
     check_state();
-    std::lock_guard callbackLock{m_callback_invoke_mutex};
+    std::lock_guard lock{m_mutex};
     m_callback = std::move(callback);
 }
 
@@ -160,20 +160,27 @@ ov::threading::Task ov::IAsyncInferRequest::make_next_stage_task(
             if ((itEndStage == itNextStage) || (nullptr != currentException)) {
                 auto lastStageTask = [this, currentException]() mutable {
                     std::promise<void> promise;
+                    std::function<void(std::exception_ptr)> callback;
                     {
-                        std::lock_guard<std::mutex> lock{m_mutex};
-                        m_state = InferState::IDLE;
-                        promise = std::move(m_promise);
-                    }
-
-                    if (std::lock_guard callbackLock{m_callback_invoke_mutex}; m_callback) {
-                        try {
-                            m_callback(currentException);
-                        } catch (...) {
-                            currentException = std::current_exception();
+                        std::lock_guard callbackLock{m_callback_invoke_mutex};
+                        {
+                            std::lock_guard lock{m_mutex};
+                            m_state = InferState::IDLE;
+                            promise = std::move(m_promise);
+                            std::swap(callback, m_callback);
+                        }
+                        if (callback) {
+                            try {
+                                callback(currentException);
+                            } catch (...) {
+                                currentException = std::current_exception();
+                            }
+                            std::lock_guard lock{m_mutex};
+                            if (!m_callback) {
+                                std::swap(callback, m_callback);
+                            }
                         }
                     }
-
                     if (nullptr == currentException) {
                         promise.set_value();
                     } else {
@@ -248,10 +255,7 @@ void ov::IAsyncInferRequest::stop_and_wait() {
         std::lock_guard<std::mutex> lock{m_mutex};
         state = m_state;
         if (state != InferState::STOP) {
-            {
-                std::lock_guard callbackLock{m_callback_invoke_mutex};
-                m_callback = {};
-            }
+            m_callback = {};
             m_state = InferState::STOP;
             futures = std::move(m_futures);
         }