feat: allow freeing of ustring table for debugging cases

src/include/OpenImageIO/imageio.h

@@ -3953,6 +3953,15 @@ OIIO_API std::string geterror(bool clear = true);
 ///   enable globally in an environment where security is a higher priority
 ///   than being tolerant of partially broken image files.
 ///
+/// - `ustring:cleanup` (int: 0)
+///
+///    If nonzero, upon exit, do a thorough (and possibly expensive) teardown
+///    of ustring internal resources to ensure that there are no apparent
+///    memory leaks. This is only desirable in certain debugging situations.
+///    Ordinarily, it is better to finish as quickly as possible, so the
+///    default of 0 skips a time consuming and pointless teardown of the
+///    ustring internal allocations when the app exits.
+///
 /// EXAMPLES:
 /// ```
 ///     // Setting single simple values simply:

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,16 +2,37 @@
 // 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>
 
+
+
+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"));
+
+}  // namespace pvt
+OIIO_NAMESPACE_END
+
+
 OIIO_NAMESPACE_3_1_BEGIN
 
 // Use rw spin locks
@@ -37,16 +58,28 @@ 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 */
+    {
+        if (OIIO::pvt::oiio_ustring_cleanup) {
+            // If requested, take the time to properly destroy all the
+            // entries, and also the unique_ptr arrays all_pools and
+            // large_allocs will naturally free their contents after this
+            // destructor body ends.
+            destroy_entries();
+        } else {
+            // If no ustring cleanup was requested, take the fastest possible
+            // route to exit, just release the pointers and let them leak!
+            for (auto& p : all_pools)
+                (void)p.release();
+            for (auto& p : large_allocs)
+                (void)p.release();
+        }
     }
 
     size_t get_memory_usage()
@@ -146,31 +179,36 @@ template<unsigned BASE_CAPACITY, unsigned POOL_SIZE> struct TableRepMap {
 private:
     void grow()
     {
-        size_t new_mask = mask * 2 + 1;
-
-        // NOTE: only increment by half because we are doubling the entries and freeing the old
-        memory_usage += (mask + 1) * sizeof(ustring::TableRep*);
-
-        ustring::TableRep** new_entries = static_cast<ustring::TableRep**>(
-            calloc(new_mask + 1, sizeof(ustring::TableRep*)));
-        size_t to_copy = num_entries;
+        // 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;
+
+        // Make bigger space for new entries table and new mask
+        resize_entries(2 * old_entries.size());
+        num_entries = old_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 +217,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 +228,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