Merge pull request #3134 from btw616:fix/assignCPU

PiperOrigin-RevId: 586091363

pkg/sentry/kernel/task_sched.go

@@ -621,9 +621,10 @@ func assignCPU(allowed sched.CPUSet, tid ThreadID) (cpu int32) {
 	n := int(tid) % int(allowed.NumCPUs())
 	// ... then pick the nth CPU in allowed.
 	allowed.ForEachCPU(func(c uint) {
-		if n--; n == 0 {
+		if n == 0 {
 			cpu = int32(c)
 		}
+		n--
 	})
 	return cpu
 }

pkg/sentry/kernel/task_test.go

@@ -29,35 +29,41 @@ func TestTaskCPU(t *testing.T) {
 		{
 			mask: []byte{0xff},
 			tid:  1,
-			cpu:  0,
+			cpu:  1,
 		},
 		{
 			mask: []byte{0xff},
 			tid:  10,
-			cpu:  1,
+			cpu:  2,
 		},
 		{
 			// more than 8 cpus.
 			mask: []byte{0xff, 0xff},
 			tid:  10,
-			cpu:  9,
+			cpu:  10,
 		},
 		{
 			// missing the first cpu.
 			mask: []byte{0xfe},
 			tid:  1,
-			cpu:  1,
+			cpu:  2,
 		},
 		{
 			mask: []byte{0xfe},
 			tid:  10,
-			cpu:  3,
+			cpu:  4,
 		},
 		{
 			// missing the fifth cpu.
 			mask: []byte{0xef},
 			tid:  10,
-			cpu:  2,
+			cpu:  3,
+		},
+		{
+			// only the fifth cpu.
+			mask: []byte{0x10},
+			tid:  10,
+			cpu:  4,
 		},
 	} {
 		assigned := assignCPU(test.mask, test.tid)

test/syscalls/linux/getcpu.cc

@@ -34,6 +34,26 @@ TEST(GetcpuTest, IsValidCpuStress) {
   }
 }
 
+TEST(GetcpuTest, IsValidCpu) {
+  const int num_cpus = NumCPUs();
+  cpu_set_t orig_set;
+  ASSERT_THAT(sched_getaffinity(getpid(), sizeof(orig_set), &orig_set),
+              SyscallSucceeds());
+  for (int i = 0; i < num_cpus; i++) {
+    if (CPU_ISSET(i, &orig_set) == 0) continue;
+    cpu_set_t set = {};
+    int cpu;
+    CPU_SET(i, &set);
+    ASSERT_THAT(sched_setaffinity(getpid(), sizeof(set), &set),
+                SyscallSucceeds());
+    ASSERT_THAT(cpu = sched_getcpu(), SyscallSucceeds());
+    // sched_setaffinity doesn't work if Kernel.useHostCores is true.
+    ASSERT_THAT(sched_getaffinity(getpid(), sizeof(set), &set),
+                SyscallSucceeds());
+    ASSERT_NE(CPU_ISSET(cpu, &set), 0);
+  }
+}
+
 }  // namespace
 
 }  // namespace testing