util: process: Refactor and fix virProcessSetAffinity

Refactor the function to return the bitmap instead of an integer and the
inner workings so that they make more sense.

This patch also fixes possible segfault on old systems that was
introduced by commit:

commit f1a43a8e41
Author: Hu Tao <hutao@cn.fujitsu.com>
Date:   Fri Sep 14 15:46:59 2012 +0800

    use virBitmap to store cpu affinity info

src/qemu/qemu_driver.c

@@ -1457,8 +1457,7 @@ qemuDomainHelperGetVcpus(virDomainObjPtr vm, virVcpuInfoPtr info, int maxinfo,
                 unsigned char *cpumap = VIR_GET_CPUMAP(cpumaps, maplen, v);
                 virBitmapPtr map = NULL;
 
-                if (virProcessGetAffinity(priv->vcpupids[v],
-                                          &map, hostcpus) < 0)
+                if (!(map = virProcessGetAffinity(priv->vcpupids[v])))
                     return -1;
 
                 virBitmapToDataBuf(map, cpumap, maplen);
@@ -5727,7 +5726,7 @@ qemuDomainGetIOThreadsLive(virQEMUDriverPtr driver,
             goto endjob;
         info_ret[i]->iothread_id = iothreads[i]->iothread_id;
 
-        if (virProcessGetAffinity(iothreads[i]->thread_id, &map, hostcpus) < 0)
+        if (!(map = virProcessGetAffinity(iothreads[i]->thread_id)))
             goto endjob;
 
         if (virBitmapToData(map, &info_ret[i]->cpumap,

src/util/virprocess.c

@@ -468,71 +468,60 @@ int virProcessSetAffinity(pid_t pid, virBitmapPtr map)
     return 0;
 }
 
-int virProcessGetAffinity(pid_t pid,
-                          virBitmapPtr *map,
-                          int maxcpu)
+virBitmapPtr
+virProcessGetAffinity(pid_t pid)
 {
     size_t i;
-# ifdef CPU_ALLOC
-    /* New method dynamically allocates cpu mask, allowing unlimted cpus */
-    int numcpus = 1024;
-    size_t masklen;
     cpu_set_t *mask;
+    size_t masklen;
+    virBitmapPtr ret = NULL;
 
-    /* Not only may the statically allocated cpu_set_t be too small,
-     * but there is no way to ask the kernel what size is large enough.
-     * So you have no option but to pick a size, try, catch EINVAL,
-     * enlarge, and re-try.
-     *
-     * http://lkml.org/lkml/2009/7/28/620
-     */
- realloc:
-    masklen = CPU_ALLOC_SIZE(numcpus);
-    mask = CPU_ALLOC(numcpus);
+# ifdef CPU_ALLOC
+    /* 262144 cpus ought to be enough for anyone */
+    masklen = CPU_ALLOC_SIZE(1024 << 8);
+    mask = CPU_ALLOC(1024 << 8);
 
     if (!mask) {
         virReportOOMError();
-        return -1;
+        return NULL;
     }
 
     CPU_ZERO_S(masklen, mask);
-    if (sched_getaffinity(pid, masklen, mask) < 0) {
-        CPU_FREE(mask);
-        if (errno == EINVAL &&
-            numcpus < (1024 << 8)) { /* 262144 cpus ought to be enough for anyone */
-            numcpus = numcpus << 2;
-            goto realloc;
-        }
-        virReportSystemError(errno,
-                             _("cannot get CPU affinity of process %d"), pid);
-        return -1;
-    }
-
-    *map = virBitmapNew(maxcpu);
-    if (!*map)
-        return -1;
-
-    for (i = 0; i < maxcpu; i++)
-        if (CPU_ISSET_S(i, masklen, mask))
-            ignore_value(virBitmapSetBit(*map, i));
-    CPU_FREE(mask);
 # else
-    /* Legacy method uses a fixed size cpu mask, only allows up to 1024 cpus */
-    cpu_set_t mask;
+    if (VIR_ALLOC(mask) < 0)
+        return NULL;
 
-    CPU_ZERO(&mask);
-    if (sched_getaffinity(pid, sizeof(mask), &mask) < 0) {
-        virReportSystemError(errno,
-                             _("cannot get CPU affinity of process %d"), pid);
-        return -1;
-    }
-
-    for (i = 0; i < maxcpu; i++)
-        if (CPU_ISSET(i, &mask))
-            ignore_value(virBitmapSetBit(*map, i));
+    masklen = sizeof(*mask);
+    CPU_ZERO(mask);
 # endif
 
-    return 0;
+    if (sched_getaffinity(pid, masklen, mask) < 0) {
+        virReportSystemError(errno,
+                             _("cannot get CPU affinity of process %d"), pid);
+        goto cleanup;
+    }
+
+    if (!(ret = virBitmapNew(masklen * 8)))
+          goto cleanup;
+
+    for (i = 0; i < masklen * 8; i++) {
+# ifdef CPU_ALLOC
+        if (CPU_ISSET_S(i, masklen, mask))
+            ignore_value(virBitmapSetBit(ret, i));
+# else
+        if (CPU_ISSET(i, mask))
+            ignore_value(virBitmapSetBit(ret, i));
+# endif
+    }
+
+ cleanup:
+# ifdef CPU_ALLOC
+    CPU_FREE(mask);
+# else
+    VIR_FREE(mask);
+# endif
+
+    return ret;
 }
 
 #elif defined(HAVE_BSD_CPU_AFFINITY)
@@ -562,29 +551,29 @@ int virProcessSetAffinity(pid_t pid,
     return 0;
 }
 
-int virProcessGetAffinity(pid_t pid,
-                          virBitmapPtr *map,
-                          int maxcpu)
+virBitmapPtr
+virProcessGetAffinity(pid_t pid)
 {
     size_t i;
     cpuset_t mask;
-
-    if (!(*map = virBitmapNew(maxcpu)))
-        return -1;
+    virBitmapPtr ret = NULL;
 
     CPU_ZERO(&mask);
     if (cpuset_getaffinity(CPU_LEVEL_WHICH, CPU_WHICH_PID, pid,
                            sizeof(mask), &mask) != 0) {
         virReportSystemError(errno,
                              _("cannot get CPU affinity of process %d"), pid);
-        return -1;
+        return NULL;
     }
 
-    for (i = 0; i < maxcpu; i++)
-        if (CPU_ISSET(i, &mask))
-            ignore_value(virBitmapSetBit(*map, i));
+    if (!(*map = virBitmapNew(sizeof(mask) * 8)))
+        return NULL;
 
-    return 0;
+    for (i = 0; i < sizeof(mask) * 8; i++)
+        if (CPU_ISSET(i, &mask))
+            ignore_value(virBitmapSetBit(ret, i));
+
+    return ret;
 }
 
 #else /* HAVE_SCHED_GETAFFINITY */
@@ -597,13 +586,12 @@ int virProcessSetAffinity(pid_t pid ATTRIBUTE_UNUSED,
     return -1;
 }
 
-int virProcessGetAffinity(pid_t pid ATTRIBUTE_UNUSED,
-                          virBitmapPtr *map ATTRIBUTE_UNUSED,
-                          int maxcpu ATTRIBUTE_UNUSED)
+virBitmapPtr
+virProcessGetAffinity(pid_t pid ATTRIBUTE_UNUSED)
 {
     virReportSystemError(ENOSYS, "%s",
                          _("Process CPU affinity is not supported on this platform"));
-    return -1;
+    return NULL;
 }
 #endif /* HAVE_SCHED_GETAFFINITY */
 

src/util/virprocess.h

@@ -58,9 +58,7 @@ int virProcessKillPainfully(pid_t pid, bool force);
 
 int virProcessSetAffinity(pid_t pid, virBitmapPtr map);
 
-int virProcessGetAffinity(pid_t pid,
-                          virBitmapPtr *map,
-                          int maxcpu);
+virBitmapPtr virProcessGetAffinity(pid_t pid);
 
 int virProcessGetStartTime(pid_t pid,
                            unsigned long long *timestamp);