diff --git a/cpp/include/ray/api/object_ref.h b/cpp/include/ray/api/object_ref.h index 71500202b..c4ceff90c 100644 --- a/cpp/include/ray/api/object_ref.h +++ b/cpp/include/ray/api/object_ref.h @@ -16,6 +16,7 @@ template class ObjectRef { public: ObjectRef(); + ~ObjectRef(); ObjectRef(const ObjectID &id); @@ -46,6 +47,17 @@ ObjectRef::ObjectRef() {} template ObjectRef::ObjectRef(const ObjectID &id) { id_ = id; + if (CoreWorkerProcess::IsInitialized()) { + auto &core_worker = CoreWorkerProcess::GetCoreWorker(); + core_worker.AddLocalReference(id_); + } +} +template +ObjectRef::~ObjectRef() { + if (CoreWorkerProcess::IsInitialized()) { + auto &core_worker = CoreWorkerProcess::GetCoreWorker(); + core_worker.RemoveLocalReference(id_); + } } template diff --git a/cpp/src/ray/test/cluster/cluster_mode_test.cc b/cpp/src/ray/test/cluster/cluster_mode_test.cc index 68ac2ec02..780fb0d30 100644 --- a/cpp/src/ray/test/cluster/cluster_mode_test.cc +++ b/cpp/src/ray/test/cluster/cluster_mode_test.cc @@ -91,16 +91,13 @@ TEST(RayClusterModeTest, FullTest) { auto r5 = Ray::Task(Plus, r4, r3).Remote(); auto r6 = Ray::Task(Plus, r4, 10).Remote(); - ///// TODO(ameer/guyang): All the commented code lines below should be - ///// uncommented once reference counting is added. Currently the objects - ///// are leaking from the object store. int result5 = *(Ray::Get(r5)); - // int result4 = *(Ray::Get(r4)); + int result4 = *(Ray::Get(r4)); int result6 = *(Ray::Get(r6)); - // int result3 = *(Ray::Get(r3)); + int result3 = *(Ray::Get(r3)); EXPECT_EQ(result0, 1); - // EXPECT_EQ(result3, 1); - // EXPECT_EQ(result4, 2); + EXPECT_EQ(result3, 1); + EXPECT_EQ(result4, 2); EXPECT_EQ(result5, 3); EXPECT_EQ(result6, 12); @@ -114,37 +111,34 @@ TEST(RayClusterModeTest, FullTest) { int result7 = *(Ray::Get(r7)); int result8 = *(Ray::Get(r8)); int result9 = *(Ray::Get(r9)); - // int result10 = *(Ray::Get(r10)); + int result10 = *(Ray::Get(r10)); EXPECT_EQ(result7, 15); EXPECT_EQ(result8, 16); EXPECT_EQ(result9, 19); - // EXPECT_EQ(result10, 27); + EXPECT_EQ(result10, 27); /// create actor and task function remote call with args passed by reference - // ActorHandle actor5 = Ray::Actor(Counter::FactoryCreate, r10, 0).Remote(); - ActorHandle actor5 = Ray::Actor(Counter::FactoryCreate, 27, 0).Remote(); - // auto r11 = actor5.Task(&Counter::Add, r0).Remote(); - auto r11 = actor5.Task(&Counter::Add, 1).Remote(); - // auto r12 = actor5.Task(&Counter::Add, r11).Remote(); + ActorHandle actor5 = Ray::Actor(Counter::FactoryCreate, r10, 0).Remote(); + + auto r11 = actor5.Task(&Counter::Add, r0).Remote(); + auto r12 = actor5.Task(&Counter::Add, r11).Remote(); auto r13 = actor5.Task(&Counter::Add, r10).Remote(); auto r14 = actor5.Task(&Counter::Add, r13).Remote(); - // auto r15 = Ray::Task(Plus, r0, r11).Remote(); - auto r15 = Ray::Task(Plus, 1, r11).Remote(); + auto r15 = Ray::Task(Plus, r0, r11).Remote(); auto r16 = Ray::Task(Plus1, r15).Remote(); - // int result12 = *(Ray::Get(r12)); + int result12 = *(Ray::Get(r12)); int result14 = *(Ray::Get(r14)); - // int result11 = *(Ray::Get(r11)); - // int result13 = *(Ray::Get(r13)); + int result11 = *(Ray::Get(r11)); + int result13 = *(Ray::Get(r13)); int result16 = *(Ray::Get(r16)); - // int result15 = *(Ray::Get(r15)); + int result15 = *(Ray::Get(r15)); - // EXPECT_EQ(result11, 28); - // EXPECT_EQ(result12, 56); - // EXPECT_EQ(result13, 83); - // EXPECT_EQ(result14, 166); - EXPECT_EQ(result14, 110); - // EXPECT_EQ(result15, 29); + EXPECT_EQ(result11, 28); + EXPECT_EQ(result12, 56); + EXPECT_EQ(result13, 83); + EXPECT_EQ(result14, 166); + EXPECT_EQ(result15, 29); EXPECT_EQ(result16, 30); Ray::Shutdown();