Repository: openmined/pysyft
Issue: Search remote data will cause the data being deleted

I've experimented on getting PySyft to work with remote worker with its own dataset.
During these, I found out the `search` from local will cause the data deleted on the remote.

Relevant files in [this gist](https://gist.github.com/feigaoxyz/99dffd2f4336417f83f0cf4e1c80d8fd).

To trigger the error, 
1. First run `python remote.py` in one terminal window
2. Then run `python local.py` from another terminal window

You will see the results:

```sh
$ python remote.py
Objects: {94139650912: tensor([-1.,  2.])
	Tags: x
	Shape: torch.Size([2])}
worker <WebsocketServerWorker id:0 #tensors:1> received SEARCH (b'x',)
worker <WebsocketServerWorker id:0 #tensors:1> received OBJ_DEL 94139650912
Objects: {}
```

```sh
$ python local.py
Local objects before search {}
Search result: [(Wrapper)>[PointerTensor | me:68744693343 -> 0:13527568418]]
Local objects after search {}
```

As the first one showed: the remote worker actually received two actions `SEARCH` and `OBJ_DEL` in correspondence with the search command on Line 20 of local.py.

---

This error can also be reproduced by slightly modifying the `test/test_websocket_woker.py` by searching same term twice, see [L57 of updated test](https://gist.github.com/feigaoxyz/99dffd2f4336417f83f0cf4e1c80d8fd#file-test_websocket_woker-py).

Changed files:
syft/codes.py, test/workers/test_websocket_worker.py, test/torch/tensors/test_pointer.py, syft/frameworks/torch/tensors/interpreters/pointer.py, syft/generic/object_storage.py, syft/serde.py, test/federated/test_plan.py, syft/workers/base.py

diff --git a/syft/codes.py b/syft/codes.py
index 760eb7066..425e1ded4 100644
--- a/syft/codes.py
+++ b/syft/codes.py
@@ -7,6 +7,7 @@ class MSGTYPE(object):
     IS_NONE = 6
     GET_SHAPE = 7
     SEARCH = 8
+    FORCE_OBJ_DEL = 9
 
 
 # Build automatically the reverse map from codes to msg types
diff --git a/syft/frameworks/torch/tensors/interpreters/pointer.py b/syft/frameworks/torch/tensors/interpreters/pointer.py
index 13c00bf86..c20f70735 100644
--- a/syft/frameworks/torch/tensors/interpreters/pointer.py
+++ b/syft/frameworks/torch/tensors/interpreters/pointer.py
@@ -241,10 +241,9 @@ class PointerTensor(AbstractTensor):
         # first here and not try to call self.owner.anything if self doesn't have
         # .owner anymore.
         if hasattr(self, "owner") and self.garbage_collect_data:
-
             # attribute pointers are not in charge of GC
             if self.point_to_attr is None:
-                self.owner.send_msg(MSGTYPE.OBJ_DEL, self.id_at_location, self.location)
+                self.owner.send_msg(MSGTYPE.FORCE_OBJ_DEL, self.id_at_location, self.location)
 
     @property
     def grad(self):
diff --git a/syft/generic/object_storage.py b/syft/generic/object_storage.py
index 3d7670876..b0c3b8a44 100644
--- a/syft/generic/object_storage.py
+++ b/syft/generic/object_storage.py
@@ -107,3 +107,19 @@ class ObjectStorage:
         """
         if remote_key in self._objects:
             del self._objects[remote_key]
+
+    def force_rm_obj(self, remote_key: Union[str, int]):
+        """Forces object removal.
+
+        Besides removing the object from the permanent object registry if it exists.
+        Explicitly forces removal of the object modifying the `garbage_collect_data` attribute.
+
+        Args:
+            remote_key: A string or integer representing id of the object to be
+                removed.
+        """
+        if remote_key in self._objects:
+            obj = self._objects[remote_key]
+            if hasattr(obj, "child"):
+                obj.child.garbage_collect_data = True
+            del self._objects[remote_key]
diff --git a/syft/serde.py b/syft/serde.py
index 83d144320..b7a7be42d 100644
--- a/syft/serde.py
+++ b/syft/serde.py
@@ -879,7 +879,14 @@ def _simplify_pointer_tensor(ptr: PointerTensor) -> tuple:
         data = _simplify_pointer_tensor(ptr)
     """
 
-    return (ptr.id, ptr.id_at_location, ptr.location.id, ptr.point_to_attr, ptr._shape)
+    return (
+        ptr.id,
+        ptr.id_at_location,
+        ptr.location.id,
+        ptr.point_to_attr,
+        ptr._shape,
+        ptr.garbage_collect_data,
+    )
 
     # a more general but slower/more verbose option
 
@@ -904,20 +911,16 @@ def _detail_pointer_tensor(worker: AbstractWorker, tensor_tuple: tuple) -> Point
         ptr = _detail_pointer_tensor(data)
     """
     # TODO: fix comment for this and simplifier
-    obj_id = tensor_tuple[0]
-    id_at_location = tensor_tuple[1]
-    worker_id = tensor_tuple[2]
+    obj_id, id_at_location, worker_id, point_to_attr, shape, garbage_collect_data = tensor_tuple
+
     if isinstance(worker_id, bytes):
         worker_id = worker_id.decode()
-    point_to_attr = tensor_tuple[3]
-    shape = tensor_tuple[4]
 
     if shape is not None:
         shape = torch.Size(shape)
 
     # If the pointer received is pointing at the current worker, we load the tensor instead
     if worker_id == worker.id:
-
         tensor = worker.get_obj(id_at_location)
 
         if point_to_attr is not None and tensor is not None:
@@ -950,7 +953,7 @@ def _detail_pointer_tensor(worker: AbstractWorker, tensor_tuple: tuple) -> Point
             owner=worker,
             id=obj_id,
             shape=shape,
-            garbage_collect_data=True,
+            garbage_collect_data=garbage_collect_data,
         )
 
         return ptr
diff --git a/syft/workers/base.py b/syft/workers/base.py
index e80e1f820..ac8839221 100644
--- a/syft/workers/base.py
+++ b/syft/workers/base.py
@@ -90,6 +90,7 @@ class BaseWorker(AbstractWorker, ObjectStorage):
             MSGTYPE.IS_NONE: self.is_tensor_none,
             MSGTYPE.GET_SHAPE: self.get_tensor_shape,
             MSGTYPE.SEARCH: self.deserialized_search,
+            MSGTYPE.FORCE_OBJ_DEL: self.force_rm_obj,
         }
 
         self.load_data(data)
diff --git a/test/federated/test_plan.py b/test/federated/test_plan.py
index 7510dc079..6c2c716b3 100644
--- a/test/federated/test_plan.py
+++ b/test/federated/test_plan.py
@@ -1,6 +1,5 @@
 import syft as sy
 
-import torch
 import torch as th
 import torch.nn as nn
 import torch.nn.functional as F
@@ -9,9 +8,8 @@ from syft.serde import deserialize
 from syft.serde import serialize
 import time
 
-import torch
-import mock
 import pytest
+import unittest.mock as mock
 
 from syft.workers import WebsocketClientWorker
 from syft.workers import WebsocketServerWorker
@@ -72,10 +70,10 @@ def test_plan_built_on_method(hook):
     """
     Test @sy.meth2plan and plan send / get / send
     """
-    x11 = torch.tensor([-1, 2.0]).tag("input_data")
-    x12 = torch.tensor([1, -2.0]).tag("input_data2")
-    x21 = torch.tensor([-1, 2.0]).tag("input_data")
-    x22 = torch.tensor([1, -2.0]).tag("input_data2")
+    x11 = th.tensor([-1, 2.0]).tag("input_data")
+    x12 = th.tensor([1, -2.0]).tag("input_data2")
+    x21 = th.tensor([-1, 2.0]).tag("input_data")
+    x22 = th.tensor([1, -2.0]).tag("input_data2")
 
     device_1 = sy.VirtualWorker(hook, id="device_1", data=(x11, x12))
     device_2 = sy.VirtualWorker(hook, id="device_2", data=(x21, x22))
diff --git a/test/torch/tensors/test_pointer.py b/test/torch/tensors/test_pointer.py
index fd5790eca..18239e451 100644
--- a/test/torch/tensors/test_pointer.py
+++ b/test/torch/tensors/test_pointer.py
@@ -16,6 +16,13 @@ def test_create_pointer(workers):
     x.create_pointer(location=workers["james"])
 
 
+def test_send_default_garbage_collector_true(workers):
+    """Pointer tensor should be garbage collected by default."""
+    x = torch.Tensor([-1, 2])
+    x_ptr = x.send(workers["alice"])
+    assert x_ptr.child.garbage_collect_data
+
+
 def test_send_get(workers):
     """Test several send get usages"""
     bob = workers["bob"]
diff --git a/test/workers/test_websocket_worker.py b/test/workers/test_websocket_worker.py
index 58fe67ebb..4c8ea647e 100644
--- a/test/workers/test_websocket_worker.py
+++ b/test/workers/test_websocket_worker.py
@@ -34,7 +34,6 @@ def test_websocket_worker(hook, start_proc):
 def test_websocket_workers_search(hook, start_proc):
     """Evaluates that a client can search and find tensors that belong
     to another party"""
-
     # Sample tensor to store on the server
     sample_data = torch.tensor([1, 2, 3, 4]).tag("#sample_data", "#another_tag")
     # Args for initializing the websocket server and client
@@ -54,6 +53,13 @@ def test_websocket_workers_search(hook, start_proc):
     assert results[0].owner.id == "me"
     assert results[0].location.id == "fed2"
 
+    # Search multiple times should still work
+    results = client_worker.search("#sample_data", "#another_tag")
+
+    assert results
+    assert results[0].owner.id == "me"
+    assert results[0].location.id == "fed2"
+
     client_worker.ws.shutdown()
     client_worker.ws.close()
     time.sleep(0.1)