From 0342f51cd09ac26acdb71a4107b1dd998d91c541 Mon Sep 17 00:00:00 2001
From: Krzysztof Bruniecki <krzysztof.bruniecki@intel.com>
Date: Thu, 19 Nov 2020 16:08:04 +0100
Subject: [PATCH] Remove from unwaited requests ids if request aborted (#3160)

---
 inference-engine/src/gna_plugin/gna_device.cpp | 10 +++++-----
 inference-engine/src/gna_plugin/gna_device.hpp |  3 ++-
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/inference-engine/src/gna_plugin/gna_device.cpp b/inference-engine/src/gna_plugin/gna_device.cpp
index c4e72cca57f..809388a66c3 100644
--- a/inference-engine/src/gna_plugin/gna_device.cpp
+++ b/inference-engine/src/gna_plugin/gna_device.cpp
@@ -86,7 +86,7 @@ uint32_t GNADeviceHelper::propagate(const uint32_t requestConfigId, Gna2Accelera
     const auto status2 = Gna2RequestEnqueue(requestConfigId, &reqId);
     checkGna2Status(status2, "Gna2RequestEnqueue");
 
-    unwaitedRequestIds.push_back(reqId);
+    unwaitedRequestIds.insert(reqId);
 
     return reqId;
 }
@@ -329,14 +329,14 @@ const std::map <const std::pair<Gna2OperationType, int32_t>, const std::string>
 GnaWaitStatus GNADeviceHelper::wait(uint32_t reqId, int64_t millisTimeout) {
 #if GNA_LIB_VER == 2
     const auto status = Gna2RequestWait(reqId, millisTimeout);
-    if (status == Gna2StatusDriverQoSTimeoutExceeded) {
-        return GNA_REQUEST_ABORTED;
-    }
     if (status == Gna2StatusWarningDeviceBusy) {
         return GNA_REQUEST_PENDING;
     }
+    unwaitedRequestIds.erase(reqId);
+    if (status == Gna2StatusDriverQoSTimeoutExceeded) {
+        return GNA_REQUEST_ABORTED;
+    }
     checkGna2Status(status, "Gna2RequestWait");
-    unwaitedRequestIds.erase(std::remove(unwaitedRequestIds.begin(), unwaitedRequestIds.end(), reqId));
 #else
     if (isPerformanceMeasuring) {
         nGNAStatus = GNAWaitPerfRes(nGNAHandle, millisTimeout, reqId, &nGNAPerfResults);
diff --git a/inference-engine/src/gna_plugin/gna_device.hpp b/inference-engine/src/gna_plugin/gna_device.hpp
index e8c25a7c4f7..cb6abd19ee9 100644
--- a/inference-engine/src/gna_plugin/gna_device.hpp
+++ b/inference-engine/src/gna_plugin/gna_device.hpp
@@ -9,6 +9,7 @@
 #include <mutex>
 #include <string>
 #include <map>
+#include <set>
 #include <vector>
 #include <thread>
 
@@ -61,7 +62,7 @@ class GNADeviceHelper {
     uint64_t instrumentationResults[TotalGna2InstrumentationPoints] = {};
     uint64_t instrumentationTotal[TotalGna2InstrumentationPoints] = {};
     uint32_t instrumentationConfigId = 0;
-    std::vector<uint32_t> unwaitedRequestIds;
+    std::set<uint32_t> unwaitedRequestIds;
 #define MAX_TIMEOUT 500000
 #endif
     bool isPerformanceMeasuring = false;