feat: allow freeing of ustring table for debugging cases

src/include/imageio_pvt.h

@@ -53,7 +53,11 @@ extern atomic_ll IB_local_mem_current;
 extern atomic_ll IB_local_mem_peak;
 extern std::atomic<float> IB_total_open_time;
 extern std::atomic<float> IB_total_image_read_time;
-extern OIIO_UTIL_API int oiio_use_tbb;  // This lives in libOpenImageIO_Util
+
+// These live in libOpenImageIO_Util
+extern OIIO_UTIL_API int oiio_use_tbb;
+extern OIIO_UTIL_API int oiio_ustring_cleanup;
+
 OIIO_API const std::vector<std::string>&
 font_dirs();
 OIIO_API const std::vector<std::string>&

src/libOpenImageIO/imageio.cpp

@@ -373,6 +373,10 @@ attribute(string_view name, TypeDesc type, const void* val)
 
     // Things below here need to buarded by the attrib_mutex
     std::lock_guard lock(attrib_mutex);
+    if (name == "debug" && type == TypeInt) {
+        oiio_print_debug = *(const int*)val;
+        return true;
+    }
     if (name == "read_chunk" && type == TypeInt) {
         oiio_read_chunk = *(const int*)val;
         return true;
@@ -447,10 +451,6 @@ attribute(string_view name, TypeDesc type, const void* val)
         oiio_use_tbb = *(const int*)val;
         return true;
     }
-    if (name == "debug" && type == TypeInt) {
-        oiio_print_debug = *(const int*)val;
-        return true;
-    }
     if (name == "log_times" && type == TypeInt) {
         oiio_log_times = *(const int*)val;
         return true;
@@ -471,6 +471,10 @@ attribute(string_view name, TypeDesc type, const void* val)
         oiio_try_all_readers = *(const int*)val;
         return true;
     }
+    if (name == "ustring:cleanup" && type == TypeInt) {
+        oiio_ustring_cleanup = *(const int*)val;
+        return true;
+    }
 
     return false;
 }
@@ -511,6 +515,10 @@ getattribute(string_view name, TypeDesc type, void* val)
 
     // Things below here need to buarded by the attrib_mutex
     std::lock_guard lock(attrib_mutex);
+    if (name == "debug" && type == TypeInt) {
+        *(int*)val = oiio_print_debug;
+        return true;
+    }
     if (name == "read_chunk" && type == TypeInt) {
         *(int*)val = oiio_read_chunk;
         return true;
@@ -649,8 +657,8 @@ getattribute(string_view name, TypeDesc type, void* val)
         *(int*)val = oiio_use_tbb;
         return true;
     }
-    if (name == "debug" && type == TypeInt) {
-        *(int*)val = oiio_print_debug;
+    if (name == "ustring:cleanup" && type == TypeInt) {
+        *(int*)val = oiio_ustring_cleanup;
         return true;
     }
     if (name == "log_times" && type == TypeInt) {

src/libutil/ustring.cpp

@@ -2,12 +2,15 @@
 // SPDX-License-Identifier: Apache-2.0
 // https://github.com/AcademySoftwareFoundation/OpenImageIO
 
+#include <cstdlib>
 #include <string>
 #include <unordered_map>
+#include <vector>
 
 #include <OpenImageIO/dassert.h>
 #include <OpenImageIO/export.h>
 #include <OpenImageIO/strutil.h>
+#include <OpenImageIO/sysutil.h>
 #include <OpenImageIO/thread.h>
 #include <OpenImageIO/unordered_map_concurrent.h>
 #include <OpenImageIO/ustring.h>
@@ -37,18 +40,35 @@ template<unsigned BASE_CAPACITY, unsigned POOL_SIZE> struct TableRepMap {
                   "BASE_CAPACITY must be a power of 2");
 
     TableRepMap()
-        : entries(static_cast<ustring::TableRep**>(
-            calloc(BASE_CAPACITY, sizeof(ustring::TableRep*))))
-        , pool(static_cast<char*>(malloc(POOL_SIZE)))
-        , memory_usage(sizeof(*this) + POOL_SIZE
-                       + sizeof(ustring::TableRep*) * BASE_CAPACITY)
+        : memory_usage(sizeof(*this))
     {
+        resize_entries(BASE_CAPACITY);
+        allocate_pool_block();
     }
 
     ~TableRepMap()
     { /* just let memory leak */
     }
 
+    // Free all allocated resources. Note that this leaves the TableRepMap
+    // in an unusable state.
+    void free_resources()
+    {
+        ustring_write_lock_t lock(mutex);
+        // Destroy TableRep objects FIRST, while the pool/large-alloc memory
+        // they live in is still valid. Their destructors may read/write the
+        // surrounding allocation and (on some platforms) the std::string
+        // subobject owns a separate heap buffer that only gets freed here.
+        destroy_entries();
+        entries.clear();
+        entries.shrink_to_fit();
+        // Now safe to free the backing buffers.
+        all_pools.clear();
+        all_pools.shrink_to_fit();
+        large_allocs.clear();
+        large_allocs.shrink_to_fit();
+    }
+
     size_t get_memory_usage()
     {
         ustring_read_lock_t lock(mutex);
@@ -146,31 +166,35 @@ template<unsigned BASE_CAPACITY, unsigned POOL_SIZE> struct TableRepMap {
 private:
     void grow()
     {
-        size_t new_mask = mask * 2 + 1;
+        // Temporarily hold the old entries while we are copying them
+        std::vector<ustring::TableRep*> old_entries;
+        old_entries.swap(entries);
+        size_t old_num_entries = num_entries;
+        size_t to_copy         = old_num_entries;
 
-        // NOTE: only increment by half because we are doubling the entries and freeing the old
-        memory_usage += (mask + 1) * sizeof(ustring::TableRep*);
+        // Make bigger space for new entries table and new mask
+        resize_entries(2 * old_entries.size());
 
-        ustring::TableRep** new_entries = static_cast<ustring::TableRep**>(
-            calloc(new_mask + 1, sizeof(ustring::TableRep*)));
-        size_t to_copy = num_entries;
+        // Copy each entry from old into the new, recomputing the hash because
+        // the mask has changd.
         for (size_t i = 0; to_copy != 0; i++) {
-            if (entries[i] == 0)
+            if (old_entries[i] == nullptr)
                 continue;
-            size_t pos = entries[i]->hashed & new_mask, dist = 0;
+            // i is old position, pos will be new position
+            size_t pos = old_entries[i]->hashed & mask, dist = 0;
             for (;;) {
-                if (new_entries[pos] == 0)
+                if (entries[pos] == 0)
                     break;
                 ++dist;
-                pos = (pos + dist) & new_mask;  // quadratic probing
+                pos = (pos + dist) & mask;  // quadratic probing
             }
-            new_entries[pos] = entries[i];
+            entries[pos]   = old_entries[i];
+            old_entries[i] = nullptr;
             to_copy--;
         }
 
-        free(entries);
-        entries = new_entries;
-        mask    = new_mask;
+        // old_entries will free when we exit this function
+        memory_usage -= sizeof(ustring::TableRep*) * old_entries.size();
     }
 
     ustring::TableRep* make_rep(string_view str, uint64_t hash)
@@ -179,6 +203,9 @@ template<unsigned BASE_CAPACITY, unsigned POOL_SIZE> struct TableRepMap {
         return new (repmem) ustring::TableRep(str, hash);
     }
 
+    // Allocate `len` bytes from the pool. Allocate a new pool block if len
+    // doesn't fit in the current block. In the unlikely even that len > the
+    // pool block size, do a separate allocation just for it.
     char* pool_alloc(size_t len)
     {
         // round up to nearest multiple of pointer size to guarantee proper alignment of TableRep objects
@@ -187,29 +214,62 @@ template<unsigned BASE_CAPACITY, unsigned POOL_SIZE> struct TableRepMap {
 
         if (len >= POOL_SIZE) {
             memory_usage += len;
-            return (char*)malloc(len);  // no need to try and use the pool
+            char* p = new char[len];
+            large_allocs.emplace_back(p);
+            return p;
         }
         if (pool_offset + len > POOL_SIZE) {
-            // NOTE: old pool will leak - this is ok because ustrings cannot be freed
-            memory_usage += POOL_SIZE;
-            pool        = (char*)malloc(POOL_SIZE);
-            pool_offset = 0;
+            allocate_pool_block();
         }
         char* result = pool + pool_offset;
         pool_offset += len;
         return result;
     }
 
-    OIIO_CACHE_ALIGN mutable ustring_mutex_t mutex;
-    size_t mask = BASE_CAPACITY - 1;
-    ustring::TableRep** entries;
-    size_t num_entries = 0;
-    char* pool;
-    size_t pool_offset = 0;
-    size_t memory_usage;
+    // Allocate one more standard POOL_SIZE block for `pool`
+    void allocate_pool_block()
+    {
+        memory_usage += POOL_SIZE;
+        pool        = new char[POOL_SIZE];
+        pool_offset = 0;
+        all_pools.emplace_back(pool);
+    }
+
+    void destroy_entries()
+    {
+        // Destroy all TableRep objects. The destructor safely handles the
+        // case where the internal std::string aliases the pool chars.
+        for (auto& e : entries) {
+            if (e) {
+                e->~TableRep();
+                e = nullptr;
+            }
+        }
+    }
+
+    void resize_entries(size_t newsize)
+    {
+        OIIO_CONTRACT_ASSERT(entries.empty());
+        OIIO_CONTRACT_ASSERT_MSG((newsize & (newsize - 1)) == 0,
+                                 "New entries size must be power of 2");
+        entries.resize(newsize, nullptr);
+        memory_usage += sizeof(ustring::TableRep*) * entries.size();
+        num_entries = 0;
+        mask        = newsize - 1;
+    }
+
+    std::vector<ustring::TableRep*> entries;
+    size_t mask         = BASE_CAPACITY - 1;
+    size_t num_entries  = 0;
+    char* pool          = nullptr;  // Current pool block we're using
+    size_t pool_offset  = 0;        // Next offset within current block
+    size_t memory_usage = 0;        // Total memory usage
+    std::vector<std::unique_ptr<char[]>> all_pools;
+    std::vector<std::unique_ptr<char[]>> large_allocs;
 #ifdef USTRING_TRACK_NUM_LOOKUPS
     size_t num_lookups = 0;
 #endif
+    OIIO_CACHE_ALIGN mutable ustring_mutex_t mutex;
 };
 
 #if 0
@@ -249,6 +309,12 @@ struct UstringTable {
         return num;
     }
 
+    void free_resources()
+    {
+        for (auto& bin : bins)
+            bin.free_resources();
+    }
+
 #    ifdef USTRING_TRACK_NUM_LOOKUPS
     size_t get_num_lookups()
     {
@@ -678,3 +744,36 @@ ustring::memory()
 }
 
 OIIO_NAMESPACE_3_1_END
+
+
+OIIO_NAMESPACE_BEGIN
+namespace pvt {
+
+// If nonzero, the ustring table will be freed at process exit. This is off by
+// default because cleanup is unnecessary (the OS reclaims the memory) and can
+// add measurable time at exit for large tables. Enable it when using valgrind
+// or other leak detectors to suppress false positives. Settable via
+// OIIO::attribute("ustring:cleanup",1) or the OIIO_USTRING_CLEANUP
+// environment variable.
+OIIO_UTIL_API int oiio_ustring_cleanup = Strutil::stoi(
+    Sysutil::getenv("OIIO_USTRING_CLEANUP"));
+
+// Register an atexit handler once at startup. The handler checks the flag at
+// exit time, so it covers both the env-var path (flag set here) and the
+// OIIO::attribute() path (flag set later at runtime).
+static int ustring_cleanup_atexit_registered = []() {
+    std::atexit([]() {
+        if (pvt::oiio_ustring_cleanup) {
+#ifndef NDEBUG
+            OIIO::print("ustring: freeing table resources ({} bytes)\n",
+                        v3_1::ustring_table().get_memory_usage());
+#endif
+            v3_1::ustring_table().free_resources();
+            v3_1::reverse_map().clear();
+        }
+    });
+    return 0;
+}();
+
+}  // namespace pvt
+OIIO_NAMESPACE_END