fix(rocJPEG): keep a VA surface pool entry busy until all its surfaces are released

projects/rocjpeg/src/rocjpeg_vaapi_decoder.cpp

@@ -199,7 +199,9 @@ bool RocJpegVaapiMemoryPool::DeleteIdleEntry() {
  * idle entry exists, every entry is still in flight and the pool is allowed to grow past the target:
  * on a GPU with a single JPEG core the batched decoder submits one image per call, so each image in
  * the batch holds its own busy entry and a batch larger than the target legitimately needs more live
- * entries than the target allows. Failing here would reject valid work.
+ * entries than the target allows. Failing here would reject valid work. An entry also stays busy
+ * until every one of its surfaces is released, so a caller that syncs only some images of a batched
+ * decode before submitting again can likewise grow the pool past the target until the rest are released.
  *
  * @param surface_format The surface format for which the pool entry is being added.
  * @param pool_entry The pool entry to be added.
@@ -231,10 +233,11 @@ RocJpegVaapiMemPoolEntry RocJpegVaapiMemoryPool::GetEntry(uint32_t surface_forma
     for (auto& entry : mem_pool_[surface_format]) {
         if (entry.image_width >= image_width && entry.image_height >= image_height && entry.va_surface_ids.size() == num_surfaces && entry.entry_status == kIdle) {
             entry.entry_status = kBusy;
+            std::fill(entry.surface_status.begin(), entry.surface_status.end(), kBusy);
             return entry;
         }
     }
-    return {0, 0, kIdle, {}, {}};
+    return {0, 0, kIdle, {}, {}, {}};
 }
 
 bool RocJpegVaapiMemoryPool::FindSurfaceId(VASurfaceID surface_id) {
@@ -368,8 +371,17 @@ bool RocJpegVaapiMemoryPool::SetSurfaceAsIdle(VASurfaceID surface_id) {
     std::lock_guard<std::mutex> lock(pool_mutex_);
     for (auto& pair : mem_pool_) {
         for (auto& entry : pair.second) {
-            if (std::find(entry.va_surface_ids.begin(), entry.va_surface_ids.end(), surface_id) != entry.va_surface_ids.end()) {
-                entry.entry_status = kIdle;
+            auto it = std::find(entry.va_surface_ids.begin(), entry.va_surface_ids.end(), surface_id);
+            if (it != entry.va_surface_ids.end()) {
+                size_t idx = std::distance(entry.va_surface_ids.begin(), it);
+                if (idx >= entry.surface_status.size()) {
+                    ErrorLog(g_rocjpeg_logger, "The surface_status of the pool entry for surface_id: " + ROCJPEG_TOSTR(surface_id) + " does not match its surfaces!");
+                    return false;
+                }
+                entry.surface_status[idx] = kIdle;
+                if (std::all_of(entry.surface_status.begin(), entry.surface_status.end(), [](MemPoolEntryStatus status) { return status == kIdle; })) {
+                    entry.entry_status = kIdle;
+                }
                 return true;
             }
         }
@@ -901,6 +913,7 @@ RocJpegStatus RocJpegVappiDecoder::SubmitDecode(const JpegStreamParameters *jpeg
         mem_pool_entry.image_width = surface_width;
         mem_pool_entry.image_height = surface_height;
         mem_pool_entry.hip_interops.resize(1, HipInteropDeviceMem());
+        mem_pool_entry.surface_status.resize(1, kBusy);
         surface_id = mem_pool_entry.va_surface_ids[0];
         mem_pool_entry.entry_status = kBusy;
         CHECK_ROCJPEG(vaapi_mem_pool_->AddPoolEntry(surface_pixel_format, mem_pool_entry));
@@ -1041,6 +1054,7 @@ RocJpegStatus RocJpegVappiDecoder::SubmitDecodeBatched(JpegStreamParameters *jpe
                 surface_ids[indices[i]] = mem_pool_entry.va_surface_ids[i];
             }
             mem_pool_entry.hip_interops.resize(indices.size(), HipInteropDeviceMem());
+            mem_pool_entry.surface_status.resize(indices.size(), kBusy);
             mem_pool_entry.entry_status = kBusy;
             CHECK_ROCJPEG(vaapi_mem_pool_->AddPoolEntry(key.pixel_format, mem_pool_entry));
         } else {

projects/rocjpeg/src/rocjpeg_vaapi_decoder.h

@@ -107,14 +107,17 @@ typedef enum {
  *
  * This structure holds information about a memory pool entry used by the RocJpegVaapiDecoder.
  * It contains the image width and height, the entry status, an array of VA surface IDs,
- * and an array of HipInteropDeviceMem objects.
+ * an array of HipInteropDeviceMem objects, and the status of each VA surface.
  */
 struct RocJpegVaapiMemPoolEntry {
     uint32_t image_width;
     uint32_t image_height;
     MemPoolEntryStatus entry_status;
     std::vector<VASurfaceID> va_surface_ids;
     std::vector<HipInteropDeviceMem> hip_interops;
+    // Status of each surface in va_surface_ids. Surfaces are released one at a time, so the entry only
+    // becomes idle once all of them are idle; otherwise it could be reused while some are still in flight.
+    std::vector<MemPoolEntryStatus> surface_status;
 };
 
 /**
@@ -185,7 +188,8 @@ class RocJpegVaapiMemoryPool {
         /**
          * @brief Sets a VASurfaceID as idle.
          *
-         * This function sets the specified VASurfaceID as idle, indicating that it is available for reuse.
+         * This function sets the specified VASurfaceID as idle. Its pool entry becomes available for reuse
+         * once every surface of the entry is idle.
          *
          * @param surface_id The VASurfaceID to set as idle.
          * @return true if the VASurfaceID was successfully set as idle, false otherwise.