Fix unregister issue
This commit is contained in:
@@ -193,8 +193,12 @@ static bool CheckGPUVRAM(int gpuIndex, size_t minFreeBytes = 512ULL * 1024 * 102
|
||||
return true;
|
||||
}
|
||||
|
||||
static std::unordered_map<ANSCENTER::ANSODBase*, int>& ODHandleRegistry() {
|
||||
static std::unordered_map<ANSCENTER::ANSODBase*, int> s;
|
||||
// destructionStarted: set by the first Unregister caller; blocks new Acquires
|
||||
// and makes subsequent Unregister calls return false without deleting.
|
||||
// Prevents double-free when Release is raced on the same handle.
|
||||
struct ODEntry { int refcount; bool destructionStarted; };
|
||||
static std::unordered_map<ANSCENTER::ANSODBase*, ODEntry>& ODHandleRegistry() {
|
||||
static std::unordered_map<ANSCENTER::ANSODBase*, ODEntry> s;
|
||||
return s;
|
||||
}
|
||||
static std::mutex& ODHandleRegistryMutex() {
|
||||
@@ -208,46 +212,48 @@ static std::condition_variable& ODHandleRegistryCV() {
|
||||
|
||||
static void RegisterODHandle(ANSCENTER::ANSODBase* h) {
|
||||
std::lock_guard<std::mutex> lk(ODHandleRegistryMutex());
|
||||
ODHandleRegistry()[h] = 1; // refcount = 1
|
||||
ODHandleRegistry()[h] = { 1, false };
|
||||
}
|
||||
|
||||
// Acquire a handle for use (increment refcount). Returns the handle
|
||||
// if valid, nullptr if already released.
|
||||
// if valid, nullptr if already released or being destroyed.
|
||||
static ANSCENTER::ANSODBase* AcquireODHandle(ANSCENTER::ANSODBase* h) {
|
||||
std::lock_guard<std::mutex> lk(ODHandleRegistryMutex());
|
||||
auto it = ODHandleRegistry().find(h);
|
||||
if (it == ODHandleRegistry().end()) return nullptr;
|
||||
it->second++;
|
||||
if (it->second.destructionStarted) return nullptr;
|
||||
it->second.refcount++;
|
||||
return h;
|
||||
}
|
||||
|
||||
// Release a use of the handle (decrement refcount).
|
||||
// Returns true if this was the last reference (caller should destroy).
|
||||
// Release a use of the handle (decrement refcount). Only signals the CV;
|
||||
// the object is always deleted by Unregister, never by a stray ref-drop.
|
||||
static bool ReleaseODHandleRef(ANSCENTER::ANSODBase* h) {
|
||||
std::lock_guard<std::mutex> lk(ODHandleRegistryMutex());
|
||||
auto it = ODHandleRegistry().find(h);
|
||||
if (it == ODHandleRegistry().end()) return false;
|
||||
it->second--;
|
||||
if (it->second <= 0) {
|
||||
ODHandleRegistry().erase(it);
|
||||
it->second.refcount--;
|
||||
if (it->second.refcount <= 0) {
|
||||
ODHandleRegistryCV().notify_all();
|
||||
return true; // Caller should destroy
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
// Unregister and wait for all in-flight uses to finish.
|
||||
// Decrements the creation refcount and blocks until refcount hits 0.
|
||||
// Returns true if caller should destroy the object.
|
||||
// First caller takes ownership of destruction; subsequent calls return false.
|
||||
static bool UnregisterODHandle(ANSCENTER::ANSODBase* h) {
|
||||
std::unique_lock<std::mutex> lk(ODHandleRegistryMutex());
|
||||
auto it = ODHandleRegistry().find(h);
|
||||
if (it == ODHandleRegistry().end()) return false;
|
||||
it->second--; // Remove creation ref
|
||||
if (it->second.destructionStarted) {
|
||||
return false; // Another thread already owns the delete.
|
||||
}
|
||||
it->second.destructionStarted = true;
|
||||
it->second.refcount--; // Remove creation ref
|
||||
// Wait for in-flight inferences to finish (30s timeout as safety net)
|
||||
bool ok = ODHandleRegistryCV().wait_for(lk, std::chrono::seconds(30), [&]() {
|
||||
auto it2 = ODHandleRegistry().find(h);
|
||||
return it2 == ODHandleRegistry().end() || it2->second <= 0;
|
||||
return it2 == ODHandleRegistry().end() || it2->second.refcount <= 0;
|
||||
});
|
||||
if (!ok) {
|
||||
OutputDebugStringA("WARNING: UnregisterODHandle timed out waiting for in-flight inference\n");
|
||||
|
||||
Reference in New Issue
Block a user