Cgroup Wrapper

cgroup.hpp

#pragma once

#include "../dep/PutitLib/UNIXFunc/unixfunc.hpp"
#include "../dep/PutitLib/IOLib/io.hpp"
#include <set>
#include <string>
#include <stdexcept>
#include <fstream>
#include <map>
#include <iterator>

extern void* handler;

/*

*** To use this, you must install cgroup-bin and cgroup-lite. ***

*** You should be careful of Shell Injecion. ***

*/

namespace sandbox {
	namespace cgroup{
		struct _impl_ES {
			const std::set<std::string>& list;
			
			_impl_ES():_list(), list(_list) { // Exceptionable
				this->reset();
			}

			_impl_ES(const _impl_ES&) = delete;
			_impl_ES(_impl_ES&&) = delete;
			_impl_ES& operator =(const _impl_ES&) = delete;
			_impl_ES& operator =(_impl_ES&&) = delete;

			void reset() { // Exceptionable
				_list.clear();

				const auto& res = UNIXFunc::runCmd("lssubsys");

				if(!res.first) {
					throw std::runtime_error{"Unable to execute lssubsys"};
				}

				std::stringstream ss;
				ss << res.second;

				std::string subsys;
				while(std::getline(ss, subsys)) {
					if(!subsys.empty()) {
						_list.insert(std::move(subsys));
					}
				}
			}
		private:
			std::set<std::string> _list;

		} EnabledSubsys;

		class cgroup {
			bool _isValid = false;
			bool _isCgroupCreated = false;
			std::string groupName;
			std::set<std::string> subsys;
			std::string unfoldedSubsys;
		public:
			cgroup() {
				this->_isValid = true;
			}

			template<typename T>
			cgroup(const T& subsystems, const std::string& groupName) {
				this->groupName = groupName;

				for(const auto& subs : subsystems) {
					this->addSubsys(subs);
				}
			}
			cgroup(const cgroup&) = delete;
			cgroup(cgroup&& rhs) {
				this->_isValid = rhs._isValid;
				this->_isCgroupCreated = rhs._isCgroupCreated;
	
				this->groupName = std::move(rhs.groupName);
	
				rhs._isValid = false;
			}
			cgroup& operator = (cgroup&& rhs) { // Exceptionable
				this->_isValid = rhs._isValid;
				this->_isCgroupCreated = rhs._isCgroupCreated;

				this->groupName = std::move(rhs.groupName);

				rhs._isValid = false;

				return *this;
			}
			cgroup& operator = (const cgroup&) = delete;

			~cgroup() {
				if(_isValid && _isCgroupCreated) {
					try {
						this->remove();
					}
					catch(...) {}
				}
			}

			cgroup& addSubsys(const std::string& subsysName) { // Exceptionable
				if(_isCgroupCreated) {
					throw std::runtime_error{"error: A group has already created."};
				}

				if(EnabledSubsys.list.count(subsysName) == 0) {
					throw std::runtime_error{
							IOLib::format(
								"error: Disabled subsystem(",
								subsysName,
								")"
							)
						};
				}

				subsys.insert(subsysName);

				return *this;
			}

			cgroup& create(const std::string& groupName) { // Exceptionable
				if(_isCgroupCreated) {
					throw std::runtime_error{"error: A cgroup has already created."};
				}
				if(subsys.empty()) {
					throw std::runtime_error{"error: There are no subsystems."};
				}

				for(auto it = subsys.begin(); it != subsys.end(); ++it) {
					unfoldedSubsys += *it + ',';
				}
				unfoldedSubsys = std::move(unfoldedSubsys.substr(0, unfoldedSubsys.size() - 1));

				auto&& res = UNIXFunc::runCmd(
					IOLib::format("cgcreate -g ", unfoldedSubsys, ":", groupName));

				if(!res.first || !res.second.empty()) {
					throw std::runtime_error{"error: Unable to create cgroup (output: " + res.second + ")"};
				}
				
				this->groupName = groupName;
				_isCgroupCreated = true;

				return *this;
			}

			void remove() { // Exceptionable
				if(!_isCgroupCreated) {
					throw std::runtime_error{"error: A cgroup hasn't created yet."};
				}

				auto&& res = UNIXFunc::runCmd(
					IOLib::format("cgdelete -r ", unfoldedSubsys, ":", groupName)
					);

				if(!res.first || !res.second.empty()) {
					throw std::runtime_error{"error: Unable to remove the cgroup(output: " + res.second + ")"};
				}

				_isCgroupCreated = false;
			}
			
			template<typename T>
			cgroup& setSubsysSetting(const std::string& subsysName, const std::string& param, T&& val) {
				auto&& str = CastLib::cast<std::string>(std::forward<T>(val));

				auto&& res = UNIXFunc::runCmd("cgset -r " + subsysName + "." + param + "=" + str + " " + groupName);

				if(!res.first || !res.second.empty()) {
					throw std::runtime_error{"Unable to set " + subsysName + "." + param + "=" + str +  " in " + str};
				}

				return *this;
			}

			cgroup& setSubsysSetting(const std::string& subsysName, const std::map<std::string, std::string>& settings) {
				for(auto&& set :  settings) {
					this->setSubsysSetting(subsysName, set.first, set.second);
				}
			}
			
		private:
			// *** Don't hold this as a lvalue. ***
			struct FocusedSetting {
				friend class cgroup;
			private:
				const std::string& setName;
				cgroup& cgRef;

				FocusedSetting(const std::string& sn, cgroup& cg) : setName(sn), cgRef(cg) {}
				FocusedSetting(FocusedSetting&&) = default;
				FocusedSetting& operator = (FocusedSetting&&) = default;
			public:
				template<typename T>
				FocusedSetting& set(const std::string& param, T&& val) { // Exceptionable
					cgRef.setSubsysSetting(setName, param, std::forward<T>(val));

					return *this;
				}
				cgroup& blur() {
					return cgRef;
				}
			};
		public:
			FocusedSetting focus(const std::string& subsysName) noexcept {
				return FocusedSetting(subsysName, *this);
			}

			void removeSubsys(const std::string& subsysName) { // Exceptionable
				if(_isCgroupCreated) {
					throw std::runtime_error{"error: A cgroup has already created."};
				}
				
				auto&& fres = subsys.find(subsysName);
				
				if(fres == subsys.end()) {
					throw std::runtime_error{"error: Unregistered Subsystem"};
				}

				subsys.erase(fres);
			}

			template<typename T = std::string>
			T getStatusValue(const std::string& subsysName, const std::string& param) { // Exceptionable
				if(subsys.count(subsysName) == 0) {
					throw std::runtime_error{"error: Unregistered Subsystem"};
				}
				
				// Cgroup Path
				std::ifstream ifs("/sys/fs/cgroup/" + subsysName + "/" + groupName + "/" + subsysName + "." + param);

				if(!ifs) {
					throw std::runtime_error{"error: Unable to open " + subsysName + "." + param + " in " + groupName};
				}

				return CastLib::cast<T>(
					std::string{
						std::istream_iterator<char>{ifs},
						std::istream_iterator<char>{}
					}
				);
			}

			template<typename T>
			void registerProcessId(T&& pid) { // Exceptionable
				auto&& res = UNIXFunc::runCmd(
					IOLib::format(
						"cgclassify -g ", unfoldedSubsys, ":", groupName, " ", pid
					)
				);

				if(!res.first || !res.second.empty()) {
					throw std::runtime_error{"error: Unable to register a process with " + unfoldedSubsys + ":" + groupName};
				}
			}
		};
	}
}