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>
@@ -39,16 +42,41 @@ template<unsigned BASE_CAPACITY, unsigned POOL_SIZE> struct TableRepMap {
     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
+        , memory_usage(sizeof(*this)
                        + sizeof(ustring::TableRep*) * BASE_CAPACITY)
     {
+        allocate_pool_block();
     }
 
     ~TableRepMap()
     { /* just let memory leak */
     }
 
+    void clear()
+    {
+        ustring_write_lock_t lock(mutex);
+        // Destroy all TableRep objects. The destructor safely handles the
+        // case where the internal std::string aliases the pool chars.
+        for (size_t i = 0; i <= mask; ++i) {
+            if (entries[i])
+                entries[i]->~TableRep();
+        }
+        // Free all pool allocations and large individual allocations.
+        all_pools.clear();
+        all_pools.shrink_to_fit();
+        large_allocs.clear();
+        large_allocs.shrink_to_fit();
+        free(entries);
+        // Re-initialize to a fresh, usable state.
+        mask    = BASE_CAPACITY - 1;
+        entries = static_cast<ustring::TableRep**>(
+            calloc(BASE_CAPACITY, sizeof(ustring::TableRep*)));
+        num_entries  = 0;
+        memory_usage = sizeof(*this)
+                       + sizeof(ustring::TableRep*) * BASE_CAPACITY;
+        allocate_pool_block();
+    }
+
     size_t get_memory_usage()
     {
         ustring_read_lock_t lock(mutex);
@@ -179,6 +207,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,26 +218,36 @@ 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;
     }
 
+    // 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);
+    }
+
     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;
+    size_t mask                 = BASE_CAPACITY - 1;
+    ustring::TableRep** entries = nullptr;
+    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
@@ -249,6 +290,12 @@ struct UstringTable {
         return num;
     }
 
+    void clear()
+    {
+        for (auto& bin : bins)
+            bin.clear();
+    }
+
 #    ifdef USTRING_TRACK_NUM_LOOKUPS
     size_t get_num_lookups()
     {
@@ -678,3 +725,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().clear();
+            v3_1::reverse_map().clear();
+        }
+    });
+    return 0;
+}();
+
+}  // namespace pvt
+OIIO_NAMESPACE_END