3 files changed, 92 insertions, 0 deletions

diff --git a/include/libcamera/base/thread.h b/include/libcamera/base/thread.h
index 4f33de63..3209d4f7 100644
--- a/include/libcamera/base/thread.h
+++ b/include/libcamera/base/thread.h
@@ -15,6 +15,7 @@
 
 #include <libcamera/base/message.h>
 #include <libcamera/base/signal.h>
+#include <libcamera/base/span.h>
 #include <libcamera/base/utils.h>
 
 namespace libcamera {
@@ -35,6 +36,8 @@ public:
 	void exit(int code = 0);
 	bool wait(utils::duration duration = utils::duration::max());
 
+	int setThreadAffinity(const Span<const unsigned int> &cpus);
+
 	bool isRunning();
 
 	Signal<> finished;
@@ -54,6 +57,8 @@ private:
 	void startThread();
 	void finishThread();
 
+	void setThreadAffinityInternal();
+
 	void postMessage(std::unique_ptr<Message> msg, Object *receiver);
 	void removeMessages(Object *receiver);
 
diff --git a/src/libcamera/base/thread.cpp b/src/libcamera/base/thread.cpp
index 8735670b..f6322fe3 100644
--- a/src/libcamera/base/thread.cpp
+++ b/src/libcamera/base/thread.cpp
@@ -9,6 +9,7 @@
 
 #include <atomic>
 #include <list>
+#include <optional>
 #include <sys/syscall.h>
 #include <sys/types.h>
 #include <unistd.h>
@@ -128,6 +129,8 @@ private:
 	int exitCode_;
 
 	MessageQueue messages_;
+
+	std::optional<cpu_set_t> cpuset_;
 };
 
 /**
@@ -281,6 +284,8 @@ void Thread::startThread()
 	data_->tid_ = syscall(SYS_gettid);
 	currentThreadData = data_;
 
+	setThreadAffinityInternal();
+
 	run();
 }
 
@@ -411,6 +416,48 @@ bool Thread::wait(utils::duration duration)
 }
 
 /**
+ * \brief Set the CPU affinity mask of the thread
+ * \param[in] cpus The list of CPU indices that the thread is set affinity to
+ *
+ * The CPU indices should be within [0, std::thread::hardware_concurrency()).
+ * If any index is invalid, this function won't modify the thread affinity and
+ * will return an error.
+ *
+ * \return 0 if all indices are valid, -EINVAL otherwise
+ */
+int Thread::setThreadAffinity(const Span<const unsigned int> &cpus)
+{
+	const unsigned int numCpus = std::thread::hardware_concurrency();
+
+	MutexLocker locker(data_->mutex_);
+	data_->cpuset_ = cpu_set_t();
+	CPU_ZERO(&data_->cpuset_.value());
+
+	for (const unsigned int &cpu : cpus) {
+		if (cpu >= numCpus) {
+			LOG(Thread, Error) << "Invalid CPU " << cpu << "for thread affinity";
+			return -EINVAL;
+		}
+
+		CPU_SET(cpu, &data_->cpuset_.value());
+	}
+
+	if (data_->running_)
+		setThreadAffinityInternal();
+
+	return 0;
+}
+
+void Thread::setThreadAffinityInternal()
+{
+	if (!data_->cpuset_)
+		return;
+
+	const cpu_set_t &cpuset = data_->cpuset_.value();
+	pthread_setaffinity_np(thread_.native_handle(), sizeof(cpuset), &cpuset);
+}
+
+/**
  * \brief Check if the thread is running
  *
  * A Thread instance is considered as running once the underlying thread has
diff --git a/test/threads.cpp b/test/threads.cpp
index ceb4fa0f..8d6ee151 100644
--- a/test/threads.cpp
+++ b/test/threads.cpp
@@ -9,9 +9,11 @@
 #include <iostream>
 #include <memory>
 #include <pthread.h>
+#include <sched.h>
 #include <thread>
 #include <time.h>
 
+#include <libcamera/base/object.h>
 #include <libcamera/base/thread.h>
 
 #include "test.h"
@@ -66,6 +68,27 @@ private:
 	bool &cancelled_;
 };
 
+class CpuSetTester : public Object
+{
+public:
+	CpuSetTester(unsigned int cpuset)
+		: cpuset_(cpuset) {}
+
+	bool testCpuSet()
+	{
+		int ret = sched_getcpu();
+		if (static_cast<int>(cpuset_) != ret) {
+			cout << "Invalid cpuset: " << ret << ", expecting: " << cpuset_ << endl;
+			return false;
+		}
+
+		return true;
+	}
+
+private:
+	const unsigned int cpuset_;
+};
+
 class ThreadTest : public Test
 {
 protected:
@@ -165,6 +188,23 @@ protected:
 			return TestFail;
 		}
 
+		const unsigned int numCpus = std::thread::hardware_concurrency();
+		for (unsigned int i = 0; i < numCpus; ++i) {
+			thread = std::make_unique<Thread>();
+			const std::array<const unsigned int, 1> cpus{ i };
+			thread->setThreadAffinity(cpus);
+			thread->start();
+
+			CpuSetTester tester(i);
+			tester.moveToThread(thread.get());
+
+			if (!tester.invokeMethod(&CpuSetTester::testCpuSet, ConnectionTypeBlocking))
+				return TestFail;
+
+			thread->exit(0);
+			thread->wait();
+		}
+
 		return TestPass;
 	}