Add feature: priority

Signed-off-by: Jing Wang <jingwang96@qq.com>

csrc/camem_allocator.cpp

@@ -595,10 +595,29 @@ static PyObject* python_try_lock_gpu_offload(PyObject* self, PyObject* args) {
 }
 
 static PyObject* python_unlock_gpu_offload(PyObject* self, PyObject* args) {
-  shm_worker->unlock_gpu();
+  int keep_wait = 0;
+  if (!PyArg_ParseTuple(args, "|p", &keep_wait)) {
+    return NULL;
+  }
+  shm_worker->unlock_gpu(keep_wait != 0);
   Py_RETURN_NONE;
 }
 
+static PyObject* python_start_wait_offload(PyObject* self, PyObject* args) {
+  shm_worker->start_wait();
+  Py_RETURN_NONE;
+}
+
+static PyObject* python_cancel_wait_offload(PyObject* self, PyObject* args) {
+  shm_worker->cancel_wait();
+  Py_RETURN_NONE;
+}
+
+static PyObject* python_has_higher_priority_waiter_offload(PyObject* self, PyObject* args) {
+  bool has_higher = shm_worker->has_higher_priority_waiter();
+  return PyBool_FromLong(has_higher);
+}
+
 static PyMethodDef module_methods[] = {
     {"init_module", (PyCFunction)py_init_module, METH_VARARGS,
      "Initialize module with python_malloc and python_free callables."},
@@ -619,7 +638,13 @@ static PyMethodDef module_methods[] = {
     {"python_try_lock_gpu_offload", (PyCFunction)python_try_lock_gpu_offload,
      METH_NOARGS, "Lock GPU."},
     {"python_unlock_gpu_offload", (PyCFunction)python_unlock_gpu_offload,
-     METH_NOARGS, "Unlock GPU."},
+     METH_VARARGS, "Unlock GPU."},
+    {"python_start_wait_offload", (PyCFunction)python_start_wait_offload,
+     METH_NOARGS, "Start waiting for GPU lock."},
+    {"python_cancel_wait_offload", (PyCFunction)python_cancel_wait_offload,
+     METH_NOARGS, "Cancel waiting for GPU lock."},
+    {"python_has_higher_priority_waiter_offload", (PyCFunction)python_has_higher_priority_waiter_offload,
+     METH_NOARGS, "Check if there is a higher priority waiter."},
     {NULL, NULL, 0, NULL} // sentinel
 };