Fix memory leak when copying ST tables

peterzhu2118 · peterzhu2118 · commit f0d08d11dcd4 · 2023-06-29T11:16:50.000-04:00
st_copy allocates a st_table, which is not needed for hashes since it is
allocated by VWA and embedded, so this causes a memory leak.

The following script demonstrates the issue:

```ruby
20.times do
  100_000.times do
    {a: 1, b: 2, c: 3, d: 4, e: 5, f: 6, g: 7, h: 8, i: 9}
  end

  puts `ps -o rss= -p #{$$}`
end
```
diff --git a/hash.c b/hash.c
@@ -1482,7 +1482,9 @@ hash_copy(VALUE ret, VALUE hash)
     else {
         HASH_ASSERT(sizeof(st_table) <= sizeof(ar_table));
 
-        RHASH_ST_TABLE_SET(ret, st_copy(RHASH_ST_TABLE(hash)));
+        RHASH_SET_ST_FLAG(ret);
+        st_replace(RHASH_ST_TABLE(ret), RHASH_ST_TABLE(hash));
+
         rb_gc_writebarrier_remember(ret);
     }
     return ret;
@@ -1776,7 +1778,8 @@ rb_hash_s_create(int argc, VALUE *argv, VALUE klass)
             }
             else {
                 hash = hash_alloc(klass);
-                hash_copy(hash, tmp);
+                if (!RHASH_EMPTY_P(tmp))
+                    hash_copy(hash, tmp);
                 return hash;
             }
         }
diff --git a/include/ruby/st.h b/include/ruby/st.h
@@ -162,6 +162,8 @@ void rb_st_cleanup_safe(st_table *, st_data_t);
 #define st_cleanup_safe rb_st_cleanup_safe
 void rb_st_clear(st_table *);
 #define st_clear rb_st_clear
+st_table *rb_st_replace(st_table *new_tab, st_table *old_tab);
+#define st_replace rb_st_replace
 st_table *rb_st_copy(st_table *);
 #define st_copy rb_st_copy
 CONSTFUNC(int rb_st_numcmp(st_data_t, st_data_t));
diff --git a/parser_st.c b/parser_st.c
@@ -109,6 +109,8 @@ nonempty_memcpy(void *dest, const void *src, size_t n)
 #define st_add_direct rb_parser_st_add_direct
 #undef st_insert2
 #define st_insert2 rb_parser_st_insert2
+#undef st_replace
+#define st_replace rb_parser_st_replace
 #undef st_copy
 #define st_copy rb_parser_st_copy
 #undef st_delete_safe
diff --git a/parser_st.h b/parser_st.h
@@ -137,6 +137,7 @@ void rb_parser_st_add_direct(parser_st_table *, parser_st_data_t, parser_st_data
 void rb_parser_st_free_table(parser_st_table *);
 void rb_parser_st_cleanup_safe(parser_st_table *, parser_st_data_t);
 void rb_parser_st_clear(parser_st_table *);
+parser_st_table *rb_parser_st_replace(parser_st_table *, parser_st_table *);
 parser_st_table *rb_parser_st_copy(parser_st_table *);
 CONSTFUNC(int rb_parser_st_numcmp(parser_st_data_t, parser_st_data_t));
 CONSTFUNC(parser_st_index_t rb_parser_st_numhash(parser_st_data_t));
diff --git a/st.c b/st.c
@@ -1228,25 +1228,17 @@ st_insert2(st_table *tab, st_data_t key, st_data_t value,
     return 1;
 }
 
-/* Create and return a copy of table OLD_TAB.  */
+/* Create a copy of old_tab into new_tab. */
 st_table *
-st_copy(st_table *old_tab)
+st_replace(st_table *new_tab, st_table *old_tab)
 {
-    st_table *new_tab;
-
-    new_tab = (st_table *) malloc(sizeof(st_table));
-#ifndef RUBY
-    if (new_tab == NULL)
-        return NULL;
-#endif
     *new_tab = *old_tab;
     if (old_tab->bins == NULL)
         new_tab->bins = NULL;
     else {
         new_tab->bins = (st_index_t *) malloc(bins_size(old_tab));
 #ifndef RUBY
         if (new_tab->bins == NULL) {
-            free(new_tab);
             return NULL;
         }
 #endif
@@ -1255,14 +1247,34 @@ st_copy(st_table *old_tab)
                                                  * sizeof(st_table_entry));
 #ifndef RUBY
     if (new_tab->entries == NULL) {
-        st_free_table(new_tab);
         return NULL;
     }
 #endif
     MEMCPY(new_tab->entries, old_tab->entries, st_table_entry,
            get_allocated_entries(old_tab));
     if (old_tab->bins != NULL)
         MEMCPY(new_tab->bins, old_tab->bins, char, bins_size(old_tab));
+
+    return new_tab;
+}
+
+/* Create and return a copy of table OLD_TAB.  */
+st_table *
+st_copy(st_table *old_tab)
+{
+    st_table *new_tab;
+
+    new_tab = (st_table *) malloc(sizeof(st_table));
+#ifndef RUBY
+    if (new_tab == NULL)
+        return NULL;
+#endif
+
+    if (st_replace(new_tab, old_tab) == NULL) {
+        st_free_table(new_tab);
+        return NULL;
+    }
+
     return new_tab;
 }
 

Original file line number	Diff line number	Diff line change
`@@ -1482,7 +1482,9 @@ hash_copy(VALUE ret, VALUE hash)`
`1482`	`1482`	`else {`
`1483`	`1483`	`HASH_ASSERT(sizeof(st_table) <= sizeof(ar_table));`
`1484`	`1484`
`1485`		`- RHASH_ST_TABLE_SET(ret, st_copy(RHASH_ST_TABLE(hash)));`
	`1485`	`+ RHASH_SET_ST_FLAG(ret);`
	`1486`	`+ st_replace(RHASH_ST_TABLE(ret), RHASH_ST_TABLE(hash));`
	`1487`	`+`
`1486`	`1488`	`rb_gc_writebarrier_remember(ret);`
`1487`	`1489`	`}`
`1488`	`1490`	`return ret;`
`@@ -1776,7 +1778,8 @@ rb_hash_s_create(int argc, VALUE *argv, VALUE klass)`
`1776`	`1778`	`}`
`1777`	`1779`	`else {`
`1778`	`1780`	`hash = hash_alloc(klass);`
`1779`		`- hash_copy(hash, tmp);`
	`1781`	`+ if (!RHASH_EMPTY_P(tmp))`
	`1782`	`+ hash_copy(hash, tmp);`
`1780`	`1783`	`return hash;`
`1781`	`1784`	`}`
`1782`	`1785`	`}`