From b15d92ff405320b0b0a95e536f3dfb29ce763f9f Mon Sep 17 00:00:00 2001
From: Marc Fiuczynski <mef@cs.princeton.edu>
Date: Wed, 22 Jun 2005 03:52:02 +0000
Subject: [PATCH] vserverimpl now includes support to set both task and memory
 limits. cpulimit.py is a backwards compat file to write cpu limits via ckrm;
 however, this will be deprecated when we switch to Andy's new sched for which
 the setsched support already is in vserverimpl.

---
 python/cpulimit.py   | 117 +++++++++++++++++++++++++++++++++++++++++++
 python/vserverimpl.c |  43 +++++++++++++++-
 2 files changed, 158 insertions(+), 2 deletions(-)
 create mode 100755 python/cpulimit.py

diff --git a/python/cpulimit.py b/python/cpulimit.py
new file mode 100755
index 0000000..708f092
--- /dev/null
+++ b/python/cpulimit.py
@@ -0,0 +1,117 @@
+#!/usr/bin/python2 -u
+
+import sys, os, re, string
+
+
+TCBASE="/rcfs/taskclass/"
+RULES="/rcfs/ce/rules/"
+
+SYSCLASS=TCBASE + "system"
+SYSRULE=RULES + "system"
+SYSCPUSHARE=100
+DEFAULTMAXCPUSHARE=8192
+
+def checkckrm():
+    checks = [ "/rcfs", TCBASE, RULES ]
+
+    for check in checks:
+        try:
+            answer = os.stat(check)
+        except:
+            print "%s does not exist" % check
+            return False
+    
+    return True
+
+def checkclass(tc):
+    try:
+        answer = os.stat(TCBASE + tc)
+        return True
+
+    except:
+        print "%s class does not exist" % tc
+        return False
+
+def getxid(name):
+    xid = -1
+    fp = open('/etc/passwd')
+    for line in fp.readlines():
+        rec = string.splitfields(line,':')
+        if rec[0] == name:
+            xid = int(rec[2])
+            break
+
+    fp.close()
+    
+    if xid == -1:
+        # raise an exception
+        pass
+
+    return xid
+
+def write(filename,s):
+    fp = os.open(filename,os.O_WRONLY|os.O_CREAT)
+    os.write(fp,s)
+    os.close(fp)
+
+def vs2ckrm_on(tc):
+    xid = getxid(tc)
+
+    try:
+        os.mkdir(TCBASE + tc)
+    except OSError:
+        pass # ignore oserror for file exists
+    
+    s = "xid=%d,class=%s" % (xid,TCBASE+tc)
+    fname = RULES + tc
+    write(fname, s)
+
+def vs2ckrm_off(tc):
+    fname = TCBASE + tc + "/members"
+    for i in range(1,15):
+        fp = open(fname)
+        lines = fp.readlines()
+        try:
+            lines.remove("No data to display\n")
+        except ValueError:
+            pass
+        if len(lines) == 0:
+            try:
+                answer = os.stat(RULES + tc)
+                os.unlink(RULES + tc)
+                answer = os.stat(TCBASE + tc)                
+                os.rmdir(TCBASE + tc)
+            except:
+                pass
+            break
+
+        else:
+            print "enter context 1 and kill processes", lines
+        
+
+def cpulimit(tc,limit):
+    global TCBASE
+
+    fname = TCBASE + tc + "/shares"
+    s = "res=cpu,guarantee=%d\n" % limit
+    write(fname,s)
+
+def cpuinit():
+    global TCBASE
+
+    fname = TCBASE + "shares"
+    s = "res=cpu,total_guarantee=%d\n" % DEFAULTMAXCPUSHARE
+    write(fname,s)
+
+if __name__ == "__main__":
+    try:
+        name = sys.argv[1]
+        limit = int(sys.argv[2])
+    except:
+        print "caught exception"
+
+    if checkckrm() is True:
+        cpuinit()
+        vs2ckrm_on(name)
+        cpulimit(name,limit)
+        vs2ckrm_off(name)
diff --git a/python/vserverimpl.c b/python/vserverimpl.c
index ad58d48..58b0a61 100644
--- a/python/vserverimpl.c
+++ b/python/vserverimpl.c
@@ -70,6 +70,43 @@ vserver_chcontext(PyObject *self, PyObject *args)
   return Py_None;
 }
 
+static PyObject *
+__vserver_rlimit(PyObject *self, PyObject *args, int resource) {
+	struct vc_rlimit limits;
+	int xid; 
+	PyObject *ret;
+
+	limits.min = VC_LIM_KEEP;
+	limits.soft = VC_LIM_KEEP;
+	limits.hard = VC_LIM_KEEP;
+
+	if (!PyArg_ParseTuple(args, "iL", &xid, &limits.hard))
+		return NULL;
+
+	ret = Py_None;
+	if (vc_set_rlimit(xid, resource, &limits)) 
+		ret = PyErr_SetFromErrno(PyExc_OSError);
+	else if (vc_get_rlimit(xid, resource, &limits)==-1)
+		ret = PyErr_SetFromErrno(PyExc_OSError);
+	else
+		ret = Py_BuildValue("L",limits.hard);
+
+	return ret;
+}
+
+
+static PyObject *
+vserver_memlimit(PyObject *self, PyObject *args) {
+	return __vserver_rlimit(self,args,5);
+}
+
+static PyObject *
+vserver_tasklimit(PyObject *self, PyObject *args) {
+	return __vserver_rlimit(self,args,6);
+}
+
+
+
 /*
  * setsched
  */
@@ -177,8 +214,6 @@ vserver_dlimit(PyObject *self, PyObject *args)
 	return res;
 }
 
-
-
 static PyMethodDef  methods[] = {
   { "chcontext", vserver_chcontext, METH_VARARGS,
     "Change to the given vserver context" },
@@ -186,6 +221,10 @@ static PyMethodDef  methods[] = {
     "Change vserver scheduling attributes for given vserver context" },
   { "dlimit", vserver_dlimit, METH_VARARGS,
     "Set disk limits for given vserver context" },
+  { "tasklimit", vserver_tasklimit, METH_VARARGS,
+    "Set task limits for given vserver context" },
+  { "memlimit", vserver_memlimit, METH_VARARGS,
+    "Set memory limits for given vserver context" },
   { NULL, NULL, 0, NULL }
 };
 
-- 
2.43.0