diff --git a/src/orchpylib.cc b/src/orchpylib.cc index ee47c65ab..0be35f747 100644 --- a/src/orchpylib.cc +++ b/src/orchpylib.cc @@ -637,9 +637,12 @@ PyObject* remote_call(PyObject* self, PyObject* args) { request.release_call(); // TODO: Make sure that call is not moved, otherwise capsule pointer needs to be updated int size = reply.result_size(); PyObject* list = PyList_New(size); + std::vector result_objrefs; for (int i = 0; i < size; ++i) { PyList_SetItem(list, i, make_pyobjref(worker_capsule, reply.result(i))); + result_objrefs.push_back(reply.result(i)); } + worker->decrement_reference_count(result_objrefs); // The corresponding increment is done in RemoteCall in the scheduler. return list; } diff --git a/src/scheduler.cc b/src/scheduler.cc index 5772d80ef..6e5038e97 100644 --- a/src/scheduler.cc +++ b/src/scheduler.cc @@ -25,7 +25,8 @@ Status SchedulerService::RemoteCall(ServerContext* context, const RemoteCallRequ } { std::lock_guard reference_counts_lock(reference_counts_lock_); // we grab this lock because increment_ref_count assumes it has been acquired - increment_ref_count(result_objrefs); // The corresponding decrement will happen in deserialize_call in orchpylib. + increment_ref_count(result_objrefs); // We increment once so the objrefs don't go out of scope before we reply to the worker that called RemoteCall. The corresponding decrement will happen in remote_call in orchpylib. + increment_ref_count(result_objrefs); // We increment once so the objrefs don't go out of scope before the task is scheduled on the worker. The corresponding decrement will happen in deserialize_call in orchpylib. } task_queue_lock_.lock();