blueloveTH 10 месяцев назад
Родитель
Сommit
cdbe55b539
3 измененных файлов с 58 добавлено и 4 удалено
  1. 3 0
      include/typings/pkpy.pyi
  2. 48 0
      src/modules/pkpy.c
  3. 7 4
      tests/98_thread.py

+ 3 - 0
include/typings/pkpy.pyi

@@ -38,3 +38,6 @@ class ComputeThread:
     def exec(self, source: str) -> None: ...
     def eval(self, source: str) -> None: ...
     def call(self, eval_src: str, *args, **kwargs) -> None: ...
+
+    def exec_blocked(self, source: str) -> None: ...
+    def eval_blocked(self, source: str): ...

+ 48 - 0
src/modules/pkpy.c

@@ -354,6 +354,51 @@ static bool ComputeThread_call(int argc, py_Ref argv) {
     return true;
 }
 
+static bool c11_ComputeThread__exec_blocked(c11_ComputeThread* self,
+                                            const char* source,
+                                            enum py_CompileMode mode) {
+    if(!self->is_done) return OSError("thread is not done yet");
+    self->is_done = false;
+    char* err = NULL;
+    int old_vm_index = py_currentvm();
+    py_switchvm(self->vm_index);
+    py_StackRef p0 = py_peek(0);
+    if(!py_exec(source, "<job_blocked>", mode, NULL)) goto __ERROR;
+    if(!py_pickle_dumps(py_retval())) goto __ERROR;
+
+    int retval_size;
+    unsigned char* retval_data = py_tobytes(py_retval(), &retval_size);
+    py_switchvm(old_vm_index);
+    bool ok = py_pickle_loads(retval_data, retval_size);
+    self->is_done = true;
+    return ok;
+
+__ERROR:
+    err = py_formatexc();
+    py_clearexc(p0);
+    py_switchvm(old_vm_index);
+    self->is_done = true;
+    RuntimeError("c11_ComputeThread__exec_blocked() failed:\n%s", err);
+    PK_FREE(err);
+    return false;
+}
+
+static bool ComputeThread_exec_blocked(int argc, py_Ref argv) {
+    PY_CHECK_ARGC(2);
+    c11_ComputeThread* self = py_touserdata(py_arg(0));
+    PY_CHECK_ARG_TYPE(1, tp_str);
+    const char* source = py_tostr(py_arg(1));
+    return c11_ComputeThread__exec_blocked(self, source, EXEC_MODE);
+}
+
+static bool ComputeThread_eval_blocked(int argc, py_Ref argv) {
+    PY_CHECK_ARGC(2);
+    c11_ComputeThread* self = py_touserdata(py_arg(0));
+    PY_CHECK_ARG_TYPE(1, tp_str);
+    const char* source = py_tostr(py_arg(1));
+    return c11_ComputeThread__exec_blocked(self, source, EVAL_MODE);
+}
+
 static void pk_ComputeThread__register(py_Ref mod) {
     py_Type type = py_newtype("ComputeThread", tp_object, mod, (py_Dtor)c11_ComputeThread__dtor);
 
@@ -367,6 +412,9 @@ static void pk_ComputeThread__register(py_Ref mod) {
     py_bindmethod(type, "exec", ComputeThread_exec);
     py_bindmethod(type, "eval", ComputeThread_eval);
     py_bind(py_tpobject(type), "call(self, eval_src, *args, **kwargs)", ComputeThread_call);
+
+    py_bindmethod(type, "exec_blocked", ComputeThread_exec_blocked);
+    py_bindmethod(type, "eval_blocked", ComputeThread_eval_blocked);
 }
 
 void pk__add_module_pkpy() {

+ 7 - 4
tests/98_thread.py

@@ -5,7 +5,7 @@ thread_1 = ComputeThread(1)
 thread_2 = ComputeThread(2)
 
 for t in [thread_1, thread_2]:
-    t.exec('''
+    t.exec_blocked('''
 def func(a):
     from pkpy import currentvm
     print("Hello from thread", currentvm(), "a =", a)
@@ -13,10 +13,13 @@ def func(a):
         if i % 100000 == 0:
             print(i, "from thread", currentvm())
     return a
+                   
+x = 123
 ''')
-    
-thread_1.wait_for_done()
-thread_2.wait_for_done()
+assert thread_1.eval_blocked('x') == 123
+
+# thread_1.wait_for_done()
+# thread_2.wait_for_done()
 
 thread_1.call('func', [1, 2, 3])
 thread_2.call('func', [4, 5, 6])