Provide simple clean up for DataProvider on remove and UI detach

Change-Id: Ib21fbb5a254814725400b31ccaf023e4b7e13d1c

2 files changed, 50 insertions, 2 deletions

diff --git a/server/src/com/vaadin/server/communication/data/typed/AbstractDataSource.java b/server/src/com/vaadin/server/communication/data/typed/AbstractDataSource.java
index bf16b47aa5..98f68ed536 100644
--- a/server/src/com/vaadin/server/communication/data/typed/AbstractDataSource.java
+++ b/server/src/com/vaadin/server/communication/data/typed/AbstractDataSource.java
@@ -30,7 +30,9 @@ public abstract class AbstractDataSource<T> implements DataSource<T> {
 
     @Override
     public void addDataChangeHandler(DataChangeHandler<T> handler) {
-        handlers.add(handler);
+        if (handler != null) {
+            handlers.add(handler);
+        }
     }
 
     @Override
diff --git a/server/src/com/vaadin/server/communication/data/typed/DataProvider.java b/server/src/com/vaadin/server/communication/data/typed/DataProvider.java
index 7cba66aa64..2c45c52ed5 100644
--- a/server/src/com/vaadin/server/communication/data/typed/DataProvider.java
+++ b/server/src/com/vaadin/server/communication/data/typed/DataProvider.java
@@ -22,6 +22,7 @@ import java.util.LinkedHashSet;
 import java.util.Set;
 
 import com.vaadin.server.AbstractExtension;
+import com.vaadin.server.ClientConnector;
 import com.vaadin.server.communication.data.typed.DataSource.DataChangeHandler;
 import com.vaadin.shared.data.DataProviderClientRpc;
 import com.vaadin.shared.data.DataProviderConstants;
@@ -178,13 +179,42 @@ public abstract class DataProvider<T> extends AbstractExtension {
     protected DataProviderClientRpc rpc;
 
     protected DataSource<T> dataSource;
+    private DataChangeHandler<T> dataChangeHandler;
+    private DetachListener detachListener;
 
     protected DataProvider(DataSource<T> dataSource) {
         addDataGenerator(handler);
         this.dataSource = dataSource;
         rpc = getRpcProxy(DataProviderClientRpc.class);
         registerRpc(createRpc());
-        this.dataSource.addDataChangeHandler(createDataChangeHandler());
+        dataChangeHandler = createDataChangeHandler();
+        this.dataSource.addDataChangeHandler(dataChangeHandler);
+    }
+
+    @Override
+    public void attach() {
+        super.attach();
+
+        if (detachListener == null) {
+            detachListener = new DetachListener() {
+
+                @Override
+                public void detach(DetachEvent event) {
+                    cleanUp();
+                }
+            };
+            getUI().addDetachListener(detachListener);
+        }
+    }
+
+    @Override
+    public void setParent(ClientConnector parent) {
+        if (getParent() != null && parent == null) {
+            // Removing from parent, clean up.
+            cleanUp();
+        }
+
+        super.setParent(parent);
     }
 
     /**
@@ -267,6 +297,22 @@ public abstract class DataProvider<T> extends AbstractExtension {
     }
 
     /**
+     * Clean up method for removing all listeners attached by the
+     * {@link DataProvider}. This method is called from {@link #remove()} or
+     * when the UI gets detached.
+     */
+    protected void cleanUp() {
+        if (dataSource != null) {
+            dataSource.removeDataChangeHandler(dataChangeHandler);
+            dataChangeHandler = null;
+        }
+        if (detachListener != null) {
+            getUI().removeDetachListener(detachListener);
+            detachListener = null;
+        }
+    }
+
+    /**
      * Gets the {@link DataKeyMapper} used by this {@link DataProvider}.
      * 
      * @return key mapper