2 files changed, 17 insertions, 21 deletions

diff --git a/ext/socket/ipsocket.c b/ext/socket/ipsocket.c
index 4dfd9c8a56..da42fbd27b 100644
--- a/ext/socket/ipsocket.c
+++ b/ext/socket/ipsocket.c
@@ -1159,13 +1159,19 @@ fast_fallback_inetsock_cleanup(VALUE v)
         getaddrinfo_shared->notify = -1;
 
         int shared_need_free = 0;
-        int need_free[2] = { 0, 0 };
+        struct addrinfo *ais[arg->family_size];
+        for (int i = 0; i < arg->family_size; i++) ais[i] = NULL;
 
         rb_nativethread_lock_lock(&getaddrinfo_shared->lock);
         {
             for (int i = 0; i < arg->family_size; i++) {
-                if (arg->getaddrinfo_entries[i] && --(arg->getaddrinfo_entries[i]->refcount) == 0) {
-                    need_free[i] = 1;
+                struct fast_fallback_getaddrinfo_entry *getaddrinfo_entry = arg->getaddrinfo_entries[i];
+
+                if (!getaddrinfo_entry) continue;
+
+                if (--(getaddrinfo_entry->refcount) == 0) {
+                    ais[i] = getaddrinfo_entry->ai;
+                    getaddrinfo_entry->ai = NULL;
                 }
             }
             if (--(getaddrinfo_shared->refcount) == 0) {
@@ -1175,9 +1181,7 @@ fast_fallback_inetsock_cleanup(VALUE v)
         rb_nativethread_lock_unlock(&getaddrinfo_shared->lock);
 
         for (int i = 0; i < arg->family_size; i++) {
-            if (arg->getaddrinfo_entries[i] && need_free[i]) {
-                free_fast_fallback_getaddrinfo_entry(&arg->getaddrinfo_entries[i]);
-            }
+            if (ais[i]) freeaddrinfo(ais[i]);
         }
         if (getaddrinfo_shared && shared_need_free) {
             free_fast_fallback_getaddrinfo_shared(&getaddrinfo_shared);
diff --git a/ext/socket/raddrinfo.c b/ext/socket/raddrinfo.c
index ac47b5b256..91e2be1148 100644
--- a/ext/socket/raddrinfo.c
+++ b/ext/socket/raddrinfo.c
@@ -3038,22 +3038,13 @@ free_fast_fallback_getaddrinfo_shared(struct fast_fallback_getaddrinfo_shared **
     *shared = NULL;
 }
 
-void
-free_fast_fallback_getaddrinfo_entry(struct fast_fallback_getaddrinfo_entry **entry)
-{
-    if ((*entry)->ai) {
-        freeaddrinfo((*entry)->ai);
-        (*entry)->ai = NULL;
-    }
-    *entry = NULL;
-}
-
 static void *
 do_fast_fallback_getaddrinfo(void *ptr)
 {
     struct fast_fallback_getaddrinfo_entry *entry = (struct fast_fallback_getaddrinfo_entry *)ptr;
     struct fast_fallback_getaddrinfo_shared *shared = entry->shared;
-    int err = 0, need_free = 0, shared_need_free = 0;
+    int err = 0, shared_need_free = 0;
+    struct addrinfo *ai = NULL;
 
     sigset_t set;
     sigemptyset(&set);
@@ -3102,14 +3093,15 @@ do_fast_fallback_getaddrinfo(void *ptr)
             entry->err = errno;
             entry->has_syserr = true;
         }
-        if (--(entry->refcount) == 0) need_free = 1;
+        if (--(entry->refcount) == 0) {
+            ai = entry->ai;
+            entry->ai = NULL;
+        }
         if (--(shared->refcount) == 0) shared_need_free = 1;
     }
     rb_nativethread_lock_unlock(&shared->lock);
 
-    if (need_free && entry) {
-        free_fast_fallback_getaddrinfo_entry(&entry);
-    }
+    if (ai) freeaddrinfo(ai);
     if (shared_need_free && shared) {
         free_fast_fallback_getaddrinfo_shared(&shared);
     }