From 0ecb1186f50c329696567349c4bfbe25b0e4b63a Mon Sep 17 00:00:00 2001
From: Moritz Sanft <58110325+msanft@users.noreply.github.com>
Date: Tue, 21 May 2024 15:18:24 +0200
Subject: [PATCH] bootstrapper: add etcd I/O prioritizer

---
 bootstrapper/internal/etcdio/etcdio.go | 120 +++++++++++++++++++++++++
 1 file changed, 120 insertions(+)
 create mode 100644 bootstrapper/internal/etcdio/etcdio.go

diff --git a/bootstrapper/internal/etcdio/etcdio.go b/bootstrapper/internal/etcdio/etcdio.go
new file mode 100644
index 000000000..23816f0d8
--- /dev/null
+++ b/bootstrapper/internal/etcdio/etcdio.go
@@ -0,0 +1,120 @@
+/*
+Copyright (c) Edgeless Systems GmbH
+
+SPDX-License-Identifier: AGPL-3.0-only
+*/
+
+// The etcdio package provides utilities to manage etcd I/O.
+package etcdio
+
+import (
+	"errors"
+	"fmt"
+	"log/slog"
+	"os"
+	"path"
+	"strconv"
+
+	"golang.org/x/sys/unix"
+)
+
+var (
+	ErrNoEtcdProcess         = errors.New("no etcd process found on node")
+	ErrMultipleEtcdProcesses = errors.New("multiple etcd processes found on node")
+)
+
+const (
+	// Tells the syscall that a process' priority is going to be set.
+	// See https://elixir.bootlin.com/linux/v6.9.1/source/include/uapi/linux/ioprio.h#L54.
+	ioPrioWhoProcess = 1
+
+	// See https://elixir.bootlin.com/linux/v6.9.1/source/include/uapi/linux/ioprio.h#L11.
+	ioPrioClassShift = 13
+	ioPrioNrClasses  = 8
+	ioPrioClassMask  = ioPrioNrClasses - 1
+	ioPrioPrioMask   = (1 << ioPrioClassShift) - 1
+
+	targetClass = 1 // Realtime IO class for best scheduling prio
+	targetPrio  = 0 // Highest priority within the class
+)
+
+// EtcdIOClient is a client for managing etcd I/O.
+type EtcdIOClient struct {
+	log *slog.Logger
+}
+
+// NewClient creates a new etcd I/O management client.
+func NewClient(log *slog.Logger) *EtcdIOClient {
+	return &EtcdIOClient{log: log}
+}
+
+// PrioritizeIO tries to find the etcd process on the node and prioritizes its I/O.
+func (c *EtcdIOClient) PrioritizeIO() error {
+	// find etcd process(es)
+	pid, err := c.findEtcdProcess()
+	if err != nil {
+		return fmt.Errorf("finding etcd process: %w", err)
+	}
+
+	// Highest realtime priority value for the etcd process, see https://elixir.bootlin.com/linux/v6.9.1/source/include/uapi/linux/ioprio.h
+	// for the calculation details.
+	prioVal := ((targetClass & ioPrioClassMask) << ioPrioClassShift) | (targetPrio & ioPrioPrioMask)
+
+	// see https://man7.org/linux/man-pages/man2/ioprio_set.2.html
+	ret, _, errno := unix.Syscall(unix.SYS_IOPRIO_SET, ioPrioWhoProcess, uintptr(pid), uintptr(prioVal))
+	if ret != 0 {
+		return fmt.Errorf("setting I/O priority for etcd: %w", errno)
+	}
+
+	return nil
+}
+
+// findEtcdProcess tries to find the etcd process on the node.
+func (c *EtcdIOClient) findEtcdProcess() (int, error) {
+	procDir, err := os.Open("/proc")
+	if err != nil {
+		return 0, fmt.Errorf("opening /proc: %w", err)
+	}
+	defer procDir.Close()
+
+	procEntries, err := procDir.Readdirnames(0)
+	if err != nil {
+		return 0, fmt.Errorf("reading /proc: %w", err)
+	}
+
+	// find etcd process(es)
+	etcdPIDs := []int{}
+	for _, f := range procEntries {
+		// exclude non-pid dirs
+		if f[0] < '0' || f[0] > '9' {
+			continue
+		}
+
+		exe, err := os.Readlink(fmt.Sprintf("/proc/%s/exe", f))
+		if err != nil {
+			continue
+		}
+
+		if path.Base(exe) != "etcd" {
+			continue
+		}
+
+		pid, err := strconv.Atoi(f)
+		if err != nil {
+			continue
+		}
+
+		// add the PID to the list of etcd PIDs
+		etcdPIDs = append(etcdPIDs, pid)
+	}
+
+	if len(etcdPIDs) == 0 {
+		return 0, ErrNoEtcdProcess
+	}
+
+	if len(etcdPIDs) > 1 {
+		return 0, ErrMultipleEtcdProcesses
+	}
+
+	return etcdPIDs[0], nil
+}