From daa910ae4f529a32fcfcc6ffbe64352eb00c1488 Mon Sep 17 00:00:00 2001
From: Abhisar Sinha <63767682+abh1sar@users.noreply.github.com>
Date: Wed, 15 Apr 2026 23:47:53 +0530
Subject: [PATCH] wait for inflight request before de-registering image
 transfer

---
 scripts/vm/hypervisor/kvm/imageserver/config.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/scripts/vm/hypervisor/kvm/imageserver/config.py b/scripts/vm/hypervisor/kvm/imageserver/config.py
index 1c92fd12937..4bd41c2db56 100644
--- a/scripts/vm/hypervisor/kvm/imageserver/config.py
+++ b/scripts/vm/hypervisor/kvm/imageserver/config.py
@@ -104,6 +104,7 @@ class TransferRegistry:
 
     def __init__(self) -> None:
         self._lock = threading.Lock()
+        self._cv = threading.Condition(self._lock)
         self._transfers: Dict[str, Dict[str, Any]] = {}
         self._last_activity: Dict[str, float] = {}
         self._inflight: Dict[str, int] = {}
@@ -127,7 +128,9 @@ class TransferRegistry:
             logging.error("unregister rejected invalid transfer_id=%r", transfer_id)
             with self._lock:
                 return len(self._transfers)
-        with self._lock:
+        with self._cv:
+            while self._inflight.get(safe_id, 0) > 0:
+                self._cv.wait()
             self._transfers.pop(safe_id, None)
             self._last_activity.pop(safe_id, None)
             self._inflight.pop(safe_id, None)
@@ -168,12 +171,13 @@ class TransferRegistry:
             yield
         finally:
             now = time.monotonic()
-            with self._lock:
+            with self._cv:
                 count = self._inflight.get(safe_id, 1) - 1
                 if count <= 0:
                     self._inflight.pop(safe_id, None)
                     if safe_id in self._transfers:
                         self._last_activity[safe_id] = now
+                    self._cv.notify_all()
                 else:
                     self._inflight[safe_id] = count