3 files changed, 90 insertions, 30 deletions

diff --git a/include/migration.h b/include/migration.h
index f60dfe9..bd52dd4 100644
--- a/include/migration.h
+++ b/include/migration.h
@@ -3,6 +3,7 @@
  * Functions to migrate tasks to different CPUs, partitions, clusters...
  */
 
+#include <sched.h>
 typedef int pid_t;
 
 /**
@@ -58,6 +59,20 @@ int be_migrate_to_cluster(int cluster, int cluster_sz);
 int be_migrate_to_domain(int domain);
 
 /**
+  * Parse CPU set given as string, and return corresponding cpu_set_t variable
+  * and its size as parameters.
+  * @param buf CPU mapping as string. Specific format is explained in
+  * src/migration.c:82 
+  * @param len Length of input string in bytes
+  * @param set cpu_set_t type variable with bits set as according to input
+  * @param sz Size of CPU set in bytes
+  * @return void
+*/
+
+void set_mapping(char* buf, int len, cpu_set_t** set, size_t *sz);
+
+
+/**
  * Return the number of CPUs currently online
  * @return The number of online CPUs
  */
diff --git a/src/migration.c b/src/migration.c
index c1e5a80..efd8d81 100644
--- a/src/migration.c
+++ b/src/migration.c
@@ -27,43 +27,17 @@ int num_online_cpus()
         return sysconf(_SC_NPROCESSORS_ONLN);
 }
 
-static int read_mapping(int idx, const char* which, cpu_set_t** set, size_t *sz)
+void set_mapping(char* buf, int len, cpu_set_t** set, size_t *sz)
 {
-        /* Max CPUs = 4096 */
-
-        int     ret = -1;
-        char buf[4096/4     /* enough chars for hex data (4 CPUs per char) */
-               + 4096/(4*8) /* for commas (separate groups of 8 chars) */
-               + 1] = {0};  /* for \0 */
-        char fname[80] = {0};
-
-        char* chunk_str;
-        int len, nbits;
+        int nbits;
+        char *chunk_str;
         int i;
 
         /* init vals returned to callee */
         *set = NULL;
         *sz = 0;
 
-        if (num_online_cpus() > 4096)
-                goto out;
-
-        /* Read string is in the format of <mask>[,<mask>]*. All <mask>s following
-           a comma are 8 chars (representing a 32-bit mask). The first <mask> may
-           have fewer chars. Bits are MSB to LSB, left to right. */
-        snprintf(fname, sizeof(fname), "/proc/litmus/%s/%d", which, idx);
-        ret = read_file(fname, &buf, sizeof(buf)-1);
-        if (ret <= 0)
-                goto out;
-
-        len = strnlen(buf, sizeof(buf));
-        /* if there is, omit newline at the end of string */
-        if (buf[len-1] == '\n') {
-                buf[len-1] = '\0';
-                len -= 1;
-        }
         nbits = 32*(len/9) + 4*(len%9); /* compute bits, accounting for commas */
-
         *set = CPU_ALLOC(nbits);
         *sz = CPU_ALLOC_SIZE(nbits);
         CPU_ZERO_S(*sz, *set);
@@ -74,7 +48,7 @@ static int read_mapping(int idx, const char* which, cpu_set_t** set, size_t *sz)
         do {
                 unsigned long chunk;
                 /* since strtoul stops processing the string with occurrence of
-                first non-digit character, it is necessary to read 8-bytes 
+                first non-digit character, it is necessary to read 8-bytes
                 on first iteration for ignoring the leading comma*/
                 chunk_str -= (9 + ((i == 0) ? -1 : 0));
                 if (chunk_str < buf)
@@ -88,6 +62,39 @@ static int read_mapping(int idx, const char* which, cpu_set_t** set, size_t *sz)
                 }
                 i += 1;
         } while (chunk_str > buf);
+}
+
+static int read_mapping(int idx, const char* which, cpu_set_t** set, size_t *sz)
+{
+        /* Max CPUs = 4096 */
+
+        int     ret = -1;
+        char buf[4096/4     /* enough chars for hex data (4 CPUs per char) */
+               + 4096/(4*8) /* for commas (separate groups of 8 chars) */
+               + 1] = {0};  /* for \0 */
+        char fname[80] = {0};
+
+        int len;
+
+        if (num_online_cpus() > 4096)
+                goto out;
+
+        /* Read string is in the format of <mask>[,<mask>]*. All <mask>s following
+           a comma are 8 chars (representing a 32-bit mask). The first <mask> may
+           have fewer chars. Bits are MSB to LSB, left to right. */
+        snprintf(fname, sizeof(fname), "/proc/litmus/%s/%d", which, idx);
+        ret = read_file(fname, &buf, sizeof(buf)-1);
+        if (ret <= 0)
+                goto out;
+
+        len = strnlen(buf, sizeof(buf));
+        /* if there is, omit newline at the end of string */
+        if (buf[len-1] == '\n') {
+                buf[len-1] = '\0';
+                len -= 1;
+        }
+
+        set_mapping(buf, len, set, sz);
 
         ret = 0;
 
diff --git a/tests/core_api.c b/tests/core_api.c
index feb60c5..065c1fc 100644
--- a/tests/core_api.c
+++ b/tests/core_api.c
@@ -193,6 +193,44 @@ TESTCASE(ctrl_page_writable, ALL,
         ctrl_page[32] = 0x12345678;
 }
 
+TESTCASE(set_cpu_mapping, LITMUS,
+         "task's CPU affinity is set to CPU set that is read from file")
+{
+        char buf[4096/4 /* enough chars for hex data (4 CPUs per char) */
+                + 4096/(4*8) /* for commas (separate groups of 8 chars) */
+                + 1] = {0}; /* for \0 */
+        int len;
+        cpu_set_t *set;
+        size_t sz;
+        
+        /*set affinity to CPU 0 */
+        strcpy(buf, "20");
+        len = strnlen(buf, sizeof(buf));
+        set_mapping(buf, len, &set, &sz);
+        ASSERT( CPU_ISSET_S(5, sz, set) );
+
+        /*set affinity to CPU 29 */
+        strcpy(buf, "20000000");
+        len = strnlen(buf, sizeof(buf));
+        set_mapping(buf, len, &set, &sz);
+        ASSERT( CPU_ISSET_S(29, sz, set) );
+
+        /*set affinity to CPUS 27 and 39 */
+        strcpy(buf, "80,08000000");
+        len = strnlen(buf, sizeof(buf));
+        set_mapping(buf, len, &set, &sz);
+        ASSERT( CPU_ISSET_S(27, sz, set) );
+        ASSERT( CPU_ISSET_S(39, sz, set) );
+
+        /*set affinity to CPUS 96 */
+        strcpy(buf, "1,00000000,00000000,00000000");
+        len = strnlen(buf, sizeof(buf));
+        set_mapping(buf, len, &set, &sz);
+        ASSERT( CPU_ISSET_S(96, sz, set) );
+
+}
+
+
 
 TESTCASE(suspended_admission, LITMUS,
          "admission control handles suspended tasks correctly")