Only gc on allocation failure.

Author: samuel-williams-shopify

cont.c

@@ -507,13 +507,14 @@ fiber_pool_allocate_memory(size_t * count, size_t stride)
 
 // Given an existing fiber pool, expand it by the specified number of stacks.
 // @param count the maximum number of stacks to allocate.
-// @return the allocated fiber pool.
+// @return the new allocation on success, or NULL on failure with errno set.
 // @sa fiber_pool_allocation_free
 static struct fiber_pool_allocation *
 fiber_pool_expand(struct fiber_pool * fiber_pool, size_t count)
 {
     struct fiber_pool_allocation * allocation;
-    RB_VM_LOCK_ENTER();
+    unsigned int lev;
+    RB_VM_LOCK_ENTER_LEV(&lev);
     {
         STACK_GROW_DIR_DETECTION;
 
@@ -524,7 +525,11 @@ fiber_pool_expand(struct fiber_pool * fiber_pool, size_t count)
         void * base = fiber_pool_allocate_memory(&count, stride);
 
         if (base == NULL) {
-            rb_raise(rb_eFiberError, "can't alloc machine stack to fiber (%"PRIuSIZE" x %"PRIuSIZE" bytes): %s", count, size, ERRNOMSG);
+            int err = errno;
+            if (!err) err = ENOMEM;
+            errno = err;
+            RB_VM_LOCK_LEAVE_LEV(&lev);
+            return NULL;
         }
 
         struct fiber_pool_vacancy * vacancies = fiber_pool->vacancies;
@@ -553,18 +558,26 @@ fiber_pool_expand(struct fiber_pool * fiber_pool, size_t count)
             DWORD old_protect;
 
             if (!VirtualProtect(page, RB_PAGE_SIZE, PAGE_READWRITE | PAGE_GUARD, &old_protect)) {
+                int err = errno;
+                if (!err) err = ENOMEM;
                 VirtualFree(allocation->base, 0, MEM_RELEASE);
                 ruby_xfree(allocation);
-                rb_raise(rb_eFiberError, "can't set a guard page: %s", ERRNOMSG);
+                errno = err;
+                RB_VM_LOCK_LEAVE_LEV(&lev);
+                return NULL;
             }
 #elif defined(__wasi__)
             // wasi-libc's mprotect emulation doesn't support PROT_NONE.
             (void)page;
 #else
             if (mprotect(page, RB_PAGE_SIZE, PROT_NONE) < 0) {
+                int err = errno;
+                if (!err) err = ENOMEM;
                 munmap(allocation->base, count*stride);
                 ruby_xfree(allocation);
-                rb_raise(rb_eFiberError, "can't set a guard page: %s", ERRNOMSG);
+                errno = err;
+                RB_VM_LOCK_LEAVE_LEV(&lev);
+                return NULL;
             }
 #endif
 
@@ -594,7 +607,7 @@ fiber_pool_expand(struct fiber_pool * fiber_pool, size_t count)
         fiber_pool->vacancies = vacancies;
         fiber_pool->count += count;
     }
-    RB_VM_LOCK_LEAVE();
+    RB_VM_LOCK_LEAVE_LEV(&lev);
 
     return allocation;
 }
@@ -616,7 +629,10 @@ fiber_pool_initialize(struct fiber_pool * fiber_pool, size_t size, size_t count,
 
     fiber_pool->vm_stack_size = vm_stack_size;
 
-    fiber_pool_expand(fiber_pool, count);
+    if (RB_UNLIKELY(!fiber_pool_expand(fiber_pool, count))) {
+        rb_raise(rb_eFiberError, "can't alloc machine stack to fiber (%"PRIuSIZE" x %"PRIuSIZE" bytes): %s",
+                 count, fiber_pool->size, strerror(errno));
+    }
 }
 
 #ifdef FIBER_POOL_ALLOCATION_FREE
@@ -664,42 +680,72 @@ fiber_pool_allocation_free(struct fiber_pool_allocation * allocation)
 }
 #endif
 
+// Number of stacks to request when expanding the pool (clamped to min/max).
+static inline size_t
+fiber_pool_stack_expand_count(const struct fiber_pool *pool)
+{
+    size_t count = pool->count;
+    if (count > FIBER_POOL_ALLOCATION_MAXIMUM_SIZE) count = FIBER_POOL_ALLOCATION_MAXIMUM_SIZE;
+    if (count < pool->initial_count) count = pool->initial_count;
+    return count;
+}
+
+// When the vacancy list is empty, grow the pool (and run GC only if mmap fails). Caller holds the VM lock.
+// Returns NULL if expansion failed after GC + retry; errno is set. Otherwise returns a vacancy.
+static struct fiber_pool_vacancy *
+fiber_pool_stack_acquire_expand(struct fiber_pool *fiber_pool)
+{
+    size_t count = fiber_pool_stack_expand_count(fiber_pool);
+
+    if (DEBUG) fprintf(stderr, "fiber_pool_stack_acquire: expanding fiber pool to %"PRIuSIZE" stacks\n", count);
+
+    struct fiber_pool_vacancy *vacancy = NULL;
+
+    if (RB_LIKELY(fiber_pool_expand(fiber_pool, count))) {
+        return fiber_pool_vacancy_pop(fiber_pool);
+    } else {
+        if (DEBUG) fprintf(stderr, "fiber_pool_stack_acquire: expand failed (%s), collecting garbage\n", strerror(errno));
+
+        rb_gc();
+
+        // After running GC, the vacancy list may have some stacks:
+        vacancy = fiber_pool_vacancy_pop(fiber_pool);
+        if (RB_LIKELY(vacancy)) {
+            return vacancy;
+        }
+
+        // Try to expand the fiber pool again:
+        if (RB_LIKELY(fiber_pool_expand(fiber_pool, count))) {
+            return fiber_pool_vacancy_pop(fiber_pool);
+        } else {
+            return NULL;
+        }
+    }
+}
+
 // Acquire a stack from the given fiber pool. If none are available, allocate more.
 static struct fiber_pool_stack
 fiber_pool_stack_acquire(struct fiber_pool * fiber_pool)
 {
     struct fiber_pool_vacancy * vacancy;
 
-    RB_VM_LOCK_ENTER();
+    unsigned int lev;
+    RB_VM_LOCK_ENTER_LEV(&lev);
     {
+        // Fast path: try to acquire a stack from the vacancy list:
         vacancy = fiber_pool_vacancy_pop(fiber_pool);
 
         if (DEBUG) fprintf(stderr, "fiber_pool_stack_acquire: %p used=%"PRIuSIZE"\n", (void*)fiber_pool->vacancies, fiber_pool->used);
 
-        // During allocation, if we run out of stacks, we may need to collect garbage to free up some stacks before trying to allocate more.
-        if (!vacancy) {
-            if (DEBUG) fprintf(stderr, "fiber_pool_stack_acquire: no stacks available, collecting garbage\n");
-            rb_gc();
-
-            // In the worst case, we garbage collect, but all fibers are reachable, so we don't free any stacks. We will try to expand the fiber pool:
-            vacancy = fiber_pool_vacancy_pop(fiber_pool);
-        }
-
-        if (!vacancy) {
-            const size_t maximum = FIBER_POOL_ALLOCATION_MAXIMUM_SIZE;
-            const size_t minimum = fiber_pool->initial_count;
-
-            size_t count = fiber_pool->count;
-            if (count > maximum) count = maximum;
-            if (count < minimum) count = minimum;
+        // Slow path: If the pool has no vacancies, expand first. Only run GC when expansion fails (e.g. mmap), so we can reclaim stacks from dead fibers before retrying:
+        if (RB_UNLIKELY(!vacancy)) {
+            vacancy = fiber_pool_stack_acquire_expand(fiber_pool);
 
-            if (DEBUG) fprintf(stderr, "fiber_pool_stack_acquire: expanding fiber pool to %"PRIuSIZE" stacks\n", count);
-            fiber_pool_expand(fiber_pool, count);
-
-            // The free list should now contain some stacks:
-            VM_ASSERT(fiber_pool->vacancies);
-
-            vacancy = fiber_pool_vacancy_pop(fiber_pool);
+            // If expansion failed, raise an error:
+            if (RB_UNLIKELY(!vacancy)) {
+                RB_VM_LOCK_LEAVE_LEV(&lev);
+                rb_raise(rb_eFiberError, "can't allocate fiber stack: %s", strerror(errno));
+            }
         }
 
         VM_ASSERT(vacancy);
@@ -718,7 +764,7 @@ fiber_pool_stack_acquire(struct fiber_pool * fiber_pool)
 
         fiber_pool_stack_reset(&vacancy->stack);
     }
-    RB_VM_LOCK_LEAVE();
+    RB_VM_LOCK_LEAVE_LEV(&lev);
 
     return vacancy->stack;
 }