Implement core tagging

Links to some possibly useful reference material on core tagging:
* LWN article: https://lwn.net/Articles/861251/
* Kernel docs: https://www.kernel.org/doc/Documentation/admin-guide/hw-vuln/core-scheduling.rst

PiperOrigin-RevId: 431093418

Author: misterwilliam

WORKSPACE

@@ -88,8 +88,8 @@ go_repository(
 go_repository(
     name = "org_golang_x_sys",
     importpath = "golang.org/x/sys",
-    sum = "h1:oN6lz7iLW/YC7un8pq+9bOLyXrprv2+DKfkJY+2LJJw=",
-    version = "v0.0.0-20211007075335-d3039528d8ac",
+    sum = "h1:BXxu8t6QN0G1uff4bzZzSkpsax8+ALqTGUtz08QrV00=",
+    version = "v0.0.0-20220224120231-95c6836cb0e7",
 )
 
 go_repository(

pkg/abi/linux/prctl.go

@@ -145,6 +145,10 @@ const (
 	// Protection eXtensions (MPX) bounds tables.
 	PR_MPX_DISABLE_MANAGEMENT = 44
 
+	// The following constants are used to control thread scheduling on cores.
+	PR_SCHED_CORE_SCOPE_THREAD       = 0
+	PR_SCHED_CORE_SCOPE_THREAD_GROUP = 1
+
 	// PR_SET_PTRACER allows a specific process (or any, if PR_SET_PTRACER_ANY is
 	// specified) to ptrace the current task.
 	PR_SET_PTRACER     = 0x59616d61

pkg/coretag/BUILD

@@ -0,0 +1,28 @@
+load("//tools:defs.bzl", "go_library", "go_test")
+
+package(licenses = ["notice"])
+
+go_library(
+    name = "coretag",
+    srcs = [
+        "coretag.go",
+        "coretag_unsafe.go",
+    ],
+    visibility = ["//:sandbox"],
+    deps = [
+        "//pkg/abi/linux",
+        "@org_golang_x_sys//unix:go_default_library",
+    ],
+)
+
+go_test(
+    name = "coretag_test",
+    size = "small",
+    srcs = [
+        "coretag_test.go",
+    ],
+    library = ":coretag",
+    deps = [
+        "//pkg/hostos",
+    ],
+)

pkg/coretag/coretag.go

@@ -0,0 +1,93 @@
+// Copyright 2022 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Package coretag implements core tagging.
+package coretag
+
+import (
+	"fmt"
+	"io/ioutil"
+	"strconv"
+
+	"golang.org/x/sys/unix"
+	"gvisor.dev/gvisor/pkg/abi/linux"
+)
+
+// Enable core tagging. If this returns with no error, all threads in the
+// current thread group will be run in a core tagged thread. Only available on
+// linux kernel >= 5.14.
+func Enable() error {
+	// Set core tag on current thread group.
+	// prctl(PR_SCHED_CORE, PR_SCHED_CORE_CREATE, pid=0,
+	//       PR_SCHED_CORE_SCOPE_THREAD_GROUP, cookie=nullptr)
+	// pid=0 means current pid.
+	// cookie=nullptr is required for PR_SCHED_CORE_CREATE.
+	if _, _, errno := unix.Syscall6(unix.SYS_PRCTL, unix.PR_SCHED_CORE,
+		unix.PR_SCHED_CORE_CREATE, 0 /*pid*/, linux.PR_SCHED_CORE_SCOPE_THREAD_GROUP, 0, 0); errno != 0 {
+		return fmt.Errorf("failed to core tag sentry: %w", errno)
+	}
+	return nil
+}
+
+// GetAllCoreTags returns the core tag of all the threads in the thread group.
+func GetAllCoreTags(pid int) ([]uint64, error) {
+	// prctl(PR_SCHED_CORE_GET, PR_SCHED_CORE_SCOPE_THREAD_GROUP, ...) is not supported
+	// in linux. So instead we get all threads from /proc/<pid>/task and get all the
+	// core tags individually.
+	tagSet := make(map[uint64]struct{})
+	// Get current pid core tag.
+	tag, err := getCoreTag(pid)
+	if err != nil {
+		return nil, err
+	}
+	tagSet[tag] = struct{}{}
+
+	// Get core tags of tids.
+	tids, err := getTids(pid)
+	if err != nil {
+		return nil, err
+	}
+	for tid := range tids {
+		tag, err := getCoreTag(tid)
+		if err != nil {
+			return nil, err
+		}
+		tagSet[tag] = struct{}{}
+	}
+
+	// Return set of tags as a slice.
+	tags := make([]uint64, 0, len(tagSet))
+	for t := range tagSet {
+		tags = append(tags, t)
+	}
+	return tags, nil
+}
+
+// getTids returns set of tids as reported by /proc/<pid>/task.
+func getTids(pid int) (map[int]struct{}, error) {
+	tids := make(map[int]struct{})
+	files, err := ioutil.ReadDir("/proc/" + strconv.Itoa(pid) + "/task")
+	if err != nil {
+		return nil, err
+	}
+	for _, file := range files {
+		tid, err := strconv.Atoi(file.Name())
+		if err != nil {
+			return nil, err
+		}
+		tids[tid] = struct{}{}
+	}
+
+	return tids, nil
+}

pkg/coretag/coretag_test.go

@@ -0,0 +1,46 @@
+// Copyright 2022 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package coretag
+
+import (
+	"os"
+	"testing"
+
+	"gvisor.dev/gvisor/pkg/hostos"
+)
+
+func TestEnable(t *testing.T) {
+	major, minor, err := hostos.KernelVersion()
+	if err != nil {
+		t.Fatalf("Unable to parse kernel version: %v", err)
+	}
+	// Skip running test when running on Linux kernel < 5.14 because core tagging
+	// is not available.
+	if major < 5 && minor < 14 {
+		t.Skipf("Running on Linux kernel: %d.%d < 5.14. Core tagging not available. Skipping test.", major, minor)
+		return
+	}
+	if err := Enable(); err != nil {
+		t.Fatalf("Enable() got error %v, wanted nil", err)
+	}
+
+	coreTags, err := GetAllCoreTags(os.Getpid())
+	if err != nil {
+		t.Fatalf("GetAllCoreTags() got error %v, wanted nil", err)
+	}
+	if len(coreTags) != 1 {
+		t.Fatalf("Got coreTags %v, wanted len(coreTags)=1", coreTags)
+	}
+}

pkg/coretag/coretag_unsafe.go

@@ -0,0 +1,34 @@
+// Copyright 2022 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package coretag
+
+import (
+	"fmt"
+	"unsafe"
+
+	"golang.org/x/sys/unix"
+	"gvisor.dev/gvisor/pkg/abi/linux"
+)
+
+// getCoreTag returns the core tag of the tid. Only available on linux kernel >= 5.14.
+func getCoreTag(tid int) (uint64, error) {
+	var cookie uint64
+	if _, _, errno := unix.Syscall6(unix.SYS_PRCTL, unix.PR_SCHED_CORE,
+		unix.PR_SCHED_CORE_GET, uintptr(tid), linux.PR_SCHED_CORE_SCOPE_THREAD,
+		uintptr(unsafe.Pointer(&cookie)), 0); errno != 0 {
+		return 0, fmt.Errorf("prctl(PR_SCHED_CORE, PR_SCHED_CORE_GET, %d, PR_SCHED_CORE_SCOPE_THREAD) (errno=%d)", tid, errno)
+	}
+	return cookie, nil
+}

pkg/cpuid/BUILD

@@ -42,5 +42,5 @@ go_test(
     ],
     library = ":cpuid",
     tags = ["manual"],
-    deps = ["@org_golang_x_sys//unix:go_default_library"],
+    deps = ["//pkg/hostos"],
 )

pkg/cpuid/cpuid_parse_test.go

@@ -15,48 +15,14 @@
 package cpuid
 
 import (
-	"fmt"
 	"io/ioutil"
 	"regexp"
-	"strconv"
 	"strings"
 	"testing"
 
-	"golang.org/x/sys/unix"
+	"gvisor.dev/gvisor/pkg/hostos"
 )
 
-func kernelVersion() (int, int, error) {
-	var u unix.Utsname
-	if err := unix.Uname(&u); err != nil {
-		return 0, 0, err
-	}
-
-	var sb strings.Builder
-	for _, b := range u.Release {
-		if b == 0 {
-			break
-		}
-		sb.WriteByte(byte(b))
-	}
-
-	s := strings.Split(sb.String(), ".")
-	if len(s) < 2 {
-		return 0, 0, fmt.Errorf("kernel release missing major and minor component: %s", sb.String())
-	}
-
-	major, err := strconv.Atoi(s[0])
-	if err != nil {
-		return 0, 0, fmt.Errorf("error parsing major version %q in %q: %w", s[0], sb.String(), err)
-	}
-
-	minor, err := strconv.Atoi(s[1])
-	if err != nil {
-		return 0, 0, fmt.Errorf("error parsing minor version %q in %q: %w", s[1], sb.String(), err)
-	}
-
-	return major, minor, nil
-}
-
 // TestHostFeatureFlags tests that all features detected by HostFeatureSet are
 // on the host.
 //
@@ -65,7 +31,7 @@ func kernelVersion() (int, int, error) {
 // analog in the actual CPUID feature set.
 func TestHostFeatureFlags(t *testing.T) {
 	// Extract the kernel version.
-	major, minor, err := kernelVersion()
+	major, minor, err := hostos.KernelVersion()
 	if err != nil {
 		t.Fatalf("Unable to parse kernel version: %v", err)
 	}

pkg/hostos/BUILD

@@ -0,0 +1,10 @@
+load("//tools:defs.bzl", "go_library")
+
+package(licenses = ["notice"])
+
+go_library(
+    name = "hostos",
+    srcs = ["hostos.go"],
+    visibility = ["//:sandbox"],
+    deps = ["@org_golang_x_sys//unix:go_default_library"],
+)

pkg/hostos/hostos.go

@@ -0,0 +1,57 @@
+// Copyright 2022 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Package hostos contains utility functions for getting information about the host OS.
+package hostos
+
+import (
+	"fmt"
+	"strconv"
+	"strings"
+
+	"golang.org/x/sys/unix"
+)
+
+// KernelVersion returns the major and minor release version of the kernel using uname().
+func KernelVersion() (int, int, error) {
+	var u unix.Utsname
+	if err := unix.Uname(&u); err != nil {
+		return 0, 0, err
+	}
+
+	var sb strings.Builder
+	for _, b := range u.Release {
+		if b == 0 {
+			break
+		}
+		sb.WriteByte(byte(b))
+	}
+
+	s := strings.Split(sb.String(), ".")
+	if len(s) < 2 {
+		return 0, 0, fmt.Errorf("kernel release missing major and minor component: %s", sb.String())
+	}
+
+	major, err := strconv.Atoi(s[0])
+	if err != nil {
+		return 0, 0, fmt.Errorf("error parsing major version %q in %q: %w", s[0], sb.String(), err)
+	}
+
+	minor, err := strconv.Atoi(s[1])
+	if err != nil {
+		return 0, 0, fmt.Errorf("error parsing minor version %q in %q: %w", s[1], sb.String(), err)
+	}
+
+	return major, minor, nil
+}

runsc/cmd/BUILD

@@ -45,6 +45,7 @@ go_library(
         "//runsc:__subpackages__",
     ],
     deps = [
+        "//pkg/coretag",
         "//pkg/coverage",
         "//pkg/log",
         "//pkg/p9",